diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index dc3aa102be78..3ca9435d97e0 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -28,7 +28,7 @@ jobs: options: --shm-size "16gb" --ipc host --gpus all steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 - name: NVIDIA-SMI @@ -58,24 +58,10 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: benchmark_test_reports path: benchmarks/${{ env.BASE_PATH }} - - # TODO: enable this once the connection problem has been resolved. - - name: Update benchmarking results to DB - env: - PGDATABASE: metrics - PGHOST: ${{ secrets.DIFFUSERS_BENCHMARKS_PGHOST }} - PGUSER: transformers_benchmarks - PGPASSWORD: ${{ secrets.DIFFUSERS_BENCHMARKS_PGPASSWORD }} - BRANCH_NAME: ${{ github.head_ref || github.ref_name }} - run: | - git config --global --add safe.directory /__w/diffusers/diffusers - commit_id=$GITHUB_SHA - commit_msg=$(git show -s --format=%s "$commit_id" | cut -c1-70) - cd benchmarks && python populate_into_db.py "$BRANCH_NAME" "$commit_id" "$commit_msg" - name: Report success status if: ${{ success() }} diff --git a/.github/workflows/build_docker_images.yml b/.github/workflows/build_docker_images.yml index b1af44736730..f928e123aa8f 100644 --- a/.github/workflows/build_docker_images.yml +++ b/.github/workflows/build_docker_images.yml @@ -25,10 +25,10 @@ jobs: if: github.event_name == 'pull_request' steps: - name: Set up Docker Buildx - uses: docker/setup-buildx-action@v1 + uses: docker/setup-buildx-action@v3 - name: Check out code - uses: actions/checkout@v3 + uses: actions/checkout@v6 - name: Find Changed Dockerfiles id: file_changes @@ -99,16 +99,16 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@v3 + uses: actions/checkout@v6 - name: Set up Docker Buildx - uses: docker/setup-buildx-action@v1 + uses: docker/setup-buildx-action@v3 - name: Login to Docker Hub - uses: docker/login-action@v2 + uses: docker/login-action@v3 with: username: ${{ env.REGISTRY }} password: ${{ secrets.DOCKERHUB_TOKEN }} - name: Build and push - uses: docker/build-push-action@v3 + uses: docker/build-push-action@v6 with: no-cache: true context: ./docker/${{ matrix.image-name }} diff --git a/.github/workflows/build_pr_documentation.yml b/.github/workflows/build_pr_documentation.yml index f47645c1f659..8e8dc92cb57d 100644 --- a/.github/workflows/build_pr_documentation.yml +++ b/.github/workflows/build_pr_documentation.yml @@ -17,10 +17,10 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: '3.10' diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml new file mode 100644 index 000000000000..5ba158b46fde --- /dev/null +++ b/.github/workflows/codeql.yml @@ -0,0 +1,22 @@ +--- +name: CodeQL Security Analysis For Github Actions + +on: + push: + branches: ["main"] + workflow_dispatch: + # pull_request: + +jobs: + codeql: + name: CodeQL Analysis + uses: huggingface/security-workflows/.github/workflows/codeql-reusable.yml@v1 + permissions: + security-events: write + packages: read + actions: read + contents: read + with: + languages: '["actions","python"]' + queries: 'security-extended,security-and-quality' + runner: 'ubuntu-latest' #optional if need custom runner diff --git a/.github/workflows/mirror_community_pipeline.yml b/.github/workflows/mirror_community_pipeline.yml index ab4ded973047..73cced7c1394 100644 --- a/.github/workflows/mirror_community_pipeline.yml +++ b/.github/workflows/mirror_community_pipeline.yml @@ -24,7 +24,6 @@ jobs: mirror_community_pipeline: env: SLACK_WEBHOOK_URL: ${{ secrets.SLACK_WEBHOOK_URL_COMMUNITY_MIRROR }} - runs-on: ubuntu-22.04 steps: # Checkout to correct ref @@ -39,37 +38,41 @@ jobs: # If ref is 'refs/heads/main' => set 'main' # Else it must be a tag => set {tag} - name: Set checkout_ref and path_in_repo + env: + EVENT_NAME: ${{ github.event_name }} + EVENT_INPUT_REF: ${{ github.event.inputs.ref }} + GITHUB_REF: ${{ github.ref }} run: | - if [ "${{ github.event_name }}" == "workflow_dispatch" ]; then - if [ -z "${{ github.event.inputs.ref }}" ]; then + if [ "$EVENT_NAME" == "workflow_dispatch" ]; then + if [ -z "$EVENT_INPUT_REF" ]; then echo "Error: Missing ref input" exit 1 - elif [ "${{ github.event.inputs.ref }}" == "main" ]; then + elif [ "$EVENT_INPUT_REF" == "main" ]; then echo "CHECKOUT_REF=refs/heads/main" >> $GITHUB_ENV echo "PATH_IN_REPO=main" >> $GITHUB_ENV else - echo "CHECKOUT_REF=refs/tags/${{ github.event.inputs.ref }}" >> $GITHUB_ENV - echo "PATH_IN_REPO=${{ github.event.inputs.ref }}" >> $GITHUB_ENV + echo "CHECKOUT_REF=refs/tags/$EVENT_INPUT_REF" >> $GITHUB_ENV + echo "PATH_IN_REPO=$EVENT_INPUT_REF" >> $GITHUB_ENV fi - elif [ "${{ github.ref }}" == "refs/heads/main" ]; then - echo "CHECKOUT_REF=${{ github.ref }}" >> $GITHUB_ENV + elif [ "$GITHUB_REF" == "refs/heads/main" ]; then + echo "CHECKOUT_REF=$GITHUB_REF" >> $GITHUB_ENV echo "PATH_IN_REPO=main" >> $GITHUB_ENV else # e.g. refs/tags/v0.28.1 -> v0.28.1 - echo "CHECKOUT_REF=${{ github.ref }}" >> $GITHUB_ENV - echo "PATH_IN_REPO=$(echo ${{ github.ref }} | sed 's/^refs\/tags\///')" >> $GITHUB_ENV + echo "CHECKOUT_REF=$GITHUB_REF" >> $GITHUB_ENV + echo "PATH_IN_REPO=$(echo $GITHUB_REF | sed 's/^refs\/tags\///')" >> $GITHUB_ENV fi - name: Print env vars run: | echo "CHECKOUT_REF: ${{ env.CHECKOUT_REF }}" echo "PATH_IN_REPO: ${{ env.PATH_IN_REPO }}" - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 with: ref: ${{ env.CHECKOUT_REF }} # Setup + install dependencies - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: python-version: "3.10" - name: Install dependencies @@ -99,4 +102,4 @@ jobs: - name: Report failure status if: ${{ failure() }} run: | - pip install requests && python utils/notify_community_pipelines_mirror.py --status=failure \ No newline at end of file + pip install requests && python utils/notify_community_pipelines_mirror.py --status=failure diff --git a/.github/workflows/nightly_tests.yml b/.github/workflows/nightly_tests.yml index 8b7e57e91297..416d2af3fc2e 100644 --- a/.github/workflows/nightly_tests.yml +++ b/.github/workflows/nightly_tests.yml @@ -28,7 +28,7 @@ jobs: pipeline_test_matrix: ${{ steps.fetch_pipeline_matrix.outputs.pipeline_test_matrix }} steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 - name: Install dependencies @@ -44,7 +44,7 @@ jobs: - name: Pipeline Tests Artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: test-pipelines.json path: reports @@ -64,7 +64,7 @@ jobs: options: --shm-size "16gb" --ipc host --gpus all steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 - name: NVIDIA-SMI @@ -97,7 +97,7 @@ jobs: cat reports/tests_pipeline_${{ matrix.module }}_cuda_failures_short.txt - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: pipeline_${{ matrix.module }}_test_reports path: reports @@ -119,7 +119,7 @@ jobs: module: [models, schedulers, lora, others, single_file, examples] steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -167,7 +167,7 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: torch_${{ matrix.module }}_cuda_test_reports path: reports @@ -184,7 +184,7 @@ jobs: steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -211,7 +211,7 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: torch_compile_test_reports path: reports @@ -228,7 +228,7 @@ jobs: options: --shm-size "16gb" --ipc host --gpus all steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 - name: NVIDIA-SMI @@ -263,7 +263,7 @@ jobs: cat reports/tests_big_gpu_torch_cuda_failures_short.txt - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: torch_cuda_big_gpu_test_reports path: reports @@ -280,7 +280,7 @@ jobs: shell: bash steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -321,7 +321,7 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: torch_minimum_version_cuda_test_reports path: reports @@ -355,7 +355,7 @@ jobs: options: --shm-size "20gb" --ipc host --gpus all steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 - name: NVIDIA-SMI @@ -391,7 +391,7 @@ jobs: cat reports/tests_${{ matrix.config.backend }}_torch_cuda_failures_short.txt - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: torch_cuda_${{ matrix.config.backend }}_reports path: reports @@ -408,7 +408,7 @@ jobs: options: --shm-size "20gb" --ipc host --gpus all steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 - name: NVIDIA-SMI @@ -441,7 +441,7 @@ jobs: cat reports/tests_pipeline_level_quant_torch_cuda_failures_short.txt - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: torch_cuda_pipeline_level_quant_reports path: reports @@ -466,7 +466,7 @@ jobs: image: diffusers/diffusers-pytorch-cpu steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -474,7 +474,7 @@ jobs: run: mkdir -p combined_reports - name: Download all test reports - uses: actions/download-artifact@v4 + uses: actions/download-artifact@v7 with: path: artifacts @@ -500,7 +500,7 @@ jobs: cat $CONSOLIDATED_REPORT_PATH >> $GITHUB_STEP_SUMMARY - name: Upload consolidated report - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: consolidated_test_report path: ${{ env.CONSOLIDATED_REPORT_PATH }} @@ -514,7 +514,7 @@ jobs: # # steps: # - name: Checkout diffusers -# uses: actions/checkout@v3 +# uses: actions/checkout@v6 # with: # fetch-depth: 2 # @@ -554,7 +554,7 @@ jobs: # # - name: Test suite reports artifacts # if: ${{ always() }} -# uses: actions/upload-artifact@v4 +# uses: actions/upload-artifact@v6 # with: # name: torch_mps_test_reports # path: reports @@ -570,7 +570,7 @@ jobs: # # steps: # - name: Checkout diffusers -# uses: actions/checkout@v3 +# uses: actions/checkout@v6 # with: # fetch-depth: 2 # @@ -610,7 +610,7 @@ jobs: # # - name: Test suite reports artifacts # if: ${{ always() }} -# uses: actions/upload-artifact@v4 +# uses: actions/upload-artifact@v6 # with: # name: torch_mps_test_reports # path: reports diff --git a/.github/workflows/notify_slack_about_release.yml b/.github/workflows/notify_slack_about_release.yml index 612ad4e24503..7751827d81f5 100644 --- a/.github/workflows/notify_slack_about_release.yml +++ b/.github/workflows/notify_slack_about_release.yml @@ -10,12 +10,12 @@ jobs: runs-on: ubuntu-22.04 steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 - name: Setup Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: - python-version: '3.8' + python-version: '3.10' - name: Notify Slack about the release env: diff --git a/.github/workflows/pr_dependency_test.yml b/.github/workflows/pr_dependency_test.yml index b914d1076190..ba5cd1c76cbc 100644 --- a/.github/workflows/pr_dependency_test.yml +++ b/.github/workflows/pr_dependency_test.yml @@ -18,11 +18,11 @@ jobs: check_dependencies: runs-on: ubuntu-22.04 steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: - python-version: "3.8" + python-version: "3.10" - name: Install dependencies run: | pip install -e . diff --git a/.github/workflows/pr_modular_tests.yml b/.github/workflows/pr_modular_tests.yml index 13c228621f5c..89b502d364ec 100644 --- a/.github/workflows/pr_modular_tests.yml +++ b/.github/workflows/pr_modular_tests.yml @@ -1,3 +1,4 @@ + name: Fast PR tests for Modular on: @@ -35,9 +36,9 @@ jobs: check_code_quality: runs-on: ubuntu-22.04 steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: python-version: "3.10" - name: Install dependencies @@ -55,9 +56,9 @@ jobs: needs: check_code_quality runs-on: ubuntu-22.04 steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: python-version: "3.10" - name: Install dependencies @@ -74,26 +75,34 @@ jobs: if: ${{ failure() }} run: | echo "Repo consistency check failed. Please ensure the right dependency versions are installed with 'pip install -e .[quality]' and run 'make fix-copies'" >> $GITHUB_STEP_SUMMARY + check_auto_docs: + runs-on: ubuntu-22.04 + steps: + - uses: actions/checkout@v6 + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: "3.10" + - name: Install dependencies + run: | + pip install --upgrade pip + pip install .[quality] + - name: Check auto docs + run: make modular-autodoctrings + - name: Check if failure + if: ${{ failure() }} + run: | + echo "Auto docstring checks failed. Please run `python utils/modular_auto_docstring.py --fix_and_overwrite`." >> $GITHUB_STEP_SUMMARY run_fast_tests: - needs: [check_code_quality, check_repository_consistency] - strategy: - fail-fast: false - matrix: - config: - - name: Fast PyTorch Modular Pipeline CPU tests - framework: pytorch_pipelines - runner: aws-highmemory-32-plus - image: diffusers/diffusers-pytorch-cpu - report: torch_cpu_modular_pipelines - - name: ${{ matrix.config.name }} + needs: [check_code_quality, check_repository_consistency, check_auto_docs] + name: Fast PyTorch Modular Pipeline CPU tests runs-on: - group: ${{ matrix.config.runner }} + group: aws-highmemory-32-plus container: - image: ${{ matrix.config.image }} + image: diffusers/diffusers-pytorch-cpu options: --shm-size "16gb" --ipc host -v /mnt/hf_cache:/mnt/cache/ defaults: @@ -102,7 +111,7 @@ jobs: steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -118,22 +127,19 @@ jobs: python utils/print_env.py - name: Run fast PyTorch Pipeline CPU tests - if: ${{ matrix.config.framework == 'pytorch_pipelines' }} run: | pytest -n 8 --max-worker-restart=0 --dist=loadfile \ -k "not Flax and not Onnx" \ - --make-reports=tests_${{ matrix.config.report }} \ + --make-reports=tests_torch_cpu_modular_pipelines \ tests/modular_pipelines - name: Failure short reports if: ${{ failure() }} - run: cat reports/tests_${{ matrix.config.report }}_failures_short.txt + run: cat reports/tests_torch_cpu_modular_pipelines_failures_short.txt - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: - name: pr_${{ matrix.config.framework }}_${{ matrix.config.report }}_test_reports + name: pr_pytorch_pipelines_torch_cpu_modular_pipelines_test_reports path: reports - - diff --git a/.github/workflows/pr_test_fetcher.yml b/.github/workflows/pr_test_fetcher.yml index 83b2ab4edbf6..a02a40709fcc 100644 --- a/.github/workflows/pr_test_fetcher.yml +++ b/.github/workflows/pr_test_fetcher.yml @@ -28,7 +28,7 @@ jobs: test_map: ${{ steps.set_matrix.outputs.test_map }} steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 0 - name: Install dependencies @@ -42,7 +42,7 @@ jobs: run: | python utils/tests_fetcher.py | tee test_preparation.txt - name: Report fetched tests - uses: actions/upload-artifact@v3 + uses: actions/upload-artifact@v6 with: name: test_fetched path: test_preparation.txt @@ -83,7 +83,7 @@ jobs: shell: bash steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -109,7 +109,7 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v3 + uses: actions/upload-artifact@v6 with: name: ${{ matrix.modules }}_test_reports path: reports @@ -138,7 +138,7 @@ jobs: steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -164,7 +164,7 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: pr_${{ matrix.config.report }}_test_reports path: reports diff --git a/.github/workflows/pr_tests.yml b/.github/workflows/pr_tests.yml index 674e62ff443a..3e69dfd05eae 100644 --- a/.github/workflows/pr_tests.yml +++ b/.github/workflows/pr_tests.yml @@ -31,11 +31,11 @@ jobs: check_code_quality: runs-on: ubuntu-22.04 steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: - python-version: "3.8" + python-version: "3.10" - name: Install dependencies run: | pip install --upgrade pip @@ -51,11 +51,11 @@ jobs: needs: check_code_quality runs-on: ubuntu-22.04 steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: - python-version: "3.8" + python-version: "3.10" - name: Install dependencies run: | pip install --upgrade pip @@ -92,7 +92,6 @@ jobs: runner: aws-general-8-plus image: diffusers/diffusers-pytorch-cpu report: torch_example_cpu - name: ${{ matrix.config.name }} runs-on: @@ -108,15 +107,14 @@ jobs: steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 - name: Install dependencies run: | uv pip install -e ".[quality]" - #uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - uv pip uninstall transformers huggingface_hub && uv pip install transformers==4.57.1 + uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps - name: Environment @@ -153,7 +151,7 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: pr_${{ matrix.config.framework }}_${{ matrix.config.report }}_test_reports path: reports @@ -185,7 +183,7 @@ jobs: steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -211,15 +209,13 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: pr_${{ matrix.config.report }}_test_reports path: reports run_lora_tests: needs: [check_code_quality, check_repository_consistency] - strategy: - fail-fast: false name: LoRA tests with PEFT main @@ -236,7 +232,7 @@ jobs: steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -247,9 +243,8 @@ jobs: uv pip install -U peft@git+https://github.com/huggingface/peft.git --no-deps uv pip install -U tokenizers uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps - #uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - uv pip uninstall transformers huggingface_hub && uv pip install transformers==4.57.1 - + uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git + - name: Environment run: | python utils/print_env.py @@ -273,8 +268,8 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: - name: pr_main_test_reports + name: pr_lora_test_reports path: reports diff --git a/.github/workflows/pr_tests_gpu.yml b/.github/workflows/pr_tests_gpu.yml index 468979d379c1..b79d80f71c09 100644 --- a/.github/workflows/pr_tests_gpu.yml +++ b/.github/workflows/pr_tests_gpu.yml @@ -32,11 +32,11 @@ jobs: check_code_quality: runs-on: ubuntu-22.04 steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: - python-version: "3.8" + python-version: "3.10" - name: Install dependencies run: | pip install --upgrade pip @@ -52,11 +52,11 @@ jobs: needs: check_code_quality runs-on: ubuntu-22.04 steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: - python-version: "3.8" + python-version: "3.10" - name: Install dependencies run: | pip install --upgrade pip @@ -83,7 +83,7 @@ jobs: pipeline_test_matrix: ${{ steps.fetch_pipeline_matrix.outputs.pipeline_test_matrix }} steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 - name: Install dependencies @@ -100,7 +100,7 @@ jobs: echo "pipeline_test_matrix=$matrix" >> $GITHUB_OUTPUT - name: Pipeline Tests Artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: test-pipelines.json path: reports @@ -120,7 +120,7 @@ jobs: options: --shm-size "16gb" --ipc host --gpus all steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -131,8 +131,7 @@ jobs: run: | uv pip install -e ".[quality]" uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git - #uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - uv pip uninstall transformers huggingface_hub && uv pip install transformers==4.57.1 + uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - name: Environment run: | @@ -170,7 +169,7 @@ jobs: cat reports/tests_pipeline_${{ matrix.module }}_cuda_failures_short.txt - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: pipeline_${{ matrix.module }}_test_reports path: reports @@ -193,7 +192,7 @@ jobs: module: [models, schedulers, lora, others] steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -202,8 +201,7 @@ jobs: uv pip install -e ".[quality]" uv pip install peft@git+https://github.com/huggingface/peft.git uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git - #uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - uv pip uninstall transformers huggingface_hub && uv pip install transformers==4.57.1 + uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - name: Environment run: | @@ -239,7 +237,7 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: torch_cuda_test_reports_${{ matrix.module }} path: reports @@ -255,7 +253,7 @@ jobs: options: --gpus all --shm-size "16gb" --ipc host steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -264,8 +262,7 @@ jobs: nvidia-smi - name: Install dependencies run: | - #uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - uv pip uninstall transformers huggingface_hub && uv pip install transformers==4.57.1 + uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git uv pip install -e ".[quality,training]" - name: Environment @@ -287,7 +284,7 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: examples_test_reports path: reports diff --git a/.github/workflows/pr_torch_dependency_test.yml b/.github/workflows/pr_torch_dependency_test.yml index 4b6160ff71e2..79569488ae21 100644 --- a/.github/workflows/pr_torch_dependency_test.yml +++ b/.github/workflows/pr_torch_dependency_test.yml @@ -18,11 +18,11 @@ jobs: check_torch_dependencies: runs-on: ubuntu-22.04 steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: - python-version: "3.8" + python-version: "3.10" - name: Install dependencies run: | pip install -e . diff --git a/.github/workflows/push_tests.yml b/.github/workflows/push_tests.yml index 7b1c441d3dc0..61bb9f0ef679 100644 --- a/.github/workflows/push_tests.yml +++ b/.github/workflows/push_tests.yml @@ -29,7 +29,7 @@ jobs: pipeline_test_matrix: ${{ steps.fetch_pipeline_matrix.outputs.pipeline_test_matrix }} steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 - name: Install dependencies @@ -46,7 +46,7 @@ jobs: echo "pipeline_test_matrix=$matrix" >> $GITHUB_OUTPUT - name: Pipeline Tests Artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: test-pipelines.json path: reports @@ -66,7 +66,7 @@ jobs: options: --shm-size "16gb" --ipc host --gpus all steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 - name: NVIDIA-SMI @@ -76,8 +76,7 @@ jobs: run: | uv pip install -e ".[quality]" uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git - #uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - uv pip uninstall transformers huggingface_hub && uv pip install transformers==4.57.1 + uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - name: Environment run: | python utils/print_env.py @@ -98,7 +97,7 @@ jobs: cat reports/tests_pipeline_${{ matrix.module }}_cuda_failures_short.txt - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: pipeline_${{ matrix.module }}_test_reports path: reports @@ -120,7 +119,7 @@ jobs: module: [models, schedulers, lora, others, single_file] steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -129,8 +128,7 @@ jobs: uv pip install -e ".[quality]" uv pip install peft@git+https://github.com/huggingface/peft.git uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git - #uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - uv pip uninstall transformers huggingface_hub && uv pip install transformers==4.57.1 + uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - name: Environment run: | @@ -155,7 +153,7 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: torch_cuda_test_reports_${{ matrix.module }} path: reports @@ -172,7 +170,7 @@ jobs: steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -182,8 +180,7 @@ jobs: - name: Install dependencies run: | uv pip install -e ".[quality,training]" - #uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - uv pip uninstall transformers huggingface_hub && uv pip install transformers==4.57.1 + uv pip uninstall transformers huggingface_hub && uv pip install --prerelease allow -U transformers@git+https://github.com/huggingface/transformers.git - name: Environment run: | python utils/print_env.py @@ -199,7 +196,7 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: torch_compile_test_reports path: reports @@ -216,7 +213,7 @@ jobs: steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -240,7 +237,7 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: torch_xformers_test_reports path: reports @@ -256,7 +253,7 @@ jobs: options: --gpus all --shm-size "16gb" --ipc host steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -286,7 +283,7 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: examples_test_reports path: reports diff --git a/.github/workflows/push_tests_fast.yml b/.github/workflows/push_tests_fast.yml index 38cbffaa6315..fe6f6a265e89 100644 --- a/.github/workflows/push_tests_fast.yml +++ b/.github/workflows/push_tests_fast.yml @@ -54,7 +54,7 @@ jobs: steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -88,7 +88,7 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: pr_${{ matrix.config.report }}_test_reports path: reports diff --git a/.github/workflows/push_tests_mps.yml b/.github/workflows/push_tests_mps.yml index 2d6feb592815..7f8ce9a4b99d 100644 --- a/.github/workflows/push_tests_mps.yml +++ b/.github/workflows/push_tests_mps.yml @@ -23,7 +23,7 @@ jobs: steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -41,7 +41,7 @@ jobs: shell: arch -arch arm64 bash {0} run: | ${CONDA_RUN} python -m pip install --upgrade pip uv - ${CONDA_RUN} python -m uv pip install -e ".[quality,test]" + ${CONDA_RUN} python -m uv pip install -e ".[quality]" ${CONDA_RUN} python -m uv pip install torch torchvision torchaudio ${CONDA_RUN} python -m uv pip install accelerate@git+https://github.com/huggingface/accelerate.git ${CONDA_RUN} python -m uv pip install transformers --upgrade @@ -65,7 +65,7 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: pr_torch_mps_test_reports path: reports diff --git a/.github/workflows/pypi_publish.yaml b/.github/workflows/pypi_publish.yaml index dc36b6b024c5..99cbcc1fade2 100644 --- a/.github/workflows/pypi_publish.yaml +++ b/.github/workflows/pypi_publish.yaml @@ -15,12 +15,12 @@ jobs: latest_branch: ${{ steps.set_latest_branch.outputs.latest_branch }} steps: - name: Checkout Repo - uses: actions/checkout@v3 + uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: - python-version: '3.8' + python-version: '3.10' - name: Fetch latest branch id: fetch_latest_branch @@ -40,21 +40,20 @@ jobs: steps: - name: Checkout Repo - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: ref: ${{ needs.find-and-checkout-latest-branch.outputs.latest_branch }} - name: Setup Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: - python-version: "3.8" + python-version: "3.10" - name: Install dependencies run: | python -m pip install --upgrade pip pip install -U setuptools wheel twine pip install -U torch --index-url https://download.pytorch.org/whl/cpu - pip install -U transformers - name: Build the dist files run: python setup.py bdist_wheel && python setup.py sdist @@ -69,6 +68,8 @@ jobs: run: | pip install diffusers && pip uninstall diffusers -y pip install -i https://test.pypi.org/simple/ diffusers + pip install -U transformers + python utils/print_env.py python -c "from diffusers import __version__; print(__version__)" python -c "from diffusers import DiffusionPipeline; pipe = DiffusionPipeline.from_pretrained('fusing/unet-ldm-dummy-update'); pipe()" python -c "from diffusers import DiffusionPipeline; pipe = DiffusionPipeline.from_pretrained('hf-internal-testing/tiny-stable-diffusion-pipe', safety_checker=None); pipe('ah suh du')" diff --git a/.github/workflows/release_tests_fast.yml b/.github/workflows/release_tests_fast.yml index efdd6ea2b651..f667d715090d 100644 --- a/.github/workflows/release_tests_fast.yml +++ b/.github/workflows/release_tests_fast.yml @@ -27,7 +27,7 @@ jobs: pipeline_test_matrix: ${{ steps.fetch_pipeline_matrix.outputs.pipeline_test_matrix }} steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 - name: Install dependencies @@ -44,7 +44,7 @@ jobs: echo "pipeline_test_matrix=$matrix" >> $GITHUB_OUTPUT - name: Pipeline Tests Artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: test-pipelines.json path: reports @@ -64,7 +64,7 @@ jobs: options: --shm-size "16gb" --ipc host --gpus all steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 - name: NVIDIA-SMI @@ -94,7 +94,7 @@ jobs: cat reports/tests_pipeline_${{ matrix.module }}_cuda_failures_short.txt - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: pipeline_${{ matrix.module }}_test_reports path: reports @@ -116,7 +116,7 @@ jobs: module: [models, schedulers, lora, others, single_file] steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -149,7 +149,7 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: torch_cuda_${{ matrix.module }}_test_reports path: reports @@ -166,7 +166,7 @@ jobs: shell: bash steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -205,7 +205,7 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: torch_minimum_version_cuda_test_reports path: reports @@ -222,7 +222,7 @@ jobs: steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -247,7 +247,7 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: torch_compile_test_reports path: reports @@ -264,7 +264,7 @@ jobs: steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -288,7 +288,7 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: torch_xformers_test_reports path: reports @@ -305,7 +305,7 @@ jobs: steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 @@ -336,7 +336,7 @@ jobs: - name: Test suite reports artifacts if: ${{ always() }} - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: examples_test_reports path: reports diff --git a/.github/workflows/run_tests_from_a_pr.yml b/.github/workflows/run_tests_from_a_pr.yml index fa8c579dd768..3e5462f5100f 100644 --- a/.github/workflows/run_tests_from_a_pr.yml +++ b/.github/workflows/run_tests_from_a_pr.yml @@ -57,7 +57,7 @@ jobs: shell: bash -e {0} - name: Checkout PR branch - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: ref: refs/pull/${{ inputs.pr_number }}/head diff --git a/.github/workflows/ssh-pr-runner.yml b/.github/workflows/ssh-pr-runner.yml index 49fa9c0ad24d..27246fb61348 100644 --- a/.github/workflows/ssh-pr-runner.yml +++ b/.github/workflows/ssh-pr-runner.yml @@ -27,7 +27,7 @@ jobs: steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 diff --git a/.github/workflows/ssh-runner.yml b/.github/workflows/ssh-runner.yml index 917eb5b1b31a..4fbfad3dc7c6 100644 --- a/.github/workflows/ssh-runner.yml +++ b/.github/workflows/ssh-runner.yml @@ -35,7 +35,7 @@ jobs: steps: - name: Checkout diffusers - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: fetch-depth: 2 diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml index 27450ed4c7f2..76dd48d09931 100644 --- a/.github/workflows/stale.yml +++ b/.github/workflows/stale.yml @@ -15,12 +15,12 @@ jobs: env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v6 - name: Setup Python - uses: actions/setup-python@v1 + uses: actions/setup-python@v6 with: - python-version: 3.8 + python-version: 3.10 - name: Install requirements run: | diff --git a/.github/workflows/trufflehog.yml b/.github/workflows/trufflehog.yml index 4743dc352455..65334e086c83 100644 --- a/.github/workflows/trufflehog.yml +++ b/.github/workflows/trufflehog.yml @@ -8,7 +8,7 @@ jobs: runs-on: ubuntu-22.04 steps: - name: Checkout code - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: fetch-depth: 0 - name: Secret Scanning diff --git a/.github/workflows/typos.yml b/.github/workflows/typos.yml index 6d2f2fc8dd9a..87ea38a5bbac 100644 --- a/.github/workflows/typos.yml +++ b/.github/workflows/typos.yml @@ -8,7 +8,7 @@ jobs: runs-on: ubuntu-22.04 steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 - name: typos-action - uses: crate-ci/typos@v1.12.4 + uses: crate-ci/typos@v1.42.1 diff --git a/.github/workflows/update_metadata.yml b/.github/workflows/update_metadata.yml index 92aea0369ba8..6e608883c13a 100644 --- a/.github/workflows/update_metadata.yml +++ b/.github/workflows/update_metadata.yml @@ -15,7 +15,7 @@ jobs: shell: bash -l {0} steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 - name: Setup environment run: | diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md deleted file mode 100644 index ec18df882641..000000000000 --- a/CONTRIBUTING.md +++ /dev/null @@ -1,506 +0,0 @@ - - -# How to contribute to Diffusers 🧨 - -We ❤️ contributions from the open-source community! Everyone is welcome, and all types of participation –not just code– are valued and appreciated. Answering questions, helping others, reaching out, and improving the documentation are all immensely valuable to the community, so don't be afraid and get involved if you're up for it! - -Everyone is encouraged to start by saying 👋 in our public Discord channel. We discuss the latest trends in diffusion models, ask questions, show off personal projects, help each other with contributions, or just hang out ☕. Join us on Discord - -Whichever way you choose to contribute, we strive to be part of an open, welcoming, and kind community. Please, read our [code of conduct](https://github.com/huggingface/diffusers/blob/main/CODE_OF_CONDUCT.md) and be mindful to respect it during your interactions. We also recommend you become familiar with the [ethical guidelines](https://huggingface.co/docs/diffusers/conceptual/ethical_guidelines) that guide our project and ask you to adhere to the same principles of transparency and responsibility. - -We enormously value feedback from the community, so please do not be afraid to speak up if you believe you have valuable feedback that can help improve the library - every message, comment, issue, and pull request (PR) is read and considered. - -## Overview - -You can contribute in many ways ranging from answering questions on issues to adding new diffusion models to -the core library. - -In the following, we give an overview of different ways to contribute, ranked by difficulty in ascending order. All of them are valuable to the community. - -* 1. Asking and answering questions on [the Diffusers discussion forum](https://discuss.huggingface.co/c/discussion-related-to-httpsgithubcomhuggingfacediffusers) or on [Discord](https://discord.gg/G7tWnz98XR). -* 2. Opening new issues on [the GitHub Issues tab](https://github.com/huggingface/diffusers/issues/new/choose). -* 3. Answering issues on [the GitHub Issues tab](https://github.com/huggingface/diffusers/issues). -* 4. Fix a simple issue, marked by the "Good first issue" label, see [here](https://github.com/huggingface/diffusers/issues?q=is%3Aopen+is%3Aissue+label%3A%22good+first+issue%22). -* 5. Contribute to the [documentation](https://github.com/huggingface/diffusers/tree/main/docs/source). -* 6. Contribute a [Community Pipeline](https://github.com/huggingface/diffusers/issues?q=is%3Aopen+is%3Aissue+label%3Acommunity-examples). -* 7. Contribute to the [examples](https://github.com/huggingface/diffusers/tree/main/examples). -* 8. Fix a more difficult issue, marked by the "Good second issue" label, see [here](https://github.com/huggingface/diffusers/issues?q=is%3Aopen+is%3Aissue+label%3A%22Good+second+issue%22). -* 9. Add a new pipeline, model, or scheduler, see ["New Pipeline/Model"](https://github.com/huggingface/diffusers/issues?q=is%3Aopen+is%3Aissue+label%3A%22New+pipeline%2Fmodel%22) and ["New scheduler"](https://github.com/huggingface/diffusers/issues?q=is%3Aopen+is%3Aissue+label%3A%22New+scheduler%22) issues. For this contribution, please have a look at [Design Philosophy](https://github.com/huggingface/diffusers/blob/main/PHILOSOPHY.md). - -As said before, **all contributions are valuable to the community**. -In the following, we will explain each contribution a bit more in detail. - -For all contributions 4-9, you will need to open a PR. It is explained in detail how to do so in [Opening a pull request](#how-to-open-a-pr). - -### 1. Asking and answering questions on the Diffusers discussion forum or on the Diffusers Discord - -Any question or comment related to the Diffusers library can be asked on the [discussion forum](https://discuss.huggingface.co/c/discussion-related-to-httpsgithubcomhuggingfacediffusers/) or on [Discord](https://discord.gg/G7tWnz98XR). Such questions and comments include (but are not limited to): -- Reports of training or inference experiments in an attempt to share knowledge -- Presentation of personal projects -- Questions to non-official training examples -- Project proposals -- General feedback -- Paper summaries -- Asking for help on personal projects that build on top of the Diffusers library -- General questions -- Ethical questions regarding diffusion models -- ... - -Every question that is asked on the forum or on Discord actively encourages the community to publicly -share knowledge and might very well help a beginner in the future who has the same question you're -having. Please do pose any questions you might have. -In the same spirit, you are of immense help to the community by answering such questions because this way you are publicly documenting knowledge for everybody to learn from. - -**Please** keep in mind that the more effort you put into asking or answering a question, the higher -the quality of the publicly documented knowledge. In the same way, well-posed and well-answered questions create a high-quality knowledge database accessible to everybody, while badly posed questions or answers reduce the overall quality of the public knowledge database. -In short, a high quality question or answer is *precise*, *concise*, *relevant*, *easy-to-understand*, *accessible*, and *well-formatted/well-posed*. For more information, please have a look through the [How to write a good issue](#how-to-write-a-good-issue) section. - -**NOTE about channels**: -[*The forum*](https://discuss.huggingface.co/c/discussion-related-to-httpsgithubcomhuggingfacediffusers/63) is much better indexed by search engines, such as Google. Posts are ranked by popularity rather than chronologically. Hence, it's easier to look up questions and answers that we posted some time ago. -In addition, questions and answers posted in the forum can easily be linked to. -In contrast, *Discord* has a chat-like format that invites fast back-and-forth communication. -While it will most likely take less time for you to get an answer to your question on Discord, your -question won't be visible anymore over time. Also, it's much harder to find information that was posted a while back on Discord. We therefore strongly recommend using the forum for high-quality questions and answers in an attempt to create long-lasting knowledge for the community. If discussions on Discord lead to very interesting answers and conclusions, we recommend posting the results on the forum to make the information more available for future readers. - -### 2. Opening new issues on the GitHub issues tab - -The 🧨 Diffusers library is robust and reliable thanks to the users who notify us of -the problems they encounter. So thank you for reporting an issue. - -Remember, GitHub issues are reserved for technical questions directly related to the Diffusers library, bug reports, feature requests, or feedback on the library design. - -In a nutshell, this means that everything that is **not** related to the **code of the Diffusers library** (including the documentation) should **not** be asked on GitHub, but rather on either the [forum](https://discuss.huggingface.co/c/discussion-related-to-httpsgithubcomhuggingfacediffusers/63) or [Discord](https://discord.gg/G7tWnz98XR). - -**Please consider the following guidelines when opening a new issue**: -- Make sure you have searched whether your issue has already been asked before (use the search bar on GitHub under Issues). -- Please never report a new issue on another (related) issue. If another issue is highly related, please -open a new issue nevertheless and link to the related issue. -- Make sure your issue is written in English. Please use one of the great, free online translation services, such as [DeepL](https://www.deepl.com/translator) to translate from your native language to English if you are not comfortable in English. -- Check whether your issue might be solved by updating to the newest Diffusers version. Before posting your issue, please make sure that `python -c "import diffusers; print(diffusers.__version__)"` is higher or matches the latest Diffusers version. -- Remember that the more effort you put into opening a new issue, the higher the quality of your answer will be and the better the overall quality of the Diffusers issues. - -New issues usually include the following. - -#### 2.1. Reproducible, minimal bug reports - -A bug report should always have a reproducible code snippet and be as minimal and concise as possible. -This means in more detail: -- Narrow the bug down as much as you can, **do not just dump your whole code file**. -- Format your code. -- Do not include any external libraries except for Diffusers depending on them. -- **Always** provide all necessary information about your environment; for this, you can run: `diffusers-cli env` in your shell and copy-paste the displayed information to the issue. -- Explain the issue. If the reader doesn't know what the issue is and why it is an issue, she cannot solve it. -- **Always** make sure the reader can reproduce your issue with as little effort as possible. If your code snippet cannot be run because of missing libraries or undefined variables, the reader cannot help you. Make sure your reproducible code snippet is as minimal as possible and can be copy-pasted into a simple Python shell. -- If in order to reproduce your issue a model and/or dataset is required, make sure the reader has access to that model or dataset. You can always upload your model or dataset to the [Hub](https://huggingface.co) to make it easily downloadable. Try to keep your model and dataset as small as possible, to make the reproduction of your issue as effortless as possible. - -For more information, please have a look through the [How to write a good issue](#how-to-write-a-good-issue) section. - -You can open a bug report [here](https://github.com/huggingface/diffusers/issues/new?assignees=&labels=bug&projects=&template=bug-report.yml). - -#### 2.2. Feature requests - -A world-class feature request addresses the following points: - -1. Motivation first: -* Is it related to a problem/frustration with the library? If so, please explain -why. Providing a code snippet that demonstrates the problem is best. -* Is it related to something you would need for a project? We'd love to hear -about it! -* Is it something you worked on and think could benefit the community? -Awesome! Tell us what problem it solved for you. -2. Write a *full paragraph* describing the feature; -3. Provide a **code snippet** that demonstrates its future use; -4. In case this is related to a paper, please attach a link; -5. Attach any additional information (drawings, screenshots, etc.) you think may help. - -You can open a feature request [here](https://github.com/huggingface/diffusers/issues/new?assignees=&labels=&template=feature_request.md&title=). - -#### 2.3 Feedback - -Feedback about the library design and why it is good or not good helps the core maintainers immensely to build a user-friendly library. To understand the philosophy behind the current design philosophy, please have a look [here](https://huggingface.co/docs/diffusers/conceptual/philosophy). If you feel like a certain design choice does not fit with the current design philosophy, please explain why and how it should be changed. If a certain design choice follows the design philosophy too much, hence restricting use cases, explain why and how it should be changed. -If a certain design choice is very useful for you, please also leave a note as this is great feedback for future design decisions. - -You can open an issue about feedback [here](https://github.com/huggingface/diffusers/issues/new?assignees=&labels=&template=feedback.md&title=). - -#### 2.4 Technical questions - -Technical questions are mainly about why certain code of the library was written in a certain way, or what a certain part of the code does. Please make sure to link to the code in question and please provide detail on -why this part of the code is difficult to understand. - -You can open an issue about a technical question [here](https://github.com/huggingface/diffusers/issues/new?assignees=&labels=bug&template=bug-report.yml). - -#### 2.5 Proposal to add a new model, scheduler, or pipeline - -If the diffusion model community released a new model, pipeline, or scheduler that you would like to see in the Diffusers library, please provide the following information: - -* Short description of the diffusion pipeline, model, or scheduler and link to the paper or public release. -* Link to any of its open-source implementation. -* Link to the model weights if they are available. - -If you are willing to contribute to the model yourself, let us know so we can best guide you. Also, don't forget -to tag the original author of the component (model, scheduler, pipeline, etc.) by GitHub handle if you can find it. - -You can open a request for a model/pipeline/scheduler [here](https://github.com/huggingface/diffusers/issues/new?assignees=&labels=New+model%2Fpipeline%2Fscheduler&template=new-model-addition.yml). - -### 3. Answering issues on the GitHub issues tab - -Answering issues on GitHub might require some technical knowledge of Diffusers, but we encourage everybody to give it a try even if you are not 100% certain that your answer is correct. -Some tips to give a high-quality answer to an issue: -- Be as concise and minimal as possible. -- Stay on topic. An answer to the issue should concern the issue and only the issue. -- Provide links to code, papers, or other sources that prove or encourage your point. -- Answer in code. If a simple code snippet is the answer to the issue or shows how the issue can be solved, please provide a fully reproducible code snippet. - -Also, many issues tend to be simply off-topic, duplicates of other issues, or irrelevant. It is of great -help to the maintainers if you can answer such issues, encouraging the author of the issue to be -more precise, provide the link to a duplicated issue or redirect them to [the forum](https://discuss.huggingface.co/c/discussion-related-to-httpsgithubcomhuggingfacediffusers/63) or [Discord](https://discord.gg/G7tWnz98XR). - -If you have verified that the issued bug report is correct and requires a correction in the source code, -please have a look at the next sections. - -For all of the following contributions, you will need to open a PR. It is explained in detail how to do so in the [Opening a pull request](#how-to-open-a-pr) section. - -### 4. Fixing a "Good first issue" - -*Good first issues* are marked by the [Good first issue](https://github.com/huggingface/diffusers/issues?q=is%3Aopen+is%3Aissue+label%3A%22good+first+issue%22) label. Usually, the issue already -explains how a potential solution should look so that it is easier to fix. -If the issue hasn't been closed and you would like to try to fix this issue, you can just leave a message "I would like to try this issue.". There are usually three scenarios: -- a.) The issue description already proposes a fix. In this case and if the solution makes sense to you, you can open a PR or draft PR to fix it. -- b.) The issue description does not propose a fix. In this case, you can ask what a proposed fix could look like and someone from the Diffusers team should answer shortly. If you have a good idea of how to fix it, feel free to directly open a PR. -- c.) There is already an open PR to fix the issue, but the issue hasn't been closed yet. If the PR has gone stale, you can simply open a new PR and link to the stale PR. PRs often go stale if the original contributor who wanted to fix the issue suddenly cannot find the time anymore to proceed. This often happens in open-source and is very normal. In this case, the community will be very happy if you give it a new try and leverage the knowledge of the existing PR. If there is already a PR and it is active, you can help the author by giving suggestions, reviewing the PR or even asking whether you can contribute to the PR. - - -### 5. Contribute to the documentation - -A good library **always** has good documentation! The official documentation is often one of the first points of contact for new users of the library, and therefore contributing to the documentation is a **highly -valuable contribution**. - -Contributing to the library can have many forms: - -- Correcting spelling or grammatical errors. -- Correct incorrect formatting of the docstring. If you see that the official documentation is weirdly displayed or a link is broken, we are very happy if you take some time to correct it. -- Correct the shape or dimensions of a docstring input or output tensor. -- Clarify documentation that is hard to understand or incorrect. -- Update outdated code examples. -- Translating the documentation to another language. - -Anything displayed on [the official Diffusers doc page](https://huggingface.co/docs/diffusers/index) is part of the official documentation and can be corrected, adjusted in the respective [documentation source](https://github.com/huggingface/diffusers/tree/main/docs/source). - -Please have a look at [this page](https://github.com/huggingface/diffusers/tree/main/docs) on how to verify changes made to the documentation locally. - - -### 6. Contribute a community pipeline - -[Pipelines](https://huggingface.co/docs/diffusers/api/pipelines/overview) are usually the first point of contact between the Diffusers library and the user. -Pipelines are examples of how to use Diffusers [models](https://huggingface.co/docs/diffusers/api/models/overview) and [schedulers](https://huggingface.co/docs/diffusers/api/schedulers/overview). -We support two types of pipelines: - -- Official Pipelines -- Community Pipelines - -Both official and community pipelines follow the same design and consist of the same type of components. - -Official pipelines are tested and maintained by the core maintainers of Diffusers. Their code -resides in [src/diffusers/pipelines](https://github.com/huggingface/diffusers/tree/main/src/diffusers/pipelines). -In contrast, community pipelines are contributed and maintained purely by the **community** and are **not** tested. -They reside in [examples/community](https://github.com/huggingface/diffusers/tree/main/examples/community) and while they can be accessed via the [PyPI diffusers package](https://pypi.org/project/diffusers/), their code is not part of the PyPI distribution. - -The reason for the distinction is that the core maintainers of the Diffusers library cannot maintain and test all -possible ways diffusion models can be used for inference, but some of them may be of interest to the community. -Officially released diffusion pipelines, -such as Stable Diffusion are added to the core src/diffusers/pipelines package which ensures -high quality of maintenance, no backward-breaking code changes, and testing. -More bleeding edge pipelines should be added as community pipelines. If usage for a community pipeline is high, the pipeline can be moved to the official pipelines upon request from the community. This is one of the ways we strive to be a community-driven library. - -To add a community pipeline, one should add a .py file to [examples/community](https://github.com/huggingface/diffusers/tree/main/examples/community) and adapt the [examples/community/README.md](https://github.com/huggingface/diffusers/tree/main/examples/community/README.md) to include an example of the new pipeline. - -An example can be seen [here](https://github.com/huggingface/diffusers/pull/2400). - -Community pipeline PRs are only checked at a superficial level and ideally they should be maintained by their original authors. - -Contributing a community pipeline is a great way to understand how Diffusers models and schedulers work. Having contributed a community pipeline is usually the first stepping stone to contributing an official pipeline to the -core package. - -### 7. Contribute to training examples - -Diffusers examples are a collection of training scripts that reside in [examples](https://github.com/huggingface/diffusers/tree/main/examples). - -We support two types of training examples: - -- Official training examples -- Research training examples - -Research training examples are located in [examples/research_projects](https://github.com/huggingface/diffusers/tree/main/examples/research_projects) whereas official training examples include all folders under [examples](https://github.com/huggingface/diffusers/tree/main/examples) except the `research_projects` and `community` folders. -The official training examples are maintained by the Diffusers' core maintainers whereas the research training examples are maintained by the community. -This is because of the same reasons put forward in [6. Contribute a community pipeline](#6-contribute-a-community-pipeline) for official pipelines vs. community pipelines: It is not feasible for the core maintainers to maintain all possible training methods for diffusion models. -If the Diffusers core maintainers and the community consider a certain training paradigm to be too experimental or not popular enough, the corresponding training code should be put in the `research_projects` folder and maintained by the author. - -Both official training and research examples consist of a directory that contains one or more training scripts, a `requirements.txt` file, and a `README.md` file. In order for the user to make use of the -training examples, it is required to clone the repository: - -```bash -git clone https://github.com/huggingface/diffusers -``` - -as well as to install all additional dependencies required for training: - -```bash -cd diffusers -pip install -r examples//requirements.txt -``` - -Therefore when adding an example, the `requirements.txt` file shall define all pip dependencies required for your training example so that once all those are installed, the user can run the example's training script. See, for example, the [DreamBooth `requirements.txt` file](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/requirements.txt). - -Training examples of the Diffusers library should adhere to the following philosophy: -- All the code necessary to run the examples should be found in a single Python file. -- One should be able to run the example from the command line with `python .py --args`. -- Examples should be kept simple and serve as **an example** on how to use Diffusers for training. The purpose of example scripts is **not** to create state-of-the-art diffusion models, but rather to reproduce known training schemes without adding too much custom logic. As a byproduct of this point, our examples also strive to serve as good educational materials. - -To contribute an example, it is highly recommended to look at already existing examples such as [dreambooth](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth.py) to get an idea of how they should look like. -We strongly advise contributors to make use of the [Accelerate library](https://github.com/huggingface/accelerate) as it's tightly integrated -with Diffusers. -Once an example script works, please make sure to add a comprehensive `README.md` that states how to use the example exactly. This README should include: -- An example command on how to run the example script as shown [here e.g.](https://github.com/huggingface/diffusers/tree/main/examples/dreambooth#running-locally-with-pytorch). -- A link to some training results (logs, models, ...) that show what the user can expect as shown [here e.g.](https://api.wandb.ai/report/patrickvonplaten/xm6cd5q5). -- If you are adding a non-official/research training example, **please don't forget** to add a sentence that you are maintaining this training example which includes your git handle as shown [here](https://github.com/huggingface/diffusers/tree/main/examples/research_projects/intel_opts#diffusers-examples-with-intel-optimizations). - -If you are contributing to the official training examples, please also make sure to add a test to [examples/test_examples.py](https://github.com/huggingface/diffusers/blob/main/examples/test_examples.py). This is not necessary for non-official training examples. - -### 8. Fixing a "Good second issue" - -*Good second issues* are marked by the [Good second issue](https://github.com/huggingface/diffusers/issues?q=is%3Aopen+is%3Aissue+label%3A%22Good+second+issue%22) label. Good second issues are -usually more complicated to solve than [Good first issues](https://github.com/huggingface/diffusers/issues?q=is%3Aopen+is%3Aissue+label%3A%22good+first+issue%22). -The issue description usually gives less guidance on how to fix the issue and requires -a decent understanding of the library by the interested contributor. -If you are interested in tackling a good second issue, feel free to open a PR to fix it and link the PR to the issue. If you see that a PR has already been opened for this issue but did not get merged, have a look to understand why it wasn't merged and try to open an improved PR. -Good second issues are usually more difficult to get merged compared to good first issues, so don't hesitate to ask for help from the core maintainers. If your PR is almost finished the core maintainers can also jump into your PR and commit to it in order to get it merged. - -### 9. Adding pipelines, models, schedulers - -Pipelines, models, and schedulers are the most important pieces of the Diffusers library. -They provide easy access to state-of-the-art diffusion technologies and thus allow the community to -build powerful generative AI applications. - -By adding a new model, pipeline, or scheduler you might enable a new powerful use case for any of the user interfaces relying on Diffusers which can be of immense value for the whole generative AI ecosystem. - -Diffusers has a couple of open feature requests for all three components - feel free to gloss over them -if you don't know yet what specific component you would like to add: -- [Model or pipeline](https://github.com/huggingface/diffusers/issues?q=is%3Aopen+is%3Aissue+label%3A%22New+pipeline%2Fmodel%22) -- [Scheduler](https://github.com/huggingface/diffusers/issues?q=is%3Aopen+is%3Aissue+label%3A%22New+scheduler%22) - -Before adding any of the three components, it is strongly recommended that you give the [Philosophy guide](https://github.com/huggingface/diffusers/blob/main/PHILOSOPHY.md) a read to better understand the design of any of the three components. Please be aware that -we cannot merge model, scheduler, or pipeline additions that strongly diverge from our design philosophy -as it will lead to API inconsistencies. If you fundamentally disagree with a design choice, please -open a [Feedback issue](https://github.com/huggingface/diffusers/issues/new?assignees=&labels=&template=feedback.md&title=) instead so that it can be discussed whether a certain design -pattern/design choice shall be changed everywhere in the library and whether we shall update our design philosophy. Consistency across the library is very important for us. - -Please make sure to add links to the original codebase/paper to the PR and ideally also ping the -original author directly on the PR so that they can follow the progress and potentially help with questions. - -If you are unsure or stuck in the PR, don't hesitate to leave a message to ask for a first review or help. - -## How to write a good issue - -**The better your issue is written, the higher the chances that it will be quickly resolved.** - -1. Make sure that you've used the correct template for your issue. You can pick between *Bug Report*, *Feature Request*, *Feedback about API Design*, *New model/pipeline/scheduler addition*, *Forum*, or a blank issue. Make sure to pick the correct one when opening [a new issue](https://github.com/huggingface/diffusers/issues/new/choose). -2. **Be precise**: Give your issue a fitting title. Try to formulate your issue description as simple as possible. The more precise you are when submitting an issue, the less time it takes to understand the issue and potentially solve it. Make sure to open an issue for one issue only and not for multiple issues. If you found multiple issues, simply open multiple issues. If your issue is a bug, try to be as precise as possible about what bug it is - you should not just write "Error in diffusers". -3. **Reproducibility**: No reproducible code snippet == no solution. If you encounter a bug, maintainers **have to be able to reproduce** it. Make sure that you include a code snippet that can be copy-pasted into a Python interpreter to reproduce the issue. Make sure that your code snippet works, *i.e.* that there are no missing imports or missing links to images, ... Your issue should contain an error message **and** a code snippet that can be copy-pasted without any changes to reproduce the exact same error message. If your issue is using local model weights or local data that cannot be accessed by the reader, the issue cannot be solved. If you cannot share your data or model, try to make a dummy model or dummy data. -4. **Minimalistic**: Try to help the reader as much as you can to understand the issue as quickly as possible by staying as concise as possible. Remove all code / all information that is irrelevant to the issue. If you have found a bug, try to create the easiest code example you can to demonstrate your issue, do not just dump your whole workflow into the issue as soon as you have found a bug. E.g., if you train a model and get an error at some point during the training, you should first try to understand what part of the training code is responsible for the error and try to reproduce it with a couple of lines. Try to use dummy data instead of full datasets. -5. Add links. If you are referring to a certain naming, method, or model make sure to provide a link so that the reader can better understand what you mean. If you are referring to a specific PR or issue, make sure to link it to your issue. Do not assume that the reader knows what you are talking about. The more links you add to your issue the better. -6. Formatting. Make sure to nicely format your issue by formatting code into Python code syntax, and error messages into normal code syntax. See the [official GitHub formatting docs](https://docs.github.com/en/get-started/writing-on-github/getting-started-with-writing-and-formatting-on-github/basic-writing-and-formatting-syntax) for more information. -7. Think of your issue not as a ticket to be solved, but rather as a beautiful entry to a well-written encyclopedia. Every added issue is a contribution to publicly available knowledge. By adding a nicely written issue you not only make it easier for maintainers to solve your issue, but you are helping the whole community to better understand a certain aspect of the library. - -## How to write a good PR - -1. Be a chameleon. Understand existing design patterns and syntax and make sure your code additions flow seamlessly into the existing code base. Pull requests that significantly diverge from existing design patterns or user interfaces will not be merged. -2. Be laser focused. A pull request should solve one problem and one problem only. Make sure to not fall into the trap of "also fixing another problem while we're adding it". It is much more difficult to review pull requests that solve multiple, unrelated problems at once. -3. If helpful, try to add a code snippet that displays an example of how your addition can be used. -4. The title of your pull request should be a summary of its contribution. -5. If your pull request addresses an issue, please mention the issue number in -the pull request description to make sure they are linked (and people -consulting the issue know you are working on it); -6. To indicate a work in progress please prefix the title with `[WIP]`. These -are useful to avoid duplicated work, and to differentiate it from PRs ready -to be merged; -7. Try to formulate and format your text as explained in [How to write a good issue](#how-to-write-a-good-issue). -8. Make sure existing tests pass; -9. Add high-coverage tests. No quality testing = no merge. -- If you are adding new `@slow` tests, make sure they pass using -`RUN_SLOW=1 python -m pytest tests/test_my_new_model.py`. -CircleCI does not run the slow tests, but GitHub Actions does every night! -10. All public methods must have informative docstrings that work nicely with markdown. See [`pipeline_latent_diffusion.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py) for an example. -11. Due to the rapidly growing repository, it is important to make sure that no files that would significantly weigh down the repository are added. This includes images, videos, and other non-text files. We prefer to leverage a hf.co hosted `dataset` like -[`hf-internal-testing`](https://huggingface.co/hf-internal-testing) or [huggingface/documentation-images](https://huggingface.co/datasets/huggingface/documentation-images) to place these files. -If an external contribution, feel free to add the images to your PR and ask a Hugging Face member to migrate your images -to this dataset. - -## How to open a PR - -Before writing code, we strongly advise you to search through the existing PRs or -issues to make sure that nobody is already working on the same thing. If you are -unsure, it is always a good idea to open an issue to get some feedback. - -You will need basic `git` proficiency to be able to contribute to -🧨 Diffusers. `git` is not the easiest tool to use but it has the greatest -manual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro -Git](https://git-scm.com/book/en/v2) is a very good reference. - -Follow these steps to start contributing ([supported Python versions](https://github.com/huggingface/diffusers/blob/42f25d601a910dceadaee6c44345896b4cfa9928/setup.py#L270)): - -1. Fork the [repository](https://github.com/huggingface/diffusers) by -clicking on the 'Fork' button on the repository's page. This creates a copy of the code -under your GitHub user account. - -2. Clone your fork to your local disk, and add the base repository as a remote: - - ```bash - $ git clone git@github.com:/diffusers.git - $ cd diffusers - $ git remote add upstream https://github.com/huggingface/diffusers.git - ``` - -3. Create a new branch to hold your development changes: - - ```bash - $ git checkout -b a-descriptive-name-for-my-changes - ``` - -**Do not** work on the `main` branch. - -4. Set up a development environment by running the following command in a virtual environment: - - ```bash - $ pip install -e ".[dev]" - ``` - -If you have already cloned the repo, you might need to `git pull` to get the most recent changes in the -library. - -5. Develop the features on your branch. - -As you work on the features, you should make sure that the test suite -passes. You should run the tests impacted by your changes like this: - - ```bash - $ pytest tests/.py - ``` - -Before you run the tests, please make sure you install the dependencies required for testing. You can do so -with this command: - - ```bash - $ pip install -e ".[test]" - ``` - -You can also run the full test suite with the following command, but it takes -a beefy machine to produce a result in a decent amount of time now that -Diffusers has grown a lot. Here is the command for it: - - ```bash - $ make test - ``` - -🧨 Diffusers relies on `ruff` and `isort` to format its source code -consistently. After you make changes, apply automatic style corrections and code verifications -that can't be automated in one go with: - - ```bash - $ make style - ``` - -🧨 Diffusers also uses `ruff` and a few custom scripts to check for coding mistakes. Quality -control runs in CI, however, you can also run the same checks with: - - ```bash - $ make quality - ``` - -Once you're happy with your changes, add changed files using `git add` and -make a commit with `git commit` to record your changes locally: - - ```bash - $ git add modified_file.py - $ git commit -m "A descriptive message about your changes." - ``` - -It is a good idea to sync your copy of the code with the original -repository regularly. This way you can quickly account for changes: - - ```bash - $ git pull upstream main - ``` - -Push the changes to your account using: - - ```bash - $ git push -u origin a-descriptive-name-for-my-changes - ``` - -6. Once you are satisfied, go to the -webpage of your fork on GitHub. Click on 'Pull request' to send your changes -to the project maintainers for review. - -7. It's ok if maintainers ask you for changes. It happens to core contributors -too! So everyone can see the changes in the Pull request, work in your local -branch and push the changes to your fork. They will automatically appear in -the pull request. - -### Tests - -An extensive test suite is included to test the library behavior and several examples. Library tests can be found in -the [tests folder](https://github.com/huggingface/diffusers/tree/main/tests). - -We like `pytest` and `pytest-xdist` because it's faster. From the root of the -repository, here's how to run tests with `pytest` for the library: - -```bash -$ python -m pytest -n auto --dist=loadfile -s -v ./tests/ -``` - -In fact, that's how `make test` is implemented! - -You can specify a smaller set of tests in order to test only the feature -you're working on. - -By default, slow tests are skipped. Set the `RUN_SLOW` environment variable to -`yes` to run them. This will download many gigabytes of models — make sure you -have enough disk space and a good Internet connection, or a lot of patience! - -```bash -$ RUN_SLOW=yes python -m pytest -n auto --dist=loadfile -s -v ./tests/ -``` - -`unittest` is fully supported, here's how to run tests with it: - -```bash -$ python -m unittest discover -s tests -t . -v -$ python -m unittest discover -s examples -t examples -v -``` - -### Syncing forked main with upstream (HuggingFace) main - -To avoid pinging the upstream repository which adds reference notes to each upstream PR and sends unnecessary notifications to the developers involved in these PRs, -when syncing the main branch of a forked repository, please, follow these steps: -1. When possible, avoid syncing with the upstream using a branch and PR on the forked repository. Instead, merge directly into the forked main. -2. If a PR is absolutely necessary, use the following steps after checking out your branch: -```bash -$ git checkout -b your-branch-for-syncing -$ git pull --squash --no-commit upstream main -$ git commit -m '' -$ git push --set-upstream origin your-branch-for-syncing -``` - -### Style guide - -For documentation strings, 🧨 Diffusers follows the [Google style](https://google.github.io/styleguide/pyguide.html). diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 120000 index 000000000000..53de38ca21e3 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1 @@ +docs/source/en/conceptual/contribution.md \ No newline at end of file diff --git a/LICENSE b/LICENSE index 261eeb9e9f8b..038e32f6445e 100644 --- a/LICENSE +++ b/LICENSE @@ -144,7 +144,7 @@ agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions + implied, including, without limitation, Any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any diff --git a/Makefile b/Makefile index 9af2e8b1a5c9..b90ff82ab268 100644 --- a/Makefile +++ b/Makefile @@ -70,6 +70,10 @@ fix-copies: python utils/check_copies.py --fix_and_overwrite python utils/check_dummies.py --fix_and_overwrite +# Auto docstrings in modular blocks +modular-autodoctrings: + python utils/modular_auto_docstring.py + # Run tests for the library test: diff --git a/benchmarks/benchmarking_utils.py b/benchmarks/benchmarking_utils.py index c8c1a10ef899..141850e64f2e 100644 --- a/benchmarks/benchmarking_utils.py +++ b/benchmarks/benchmarking_utils.py @@ -6,7 +6,7 @@ import threading from contextlib import nullcontext from dataclasses import dataclass -from typing import Any, Callable, Dict, Optional, Union +from typing import Any, Callable import pandas as pd import torch @@ -91,10 +91,10 @@ def model_init_fn(model_cls, group_offload_kwargs=None, layerwise_upcasting=Fals class BenchmarkScenario: name: str model_cls: ModelMixin - model_init_kwargs: Dict[str, Any] + model_init_kwargs: dict[str, Any] model_init_fn: Callable get_model_input_dict: Callable - compile_kwargs: Optional[Dict[str, Any]] = None + compile_kwargs: dict[str, Any] | None = None @require_torch_gpu @@ -176,7 +176,7 @@ def run_benchmark(self, scenario: BenchmarkScenario): result["fullgraph"], result["mode"] = None, None return result - def run_bencmarks_and_collate(self, scenarios: Union[BenchmarkScenario, list[BenchmarkScenario]], filename: str): + def run_bencmarks_and_collate(self, scenarios: BenchmarkScenario | list[BenchmarkScenario], filename: str): if not isinstance(scenarios, list): scenarios = [scenarios] record_queue = queue.Queue() @@ -214,10 +214,10 @@ def _run_phase( *, model_cls: ModelMixin, init_fn: Callable, - init_kwargs: Dict[str, Any], + init_kwargs: dict[str, Any], get_input_fn: Callable, - compile_kwargs: Optional[Dict[str, Any]], - ) -> Dict[str, float]: + compile_kwargs: dict[str, Any] | None = None, + ) -> dict[str, float]: # setup self.pre_benchmark() diff --git a/benchmarks/populate_into_db.py b/benchmarks/populate_into_db.py deleted file mode 100644 index 55e46b058683..000000000000 --- a/benchmarks/populate_into_db.py +++ /dev/null @@ -1,166 +0,0 @@ -import argparse -import os -import sys - -import gpustat -import pandas as pd -import psycopg2 -import psycopg2.extras -from psycopg2.extensions import register_adapter -from psycopg2.extras import Json - - -register_adapter(dict, Json) - -FINAL_CSV_FILENAME = "collated_results.csv" -# https://github.com/huggingface/transformers/blob/593e29c5e2a9b17baec010e8dc7c1431fed6e841/benchmark/init_db.sql#L27 -BENCHMARKS_TABLE_NAME = "benchmarks" -MEASUREMENTS_TABLE_NAME = "model_measurements" - - -def _init_benchmark(conn, branch, commit_id, commit_msg): - gpu_stats = gpustat.GPUStatCollection.new_query() - metadata = {"gpu_name": gpu_stats[0]["name"]} - repository = "huggingface/diffusers" - with conn.cursor() as cur: - cur.execute( - f"INSERT INTO {BENCHMARKS_TABLE_NAME} (repository, branch, commit_id, commit_message, metadata) VALUES (%s, %s, %s, %s, %s) RETURNING benchmark_id", - (repository, branch, commit_id, commit_msg, metadata), - ) - benchmark_id = cur.fetchone()[0] - print(f"Initialised benchmark #{benchmark_id}") - return benchmark_id - - -def parse_args(): - parser = argparse.ArgumentParser() - parser.add_argument( - "branch", - type=str, - help="The branch name on which the benchmarking is performed.", - ) - - parser.add_argument( - "commit_id", - type=str, - help="The commit hash on which the benchmarking is performed.", - ) - - parser.add_argument( - "commit_msg", - type=str, - help="The commit message associated with the commit, truncated to 70 characters.", - ) - args = parser.parse_args() - return args - - -if __name__ == "__main__": - args = parse_args() - try: - conn = psycopg2.connect( - host=os.getenv("PGHOST"), - database=os.getenv("PGDATABASE"), - user=os.getenv("PGUSER"), - password=os.getenv("PGPASSWORD"), - ) - print("DB connection established successfully.") - except Exception as e: - print(f"Problem during DB init: {e}") - sys.exit(1) - - try: - benchmark_id = _init_benchmark( - conn=conn, - branch=args.branch, - commit_id=args.commit_id, - commit_msg=args.commit_msg, - ) - except Exception as e: - print(f"Problem during initializing benchmark: {e}") - sys.exit(1) - - cur = conn.cursor() - - df = pd.read_csv(FINAL_CSV_FILENAME) - - # Helper to cast values (or None) given a dtype - def _cast_value(val, dtype: str): - if pd.isna(val): - return None - - if dtype == "text": - return str(val).strip() - - if dtype == "float": - try: - return float(val) - except ValueError: - return None - - if dtype == "bool": - s = str(val).strip().lower() - if s in ("true", "t", "yes", "1"): - return True - if s in ("false", "f", "no", "0"): - return False - if val in (1, 1.0): - return True - if val in (0, 0.0): - return False - return None - - return val - - try: - rows_to_insert = [] - for _, row in df.iterrows(): - scenario = _cast_value(row.get("scenario"), "text") - model_cls = _cast_value(row.get("model_cls"), "text") - num_params_B = _cast_value(row.get("num_params_B"), "float") - flops_G = _cast_value(row.get("flops_G"), "float") - time_plain_s = _cast_value(row.get("time_plain_s"), "float") - mem_plain_GB = _cast_value(row.get("mem_plain_GB"), "float") - time_compile_s = _cast_value(row.get("time_compile_s"), "float") - mem_compile_GB = _cast_value(row.get("mem_compile_GB"), "float") - fullgraph = _cast_value(row.get("fullgraph"), "bool") - mode = _cast_value(row.get("mode"), "text") - - # If "github_sha" column exists in the CSV, cast it; else default to None - if "github_sha" in df.columns: - github_sha = _cast_value(row.get("github_sha"), "text") - else: - github_sha = None - - measurements = { - "scenario": scenario, - "model_cls": model_cls, - "num_params_B": num_params_B, - "flops_G": flops_G, - "time_plain_s": time_plain_s, - "mem_plain_GB": mem_plain_GB, - "time_compile_s": time_compile_s, - "mem_compile_GB": mem_compile_GB, - "fullgraph": fullgraph, - "mode": mode, - "github_sha": github_sha, - } - rows_to_insert.append((benchmark_id, measurements)) - - # Batch-insert all rows - insert_sql = f""" - INSERT INTO {MEASUREMENTS_TABLE_NAME} ( - benchmark_id, - measurements - ) - VALUES (%s, %s); - """ - - psycopg2.extras.execute_batch(cur, insert_sql, rows_to_insert) - conn.commit() - - cur.close() - conn.close() - except Exception as e: - print(f"Exception: {e}") - sys.exit(1) diff --git a/docker/diffusers-pytorch-cuda/Dockerfile b/docker/diffusers-pytorch-cuda/Dockerfile index a700d1db72bc..b11d8a491168 100644 --- a/docker/diffusers-pytorch-cuda/Dockerfile +++ b/docker/diffusers-pytorch-cuda/Dockerfile @@ -1,8 +1,8 @@ -FROM nvidia/cuda:12.1.0-runtime-ubuntu20.04 +FROM nvidia/cuda:12.9.1-runtime-ubuntu24.04 LABEL maintainer="Hugging Face" LABEL repository="diffusers" -ARG PYTHON_VERSION=3.12 +ARG PYTHON_VERSION=3.10 ENV DEBIAN_FRONTEND=noninteractive RUN apt-get -y update \ @@ -32,10 +32,17 @@ RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV} ENV PATH="$VIRTUAL_ENV/bin:$PATH" # pre-install the heavy dependencies (these can later be overridden by the deps from setup.py) +# Install torch, torchvision, and torchaudio together to ensure compatibility RUN uv pip install --no-cache-dir \ torch \ torchvision \ - torchaudio + torchaudio \ + --index-url https://download.pytorch.org/whl/cu129 + +# Install compatible versions of numba/llvmlite for Python 3.10+ +RUN uv pip install --no-cache-dir \ + "llvmlite>=0.40.0" \ + "numba>=0.57.0" RUN uv pip install --no-cache-dir "git+https://github.com/huggingface/diffusers.git@main#egg=diffusers[test]" diff --git a/docker/diffusers-pytorch-xformers-cuda/Dockerfile b/docker/diffusers-pytorch-xformers-cuda/Dockerfile index eae7eaf4faf1..33cae319b93c 100644 --- a/docker/diffusers-pytorch-xformers-cuda/Dockerfile +++ b/docker/diffusers-pytorch-xformers-cuda/Dockerfile @@ -1,8 +1,8 @@ -FROM nvidia/cuda:12.1.0-runtime-ubuntu20.04 +FROM nvidia/cuda:12.9.1-runtime-ubuntu24.04 LABEL maintainer="Hugging Face" LABEL repository="diffusers" -ARG PYTHON_VERSION=3.12 +ARG PYTHON_VERSION=3.10 ENV DEBIAN_FRONTEND=noninteractive RUN apt-get -y update \ @@ -32,10 +32,17 @@ RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV} ENV PATH="$VIRTUAL_ENV/bin:$PATH" # pre-install the heavy dependencies (these can later be overridden by the deps from setup.py) +# Install torch, torchvision, and torchaudio together to ensure compatibility RUN uv pip install --no-cache-dir \ torch \ torchvision \ - torchaudio + torchaudio \ + --index-url https://download.pytorch.org/whl/cu129 + +# Install compatible versions of numba/llvmlite for Python 3.10+ +RUN uv pip install --no-cache-dir \ + "llvmlite>=0.40.0" \ + "numba>=0.57.0" RUN uv pip install --no-cache-dir "git+https://github.com/huggingface/diffusers.git@main#egg=diffusers[test]" diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index f0cb0164436e..e0b7af4898b2 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -54,6 +54,8 @@ title: Batch inference - local: training/distributed_inference title: Distributed inference + - local: hybrid_inference/overview + title: Remote inference title: Inference - isExpanded: false sections: @@ -88,17 +90,6 @@ title: FreeU title: Community optimizations title: Inference optimization -- isExpanded: false - sections: - - local: hybrid_inference/overview - title: Overview - - local: hybrid_inference/vae_decode - title: VAE Decode - - local: hybrid_inference/vae_encode - title: VAE Encode - - local: hybrid_inference/api_reference - title: API Reference - title: Hybrid Inference - isExpanded: false sections: - local: modular_diffusers/overview @@ -123,6 +114,8 @@ title: Guiders - local: modular_diffusers/custom_blocks title: Building Custom Blocks + - local: modular_diffusers/mellon + title: Using Custom Blocks with Mellon title: Modular Diffusers - isExpanded: false sections: @@ -201,6 +194,8 @@ title: Model accelerators and hardware - isExpanded: false sections: + - local: using-diffusers/helios + title: Helios - local: using-diffusers/consisid title: ConsisID - local: using-diffusers/sdxl @@ -270,6 +265,8 @@ title: Outputs - local: api/quantization title: Quantization + - local: hybrid_inference/api_reference + title: Remote inference - local: api/parallel title: Parallel inference title: Main Classes @@ -353,6 +350,10 @@ title: Flux2Transformer2DModel - local: api/models/flux_transformer title: FluxTransformer2DModel + - local: api/models/glm_image_transformer2d + title: GlmImageTransformer2DModel + - local: api/models/helios_transformer3d + title: HeliosTransformer3DModel - local: api/models/hidream_image_transformer title: HiDreamImageTransformer2DModel - local: api/models/hunyuan_transformer2d @@ -367,6 +368,8 @@ title: LatteTransformer3DModel - local: api/models/longcat_image_transformer2d title: LongCatImageTransformer2DModel + - local: api/models/ltx2_video_transformer3d + title: LTX2VideoTransformer3DModel - local: api/models/ltx_video_transformer3d title: LTXVideoTransformer3DModel - local: api/models/lumina2_transformer2d @@ -443,6 +446,10 @@ title: AutoencoderKLHunyuanVideo - local: api/models/autoencoder_kl_hunyuan_video15 title: AutoencoderKLHunyuanVideo15 + - local: api/models/autoencoderkl_audio_ltx_2 + title: AutoencoderKLLTX2Audio + - local: api/models/autoencoderkl_ltx_2 + title: AutoencoderKLLTX2Video - local: api/models/autoencoderkl_ltx_video title: AutoencoderKLLTXVideo - local: api/models/autoencoderkl_magvit @@ -453,6 +460,8 @@ title: AutoencoderKLQwenImage - local: api/models/autoencoder_kl_wan title: AutoencoderKLWan + - local: api/models/autoencoder_rae + title: AutoencoderRAE - local: api/models/consistency_decoder_vae title: ConsistencyDecoderVAE - local: api/models/autoencoder_oobleck @@ -495,6 +504,8 @@ title: Bria 3.2 - local: api/pipelines/bria_fibo title: Bria Fibo + - local: api/pipelines/bria_fibo_edit + title: Bria Fibo Edit - local: api/pipelines/chroma title: Chroma - local: api/pipelines/cogview3 @@ -541,6 +552,8 @@ title: Flux2 - local: api/pipelines/control_flux_inpaint title: FluxControlInpaint + - local: api/pipelines/glm_image + title: GLM-Image - local: api/pipelines/hidream title: HiDream-I1 - local: api/pipelines/hunyuandit @@ -618,8 +631,6 @@ title: Image-to-image - local: api/pipelines/stable_diffusion/inpaint title: Inpainting - - local: api/pipelines/stable_diffusion/k_diffusion - title: K-Diffusion - local: api/pipelines/stable_diffusion/latent_upscale title: Latent upscaler - local: api/pipelines/stable_diffusion/ldm3d_diffusion @@ -668,6 +679,8 @@ title: ConsisID - local: api/pipelines/framepack title: Framepack + - local: api/pipelines/helios + title: Helios - local: api/pipelines/hunyuan_video title: HunyuanVideo - local: api/pipelines/hunyuan_video15 @@ -678,6 +691,8 @@ title: Kandinsky 5.0 Video - local: api/pipelines/latte title: Latte + - local: api/pipelines/ltx2 + title: LTX-2 - local: api/pipelines/ltx_video title: LTXVideo - local: api/pipelines/mochi @@ -737,6 +752,10 @@ title: FlowMatchEulerDiscreteScheduler - local: api/schedulers/flow_match_heun_discrete title: FlowMatchHeunDiscreteScheduler + - local: api/schedulers/helios_dmd + title: HeliosDMDScheduler + - local: api/schedulers/helios + title: HeliosScheduler - local: api/schedulers/heun title: HeunDiscreteScheduler - local: api/schedulers/ipndm diff --git a/docs/source/en/api/cache.md b/docs/source/en/api/cache.md index c93dcad43821..6a2d74892cfa 100644 --- a/docs/source/en/api/cache.md +++ b/docs/source/en/api/cache.md @@ -29,7 +29,7 @@ Cache methods speedup diffusion transformers by storing and reusing intermediate [[autodoc]] apply_faster_cache -### FirstBlockCacheConfig +## FirstBlockCacheConfig [[autodoc]] FirstBlockCacheConfig diff --git a/docs/source/en/api/loaders/lora.md b/docs/source/en/api/loaders/lora.md index 7911bc2b2332..db1ea884558f 100644 --- a/docs/source/en/api/loaders/lora.md +++ b/docs/source/en/api/loaders/lora.md @@ -23,6 +23,7 @@ LoRA is a fast and lightweight training method that inserts and trains a signifi - [`AuraFlowLoraLoaderMixin`] provides similar functions for [AuraFlow](https://huggingface.co/fal/AuraFlow). - [`LTXVideoLoraLoaderMixin`] provides similar functions for [LTX-Video](https://huggingface.co/docs/diffusers/main/en/api/pipelines/ltx_video). - [`SanaLoraLoaderMixin`] provides similar functions for [Sana](https://huggingface.co/docs/diffusers/main/en/api/pipelines/sana). +- [`HeliosLoraLoaderMixin`] provides similar functions for [HunyuanVideo](https://huggingface.co/docs/diffusers/main/en/api/pipelines/helios). - [`HunyuanVideoLoraLoaderMixin`] provides similar functions for [HunyuanVideo](https://huggingface.co/docs/diffusers/main/en/api/pipelines/hunyuan_video). - [`Lumina2LoraLoaderMixin`] provides similar functions for [Lumina2](https://huggingface.co/docs/diffusers/main/en/api/pipelines/lumina2). - [`WanLoraLoaderMixin`] provides similar functions for [Wan](https://huggingface.co/docs/diffusers/main/en/api/pipelines/wan). @@ -33,6 +34,7 @@ LoRA is a fast and lightweight training method that inserts and trains a signifi - [`QwenImageLoraLoaderMixin`] provides similar functions for [Qwen Image](https://huggingface.co/docs/diffusers/main/en/api/pipelines/qwen). - [`ZImageLoraLoaderMixin`] provides similar functions for [Z-Image](https://huggingface.co/docs/diffusers/main/en/api/pipelines/zimage). - [`Flux2LoraLoaderMixin`] provides similar functions for [Flux2](https://huggingface.co/docs/diffusers/main/en/api/pipelines/flux2). +- [`LTX2LoraLoaderMixin`] provides similar functions for [Flux2](https://huggingface.co/docs/diffusers/main/en/api/pipelines/ltx2). - [`LoraBaseMixin`] provides a base class with several utility methods to fuse, unfuse, unload, LoRAs and more. > [!TIP] @@ -62,6 +64,10 @@ LoRA is a fast and lightweight training method that inserts and trains a signifi [[autodoc]] loaders.lora_pipeline.Flux2LoraLoaderMixin +## LTX2LoraLoaderMixin + +[[autodoc]] loaders.lora_pipeline.LTX2LoraLoaderMixin + ## CogVideoXLoraLoaderMixin [[autodoc]] loaders.lora_pipeline.CogVideoXLoraLoaderMixin @@ -81,6 +87,10 @@ LoRA is a fast and lightweight training method that inserts and trains a signifi [[autodoc]] loaders.lora_pipeline.SanaLoraLoaderMixin +## HeliosLoraLoaderMixin + +[[autodoc]] loaders.lora_pipeline.HeliosLoraLoaderMixin + ## HunyuanVideoLoraLoaderMixin [[autodoc]] loaders.lora_pipeline.HunyuanVideoLoraLoaderMixin diff --git a/docs/source/en/api/models/autoencoder_rae.md b/docs/source/en/api/models/autoencoder_rae.md new file mode 100644 index 000000000000..a8c00dd4fde2 --- /dev/null +++ b/docs/source/en/api/models/autoencoder_rae.md @@ -0,0 +1,89 @@ + + +# AutoencoderRAE + +The Representation Autoencoder (RAE) model introduced in [Diffusion Transformers with Representation Autoencoders](https://huggingface.co/papers/2510.11690) by Boyang Zheng, Nanye Ma, Shengbang Tong, Saining Xie from NYU VISIONx. + +RAE combines a frozen pretrained vision encoder (DINOv2, SigLIP2, or MAE) with a trainable ViT-MAE-style decoder. In the two-stage RAE training recipe, the autoencoder is trained in stage 1 (reconstruction), and then a diffusion model is trained on the resulting latent space in stage 2 (generation). + +The following RAE models are released and supported in Diffusers: + +| Model | Encoder | Latent shape (224px input) | +|:------|:--------|:---------------------------| +| [`nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08`](https://huggingface.co/nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08) | DINOv2-base | 768 x 16 x 16 | +| [`nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08-i512`](https://huggingface.co/nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08-i512) | DINOv2-base (512px) | 768 x 32 x 32 | +| [`nyu-visionx/RAE-dinov2-wReg-small-ViTXL-n08`](https://huggingface.co/nyu-visionx/RAE-dinov2-wReg-small-ViTXL-n08) | DINOv2-small | 384 x 16 x 16 | +| [`nyu-visionx/RAE-dinov2-wReg-large-ViTXL-n08`](https://huggingface.co/nyu-visionx/RAE-dinov2-wReg-large-ViTXL-n08) | DINOv2-large | 1024 x 16 x 16 | +| [`nyu-visionx/RAE-siglip2-base-p16-i256-ViTXL-n08`](https://huggingface.co/nyu-visionx/RAE-siglip2-base-p16-i256-ViTXL-n08) | SigLIP2-base | 768 x 16 x 16 | +| [`nyu-visionx/RAE-mae-base-p16-ViTXL-n08`](https://huggingface.co/nyu-visionx/RAE-mae-base-p16-ViTXL-n08) | MAE-base | 768 x 16 x 16 | + +## Loading a pretrained model + +```python +from diffusers import AutoencoderRAE + +model = AutoencoderRAE.from_pretrained( + "nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08" +).to("cuda").eval() +``` + +## Encoding and decoding a real image + +```python +import torch +from diffusers import AutoencoderRAE +from diffusers.utils import load_image +from torchvision.transforms.functional import to_tensor, to_pil_image + +model = AutoencoderRAE.from_pretrained( + "nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08" +).to("cuda").eval() + +image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png") +image = image.convert("RGB").resize((224, 224)) +x = to_tensor(image).unsqueeze(0).to("cuda") # (1, 3, 224, 224), values in [0, 1] + +with torch.no_grad(): + latents = model.encode(x).latent # (1, 768, 16, 16) + recon = model.decode(latents).sample # (1, 3, 256, 256) + +recon_image = to_pil_image(recon[0].clamp(0, 1).cpu()) +recon_image.save("recon.png") +``` + +## Latent normalization + +Some pretrained checkpoints include per-channel `latents_mean` and `latents_std` statistics for normalizing the latent space. When present, `encode` and `decode` automatically apply the normalization and denormalization, respectively. + +```python +model = AutoencoderRAE.from_pretrained( + "nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08" +).to("cuda").eval() + +# Latent normalization is handled automatically inside encode/decode +# when the checkpoint config includes latents_mean/latents_std. +with torch.no_grad(): + latents = model.encode(x).latent # normalized latents + recon = model.decode(latents).sample +``` + +## AutoencoderRAE + +[[autodoc]] AutoencoderRAE + - encode + - decode + - all + +## DecoderOutput + +[[autodoc]] models.autoencoders.vae.DecoderOutput diff --git a/docs/source/en/api/models/autoencoderkl_audio_ltx_2.md b/docs/source/en/api/models/autoencoderkl_audio_ltx_2.md new file mode 100644 index 000000000000..d0024474e9e0 --- /dev/null +++ b/docs/source/en/api/models/autoencoderkl_audio_ltx_2.md @@ -0,0 +1,29 @@ + + +# AutoencoderKLLTX2Audio + +The 3D variational autoencoder (VAE) model with KL loss used in [LTX-2](https://huggingface.co/Lightricks/LTX-2) was introduced by Lightricks. This is for encoding and decoding audio latent representations. + +The model can be loaded with the following code snippet. + +```python +from diffusers import AutoencoderKLLTX2Audio + +vae = AutoencoderKLLTX2Audio.from_pretrained("Lightricks/LTX-2", subfolder="vae", torch_dtype=torch.float32).to("cuda") +``` + +## AutoencoderKLLTX2Audio + +[[autodoc]] AutoencoderKLLTX2Audio + - encode + - decode + - all \ No newline at end of file diff --git a/docs/source/en/api/models/autoencoderkl_ltx_2.md b/docs/source/en/api/models/autoencoderkl_ltx_2.md new file mode 100644 index 000000000000..1dbf516c017a --- /dev/null +++ b/docs/source/en/api/models/autoencoderkl_ltx_2.md @@ -0,0 +1,29 @@ + + +# AutoencoderKLLTX2Video + +The 3D variational autoencoder (VAE) model with KL loss used in [LTX-2](https://huggingface.co/Lightricks/LTX-2) was introduced by Lightricks. + +The model can be loaded with the following code snippet. + +```python +from diffusers import AutoencoderKLLTX2Video + +vae = AutoencoderKLLTX2Video.from_pretrained("Lightricks/LTX-2", subfolder="vae", torch_dtype=torch.float32).to("cuda") +``` + +## AutoencoderKLLTX2Video + +[[autodoc]] AutoencoderKLLTX2Video + - decode + - encode + - all diff --git a/docs/source/en/api/models/controlnet_flux.md b/docs/source/en/api/models/controlnet_flux.md index 6b230d90fba3..ec0370c19e06 100644 --- a/docs/source/en/api/models/controlnet_flux.md +++ b/docs/source/en/api/models/controlnet_flux.md @@ -42,4 +42,4 @@ pipe = FluxControlNetPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", co ## FluxControlNetOutput -[[autodoc]] models.controlnet_flux.FluxControlNetOutput \ No newline at end of file +[[autodoc]] models.controlnets.controlnet_flux.FluxControlNetOutput \ No newline at end of file diff --git a/docs/source/en/api/models/controlnet_sparsectrl.md b/docs/source/en/api/models/controlnet_sparsectrl.md index b9e81dc57eeb..0aa9848d0d2b 100644 --- a/docs/source/en/api/models/controlnet_sparsectrl.md +++ b/docs/source/en/api/models/controlnet_sparsectrl.md @@ -43,4 +43,4 @@ controlnet = SparseControlNetModel.from_pretrained("guoyww/animatediff-sparsectr ## SparseControlNetOutput -[[autodoc]] models.controlnet_sparsectrl.SparseControlNetOutput +[[autodoc]] models.controlnets.controlnet_sparsectrl.SparseControlNetOutput diff --git a/docs/source/en/api/models/glm_image_transformer2d.md b/docs/source/en/api/models/glm_image_transformer2d.md new file mode 100644 index 000000000000..7a18d1050075 --- /dev/null +++ b/docs/source/en/api/models/glm_image_transformer2d.md @@ -0,0 +1,18 @@ + + +# GlmImageTransformer2DModel + +A Diffusion Transformer model for 2D data from [GlmImageTransformer2DModel] (TODO). + +## GlmImageTransformer2DModel + +[[autodoc]] GlmImageTransformer2DModel diff --git a/docs/source/en/api/models/helios_transformer3d.md b/docs/source/en/api/models/helios_transformer3d.md new file mode 100644 index 000000000000..302b91d6c829 --- /dev/null +++ b/docs/source/en/api/models/helios_transformer3d.md @@ -0,0 +1,35 @@ + + +# HeliosTransformer3DModel + +A 14B Real-Time Autogressive Diffusion Transformer model (support T2V, I2V and V2V) for 3D video-like data from [Helios](https://github.com/PKU-YuanGroup/Helios) was introduced in [Helios: Real Real-Time Long Video Generation Model](https://huggingface.co/papers/2603.04379) by Peking University & ByteDance & etc. + +The model can be loaded with the following code snippet. + +```python +from diffusers import HeliosTransformer3DModel + +# Best Quality +transformer = HeliosTransformer3DModel.from_pretrained("BestWishYsh/Helios-Base", subfolder="transformer", torch_dtype=torch.bfloat16) +# Intermediate Weight +transformer = HeliosTransformer3DModel.from_pretrained("BestWishYsh/Helios-Mid", subfolder="transformer", torch_dtype=torch.bfloat16) +# Best Efficiency +transformer = HeliosTransformer3DModel.from_pretrained("BestWishYsh/Helios-Distilled", subfolder="transformer", torch_dtype=torch.bfloat16) +``` + +## HeliosTransformer3DModel + +[[autodoc]] HeliosTransformer3DModel + +## Transformer2DModelOutput + +[[autodoc]] models.modeling_outputs.Transformer2DModelOutput diff --git a/docs/source/en/api/models/ltx2_video_transformer3d.md b/docs/source/en/api/models/ltx2_video_transformer3d.md new file mode 100644 index 000000000000..9faab8695468 --- /dev/null +++ b/docs/source/en/api/models/ltx2_video_transformer3d.md @@ -0,0 +1,26 @@ + + +# LTX2VideoTransformer3DModel + +A Diffusion Transformer model for 3D data from [LTX](https://huggingface.co/Lightricks/LTX-2) was introduced by Lightricks. + +The model can be loaded with the following code snippet. + +```python +from diffusers import LTX2VideoTransformer3DModel + +transformer = LTX2VideoTransformer3DModel.from_pretrained("Lightricks/LTX-2", subfolder="transformer", torch_dtype=torch.bfloat16).to("cuda") +``` + +## LTX2VideoTransformer3DModel + +[[autodoc]] LTX2VideoTransformer3DModel diff --git a/docs/source/en/api/modular_diffusers/pipeline_blocks.md b/docs/source/en/api/modular_diffusers/pipeline_blocks.md index 8ad581e679ac..4808f2cf3bbe 100644 --- a/docs/source/en/api/modular_diffusers/pipeline_blocks.md +++ b/docs/source/en/api/modular_diffusers/pipeline_blocks.md @@ -14,4 +14,8 @@ ## AutoPipelineBlocks -[[autodoc]] diffusers.modular_pipelines.modular_pipeline.AutoPipelineBlocks \ No newline at end of file +[[autodoc]] diffusers.modular_pipelines.modular_pipeline.AutoPipelineBlocks + +## ConditionalPipelineBlocks + +[[autodoc]] diffusers.modular_pipelines.modular_pipeline.ConditionalPipelineBlocks \ No newline at end of file diff --git a/docs/source/en/api/pipelines/bria_fibo_edit.md b/docs/source/en/api/pipelines/bria_fibo_edit.md new file mode 100644 index 000000000000..b46dd78cdb90 --- /dev/null +++ b/docs/source/en/api/pipelines/bria_fibo_edit.md @@ -0,0 +1,33 @@ + + +# Bria Fibo Edit + +Fibo Edit is an 8B parameter image-to-image model that introduces a new paradigm of structured control, operating on JSON inputs paired with source images to enable deterministic and repeatable editing workflows. +Featuring native masking for granular precision, it moves beyond simple prompt-based diffusion to offer explicit, interpretable control optimized for production environments. +Its lightweight architecture is designed for deep customization, empowering researchers to build specialized "Edit" models for domain-specific tasks while delivering top-tier aesthetic quality + +## Usage +_As the model is gated, before using it with diffusers you first need to go to the [Bria Fibo Hugging Face page](https://huggingface.co/briaai/Fibo-Edit), fill in the form and accept the gate. Once you are in, you need to login so that your system knows you’ve accepted the gate._ + +Use the command below to log in: + +```bash +hf auth login +``` + + +## BriaFiboEditPipeline + +[[autodoc]] BriaFiboEditPipeline + - all + - __call__ \ No newline at end of file diff --git a/docs/source/en/api/pipelines/chroma.md b/docs/source/en/api/pipelines/chroma.md index cc52ffa09a6d..2b3b50c25e80 100644 --- a/docs/source/en/api/pipelines/chroma.md +++ b/docs/source/en/api/pipelines/chroma.md @@ -99,3 +99,9 @@ image.save("chroma-single-file.png") [[autodoc]] ChromaImg2ImgPipeline - all - __call__ + +## ChromaInpaintPipeline + +[[autodoc]] ChromaInpaintPipeline + - all + - __call__ diff --git a/docs/source/en/api/pipelines/chronoedit.md b/docs/source/en/api/pipelines/chronoedit.md index 48e70ab9e55e..5e7057f9ccb8 100644 --- a/docs/source/en/api/pipelines/chronoedit.md +++ b/docs/source/en/api/pipelines/chronoedit.md @@ -30,6 +30,10 @@ The ChronoEdit pipeline is developed by the ChronoEdit Team. The original code is available on [GitHub](https://github.com/nv-tlabs/ChronoEdit), and pretrained models can be found in the [nvidia/ChronoEdit](https://huggingface.co/collections/nvidia/chronoedit) collection on Hugging Face. +Available Models/LoRAs: +- [nvidia/ChronoEdit-14B-Diffusers](https://huggingface.co/nvidia/ChronoEdit-14B-Diffusers) +- [nvidia/ChronoEdit-14B-Diffusers-Upscaler-Lora](https://huggingface.co/nvidia/ChronoEdit-14B-Diffusers-Upscaler-Lora) +- [nvidia/ChronoEdit-14B-Diffusers-Paint-Brush-Lora](https://huggingface.co/nvidia/ChronoEdit-14B-Diffusers-Paint-Brush-Lora) ### Image Editing @@ -100,6 +104,7 @@ Image.fromarray((output[-1] * 255).clip(0, 255).astype("uint8")).save("output.pn import torch import numpy as np from diffusers import AutoencoderKLWan, ChronoEditTransformer3DModel, ChronoEditPipeline +from diffusers.schedulers import UniPCMultistepScheduler from diffusers.utils import export_to_video, load_image from transformers import CLIPVisionModel from PIL import Image @@ -109,9 +114,8 @@ image_encoder = CLIPVisionModel.from_pretrained(model_id, subfolder="image_encod vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32) transformer = ChronoEditTransformer3DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.bfloat16) pipe = ChronoEditPipeline.from_pretrained(model_id, image_encoder=image_encoder, transformer=transformer, vae=vae, torch_dtype=torch.bfloat16) -lora_path = hf_hub_download(repo_id=model_id, filename="lora/chronoedit_distill_lora.safetensors") -pipe.load_lora_weights(lora_path) -pipe.fuse_lora(lora_scale=1.0) +pipe.load_lora_weights("nvidia/ChronoEdit-14B-Diffusers", weight_name="lora/chronoedit_distill_lora.safetensors", adapter_name="distill") +pipe.fuse_lora(adapter_names=["distill"], lora_scale=1.0) pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=2.0) pipe.to("cuda") @@ -145,6 +149,57 @@ export_to_video(output, "output.mp4", fps=16) Image.fromarray((output[-1] * 255).clip(0, 255).astype("uint8")).save("output.png") ``` +### Inference with Multiple LoRAs + +```py +import torch +import numpy as np +from diffusers import AutoencoderKLWan, ChronoEditTransformer3DModel, ChronoEditPipeline +from diffusers.schedulers import UniPCMultistepScheduler +from diffusers.utils import export_to_video, load_image +from transformers import CLIPVisionModel +from PIL import Image + +model_id = "nvidia/ChronoEdit-14B-Diffusers" +image_encoder = CLIPVisionModel.from_pretrained(model_id, subfolder="image_encoder", torch_dtype=torch.float32) +vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32) +transformer = ChronoEditTransformer3DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.bfloat16) +pipe = ChronoEditPipeline.from_pretrained(model_id, image_encoder=image_encoder, transformer=transformer, vae=vae, torch_dtype=torch.bfloat16) +pipe.load_lora_weights("nvidia/ChronoEdit-14B-Diffusers-Paint-Brush-Lora", weight_name="paintbrush_lora_diffusers.safetensors", adapter_name="paintbrush") +pipe.load_lora_weights("nvidia/ChronoEdit-14B-Diffusers", weight_name="lora/chronoedit_distill_lora.safetensors", adapter_name="distill") +pipe.fuse_lora(adapter_names=["paintbrush", "distill"], lora_scale=1.0) +pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=2.0) +pipe.to("cuda") + +image = load_image( + "https://raw.githubusercontent.com/nv-tlabs/ChronoEdit/refs/heads/main/assets/images/input_paintbrush.png" +) +max_area = 720 * 1280 +aspect_ratio = image.height / image.width +mod_value = pipe.vae_scale_factor_spatial * pipe.transformer.config.patch_size[1] +height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value +width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value +print("width", width, "height", height) +image = image.resize((width, height)) +prompt = ( + "Turn the pencil sketch in the image into an actual object that is consistent with the image’s content. The user wants to change the sketch to a crown and a hat." +) + +output = pipe( + image=image, + prompt=prompt, + height=height, + width=width, + num_frames=5, + num_inference_steps=8, + guidance_scale=1.0, + enable_temporal_reasoning=False, + num_temporal_reasoning_steps=0, +).frames[0] +export_to_video(output, "output.mp4", fps=16) +Image.fromarray((output[-1] * 255).clip(0, 255).astype("uint8")).save("output_1.png") +``` + ## ChronoEditPipeline [[autodoc]] ChronoEditPipeline diff --git a/docs/source/en/api/pipelines/cosmos.md b/docs/source/en/api/pipelines/cosmos.md index 60ecce660303..2302ed2c4a6c 100644 --- a/docs/source/en/api/pipelines/cosmos.md +++ b/docs/source/en/api/pipelines/cosmos.md @@ -46,6 +46,20 @@ output = pipe( output.save("output.png") ``` +## Cosmos2_5_TransferPipeline + +[[autodoc]] Cosmos2_5_TransferPipeline + - all + - __call__ + + +## Cosmos2_5_PredictBasePipeline + +[[autodoc]] Cosmos2_5_PredictBasePipeline + - all + - __call__ + + ## CosmosTextToWorldPipeline [[autodoc]] CosmosTextToWorldPipeline @@ -70,12 +84,6 @@ output.save("output.png") - all - __call__ -## Cosmos2_5_PredictBasePipeline - -[[autodoc]] Cosmos2_5_PredictBasePipeline - - all - - __call__ - ## CosmosPipelineOutput [[autodoc]] pipelines.cosmos.pipeline_output.CosmosPipelineOutput diff --git a/docs/source/en/api/pipelines/diffedit.md b/docs/source/en/api/pipelines/diffedit.md index 9734ca2eabc3..670b7bb4fca0 100644 --- a/docs/source/en/api/pipelines/diffedit.md +++ b/docs/source/en/api/pipelines/diffedit.md @@ -21,7 +21,7 @@ The abstract from the paper is: *Image generation has recently seen tremendous advances, with diffusion models allowing to synthesize convincing images for a large variety of text prompts. In this article, we propose DiffEdit, a method to take advantage of text-conditioned diffusion models for the task of semantic image editing, where the goal is to edit an image based on a text query. Semantic image editing is an extension of image generation, with the additional constraint that the generated image should be as similar as possible to a given input image. Current editing methods based on diffusion models usually require to provide a mask, making the task much easier by treating it as a conditional inpainting task. In contrast, our main contribution is able to automatically generate a mask highlighting regions of the input image that need to be edited, by contrasting predictions of a diffusion model conditioned on different text prompts. Moreover, we rely on latent inference to preserve content in those regions of interest and show excellent synergies with mask-based diffusion. DiffEdit achieves state-of-the-art editing performance on ImageNet. In addition, we evaluate semantic image editing in more challenging settings, using images from the COCO dataset as well as text-based generated images.* -The original codebase can be found at [Xiang-cd/DiffEdit-stable-diffusion](https://github.com/Xiang-cd/DiffEdit-stable-diffusion), and you can try it out in this [demo](https://blog.problemsolversguild.com/technical/research/2022/11/02/DiffEdit-Implementation.html). +The original codebase can be found at [Xiang-cd/DiffEdit-stable-diffusion](https://github.com/Xiang-cd/DiffEdit-stable-diffusion), and you can try it out in this [demo](https://blog.problemsolversguild.com/posts/2022-11-02-diffedit-implementation.html). This pipeline was contributed by [clarencechen](https://github.com/clarencechen). ❤️ diff --git a/docs/source/en/api/pipelines/flux2.md b/docs/source/en/api/pipelines/flux2.md index 393e0d03c341..4ace2f3b3aa0 100644 --- a/docs/source/en/api/pipelines/flux2.md +++ b/docs/source/en/api/pipelines/flux2.md @@ -35,5 +35,11 @@ The [official implementation](https://github.com/black-forest-labs/flux2/blob/5a ## Flux2Pipeline [[autodoc]] Flux2Pipeline + - all + - __call__ + +## Flux2KleinPipeline + +[[autodoc]] Flux2KleinPipeline - all - __call__ \ No newline at end of file diff --git a/docs/source/en/api/pipelines/glm_image.md b/docs/source/en/api/pipelines/glm_image.md new file mode 100644 index 000000000000..a99832787847 --- /dev/null +++ b/docs/source/en/api/pipelines/glm_image.md @@ -0,0 +1,95 @@ + + +# GLM-Image + +## Overview + +GLM-Image is an image generation model adopts a hybrid autoregressive + diffusion decoder architecture, effectively pushing the upper bound of visual fidelity and fine-grained details. In general image generation quality, it aligns with industry-standard LDM-based approaches, while demonstrating significant advantages in knowledge-intensive image generation scenarios. + +Model architecture: a hybrid autoregressive + diffusion decoder design、 + ++ Autoregressive generator: a 9B-parameter model initialized from [GLM-4-9B-0414](https://huggingface.co/zai-org/GLM-4-9B-0414), with an expanded vocabulary to incorporate visual tokens. The model first generates a compact encoding of approximately 256 tokens, then expands to 1K–4K tokens, corresponding to 1K–2K high-resolution image outputs. You can check AR model in class `GlmImageForConditionalGeneration` of `transformers` library. ++ Diffusion Decoder: a 7B-parameter decoder based on a single-stream DiT architecture for latent-space image decoding. It is equipped with a Glyph Encoder text module, significantly improving accurate text rendering within images. + +Post-training with decoupled reinforcement learning: the model introduces a fine-grained, modular feedback strategy using the GRPO algorithm, substantially enhancing both semantic understanding and visual detail quality. + ++ Autoregressive module: provides low-frequency feedback signals focused on aesthetics and semantic alignment, improving instruction following and artistic expressiveness. ++ Decoder module: delivers high-frequency feedback targeting detail fidelity and text accuracy, resulting in highly realistic textures, lighting, and color reproduction, as well as more precise text rendering. + +GLM-Image supports both text-to-image and image-to-image generation within a single model + ++ Text-to-image: generates high-detail images from textual descriptions, with particularly strong performance in information-dense scenarios. ++ Image-to-image: supports a wide range of tasks, including image editing, style transfer, multi-subject consistency, and identity-preserving generation for people and objects. + +This pipeline was contributed by [zRzRzRzRzRzRzR](https://github.com/zRzRzRzRzRzRzR). The codebase can be found [here](https://huggingface.co/zai-org/GLM-Image). + +## Usage examples + +### Text to Image Generation + +```python +import torch +from diffusers.pipelines.glm_image import GlmImagePipeline + +pipe = GlmImagePipeline.from_pretrained("zai-org/GLM-Image",torch_dtype=torch.bfloat16,device_map="cuda") +prompt = "A beautifully designed modern food magazine style dessert recipe illustration, themed around a raspberry mousse cake. The overall layout is clean and bright, divided into four main areas: the top left features a bold black title 'Raspberry Mousse Cake Recipe Guide', with a soft-lit close-up photo of the finished cake on the right, showcasing a light pink cake adorned with fresh raspberries and mint leaves; the bottom left contains an ingredient list section, titled 'Ingredients' in a simple font, listing 'Flour 150g', 'Eggs 3', 'Sugar 120g', 'Raspberry puree 200g', 'Gelatin sheets 10g', 'Whipping cream 300ml', and 'Fresh raspberries', each accompanied by minimalist line icons (like a flour bag, eggs, sugar jar, etc.); the bottom right displays four equally sized step boxes, each containing high-definition macro photos and corresponding instructions, arranged from top to bottom as follows: Step 1 shows a whisk whipping white foam (with the instruction 'Whip egg whites to stiff peaks'), Step 2 shows a red-and-white mixture being folded with a spatula (with the instruction 'Gently fold in the puree and batter'), Step 3 shows pink liquid being poured into a round mold (with the instruction 'Pour into mold and chill for 4 hours'), Step 4 shows the finished cake decorated with raspberries and mint leaves (with the instruction 'Decorate with raspberries and mint'); a light brown information bar runs along the bottom edge, with icons on the left representing 'Preparation time: 30 minutes', 'Cooking time: 20 minutes', and 'Servings: 8'. The overall color scheme is dominated by creamy white and light pink, with a subtle paper texture in the background, featuring compact and orderly text and image layout with clear information hierarchy." +image = pipe( + prompt=prompt, + height=32 * 32, + width=36 * 32, + num_inference_steps=30, + guidance_scale=1.5, + generator=torch.Generator(device="cuda").manual_seed(42), +).images[0] + +image.save("output_t2i.png") +``` + +### Image to Image Generation + +```python +import torch +from diffusers.pipelines.glm_image import GlmImagePipeline +from PIL import Image + +pipe = GlmImagePipeline.from_pretrained("zai-org/GLM-Image",torch_dtype=torch.bfloat16,device_map="cuda") +image_path = "cond.jpg" +prompt = "Replace the background of the snow forest with an underground station featuring an automatic escalator." +image = Image.open(image_path).convert("RGB") +image = pipe( + prompt=prompt, + image=[image], # can input multiple images for multi-image-to-image generation such as [image, image1] + height=33 * 32, + width=32 * 32, + num_inference_steps=30, + guidance_scale=1.5, + generator=torch.Generator(device="cuda").manual_seed(42), +).images[0] + +image.save("output_i2i.png") +``` + ++ Since the AR model used in GLM-Image is configured with `do_sample=True` and a temperature of `0.95` by default, the generated images can vary significantly across runs. We do not recommend setting do_sample=False, as this may lead to incorrect or degenerate outputs from the AR model. + +## GlmImagePipeline + +[[autodoc]] pipelines.glm_image.pipeline_glm_image.GlmImagePipeline + - all + - __call__ + +## GlmImagePipelineOutput + +[[autodoc]] pipelines.glm_image.pipeline_output.GlmImagePipelineOutput diff --git a/docs/source/en/api/pipelines/helios.md b/docs/source/en/api/pipelines/helios.md new file mode 100644 index 000000000000..b85e1dca56b0 --- /dev/null +++ b/docs/source/en/api/pipelines/helios.md @@ -0,0 +1,464 @@ + + +
+
+ + LoRA + +
+
+ +# Helios + +[Helios: Real Real-Time Long Video Generation Model](https://huggingface.co/papers/2603.04379) from Peking University & ByteDance & etc, by Shenghai Yuan, Yuanyang Yin, Zongjian Li, Xinwei Huang, Xiao Yang, Li Yuan. + +* We introduce Helios, the first 14B video generation model that runs at 17 FPS on a single NVIDIA H100 GPU and supports minute-scale generation while matching a strong baseline in quality. We make breakthroughs along three key dimensions: (1) robustness to long-video drifting without commonly used anti-drift heuristics such as self-forcing, error banks, or keyframe sampling; (2) real-time generation without standard acceleration techniques such as KV-cache, causal masking, or sparse attention; and (3) training without parallelism or sharding frameworks, enabling image-diffusion-scale batch sizes while fitting up to four 14B models within 80 GB of GPU memory. Specifically, Helios is a 14B autoregressive diffusion model with a unified input representation that natively supports T2V, I2V, and V2V tasks. To mitigate drifting in long-video generation, we characterize its typical failure modes and propose simple yet effective training strategies that explicitly simulate drifting during training, while eliminating repetitive motion at its source. For efficiency, we heavily compress the historical and noisy context and reduce the number of sampling steps, yielding computational costs comparable to—or lower than—those of 1.3B video generative models. Moreover, we introduce infrastructure-level optimizations that accelerate both inference and training while reducing memory consumption. Extensive experiments demonstrate that Helios consistently outperforms prior methods on both short- and long-video generation. All the code and models are available at [this https URL](https://pku-yuangroup.github.io/Helios-Page). + +The following Helios models are supported in Diffusers: + +- [Helios-Base](https://huggingface.co/BestWishYsh/Helios-Base): Best Quality, with v-prediction, standard CFG and custom HeliosScheduler. +- [Helios-Mid](https://huggingface.co/BestWishYsh/Helios-Mid): Intermediate Weight, with v-prediction, CFG-Zero* and custom HeliosScheduler. +- [Helios-Distilled](https://huggingface.co/BestWishYsh/Helios-Distilled): Best Efficiency, with x0-prediction and custom HeliosDMDScheduler. + +> [!TIP] +> Click on the Helios models in the right sidebar for more examples of video generation. + +### Optimizing Memory and Inference Speed + +The example below demonstrates how to generate a video from text optimized for memory or inference speed. + + + + +Refer to the [Reduce memory usage](../../optimization/memory) guide for more details about the various memory saving techniques. + +The Helios model below requires ~6GB of VRAM. + +```py +import torch +from diffusers import AutoModel, HeliosPipeline +from diffusers.hooks.group_offloading import apply_group_offloading +from diffusers.utils import export_to_video + +vae = AutoModel.from_pretrained("BestWishYsh/Helios-Base", subfolder="vae", torch_dtype=torch.float32) + +# group-offloading +pipeline = HeliosPipeline.from_pretrained( + "BestWishYsh/Helios-Base", + vae=vae, + torch_dtype=torch.bfloat16 +) +pipeline.enable_group_offload( + onload_device=torch.device("cuda"), + offload_device=torch.device("cpu"), + offload_type="leaf_level", + use_stream=True, + record_stream=True, +) + +prompt = """ +A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue +and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with +a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, +allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades +of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and +the vivid colors of its surroundings. A close-up shot with dynamic movement. +""" +negative_prompt = """ +Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, +low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, +misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards +""" + +output = pipeline( + prompt=prompt, + negative_prompt=negative_prompt, + num_frames=99, + num_inference_steps=50, + guidance_scale=5.0, + generator=torch.Generator("cuda").manual_seed(42), +).frames[0] +export_to_video(output, "helios_base_t2v_output.mp4", fps=24) +``` + + + + +[Compilation](../../optimization/fp16#torchcompile) is slow the first time but subsequent calls to the pipeline are faster. [Attention Backends](../../optimization/attention_backends) such as FlashAttention and SageAttention can significantly increase speed by optimizing the computation of the attention mechanism. [Context Parallelism](../../training/distributed_inference#context-parallelism) splits the input sequence across multiple devices to enable processing of long contexts in parallel, reducing memory pressure and latency. [Caching](../../optimization/cache) may also speed up inference by storing and reusing intermediate outputs. + +```py +import torch +from diffusers import AutoModel, HeliosPipeline +from diffusers.utils import export_to_video + +vae = AutoModel.from_pretrained("BestWishYsh/Helios-Base", subfolder="vae", torch_dtype=torch.float32) + +pipeline = HeliosPipeline.from_pretrained( + "BestWishYsh/Helios-Base", + vae=vae, + torch_dtype=torch.bfloat16 +) +pipeline.to("cuda") + +# attention backend +# pipeline.transformer.set_attention_backend("flash") +pipeline.transformer.set_attention_backend("_flash_3_hub") # For Hopper GPUs + +# torch.compile +torch.backends.cudnn.benchmark = True +pipeline.text_encoder.compile(mode="max-autotune-no-cudagraphs", dynamic=False) +pipeline.vae.compile(mode="max-autotune-no-cudagraphs", dynamic=False) +pipeline.transformer.compile(mode="max-autotune-no-cudagraphs", dynamic=False) + +prompt = """ +A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue +and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with +a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, +allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades +of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and +the vivid colors of its surroundings. A close-up shot with dynamic movement. +""" +negative_prompt = """ +Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, +low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, +misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards +""" + +output = pipeline( + prompt=prompt, + negative_prompt=negative_prompt, + num_frames=99, + num_inference_steps=50, + guidance_scale=5.0, + generator=torch.Generator("cuda").manual_seed(42), +).frames[0] +export_to_video(output, "helios_base_t2v_output.mp4", fps=24) +``` + + + + + +### Generation with Helios-Base + +The example below demonstrates how to use Helios-Base to generate video based on text, image or video. + + + + +```python +import torch +from diffusers import AutoModel, HeliosPipeline +from diffusers.utils import export_to_video, load_video, load_image + +vae = AutoModel.from_pretrained("BestWishYsh/Helios-Base", subfolder="vae", torch_dtype=torch.float32) + +pipeline = HeliosPipeline.from_pretrained( + "BestWishYsh/Helios-Base", + vae=vae, + torch_dtype=torch.bfloat16 +) +pipeline.to("cuda") + +negative_prompt = """ +Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, +low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, +misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards +""" + +# For Text-to-Video +prompt = """ +A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue +and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with +a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, +allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades +of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and +the vivid colors of its surroundings. A close-up shot with dynamic movement. +""" + +output = pipeline( + prompt=prompt, + negative_prompt=negative_prompt, + num_frames=99, + num_inference_steps=50, + guidance_scale=5.0, + generator=torch.Generator("cuda").manual_seed(42), +).frames[0] +export_to_video(output, "helios_base_t2v_output.mp4", fps=24) + +# For Image-to-Video +prompt = """ +A towering emerald wave surges forward, its crest curling with raw power and energy. Sunlight glints off the translucent water, +illuminating the intricate textures and deep green hues within the wave’s body. A thick spray erupts from the breaking crest, +casting a misty veil that dances above the churning surface. As the perspective widens, the immense scale of the wave becomes +apparent, revealing the restless expanse of the ocean stretching beyond. The scene captures the ocean’s untamed beauty and +relentless force, with every droplet and ripple shimmering in the light. The dynamic motion and vivid colors evoke both awe and +respect for nature’s might. +""" +image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/helios/wave.jpg" + +output = pipeline( + prompt=prompt, + negative_prompt=negative_prompt, + image=load_image(image_path).resize((640, 384)), + num_frames=99, + num_inference_steps=50, + guidance_scale=5.0, + generator=torch.Generator("cuda").manual_seed(42), +).frames[0] +export_to_video(output, "helios_base_i2v_output.mp4", fps=24) + +# For Video-to-Video +prompt = """ +A bright yellow Lamborghini Huracn Tecnica speeds along a curving mountain road, surrounded by lush green trees +under a partly cloudy sky. The car's sleek design and vibrant color stand out against the natural backdrop, +emphasizing its dynamic movement. The road curves gently, with a guardrail visible on one side, adding depth to +the scene. The motion blur captures the sense of speed and energy, creating a thrilling and exhilarating atmosphere. +A front-facing shot from a slightly elevated angle, highlighting the car's aggressive stance and the surrounding greenery. +""" +video_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/helios/car.mp4" + +output = pipeline( + prompt=prompt, + negative_prompt=negative_prompt, + video=load_video(video_path), + num_frames=99, + num_inference_steps=50, + guidance_scale=5.0, + generator=torch.Generator("cuda").manual_seed(42), +).frames[0] +export_to_video(output, "helios_base_v2v_output.mp4", fps=24) +``` + + + + + +### Generation with Helios-Mid + +The example below demonstrates how to use Helios-Mid to generate video based on text, image or video. + + + + +```python +import torch +from diffusers import AutoModel, HeliosPyramidPipeline +from diffusers.utils import export_to_video, load_video, load_image + +vae = AutoModel.from_pretrained("BestWishYsh/Helios-Mid", subfolder="vae", torch_dtype=torch.float32) + +pipeline = HeliosPyramidPipeline.from_pretrained( + "BestWishYsh/Helios-Mid", + vae=vae, + torch_dtype=torch.bfloat16 +) +pipeline.to("cuda") + +negative_prompt = """ +Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, +low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, +misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards +""" + +# For Text-to-Video +prompt = """ +A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue +and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with +a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, +allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades +of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and +the vivid colors of its surroundings. A close-up shot with dynamic movement. +""" + +output = pipeline( + prompt=prompt, + negative_prompt=negative_prompt, + num_frames=99, + pyramid_num_inference_steps_list=[20, 20, 20], + guidance_scale=5.0, + use_zero_init=True, + zero_steps=1, + generator=torch.Generator("cuda").manual_seed(42), +).frames[0] +export_to_video(output, "helios_pyramid_t2v_output.mp4", fps=24) + +# For Image-to-Video +prompt = """ +A towering emerald wave surges forward, its crest curling with raw power and energy. Sunlight glints off the translucent water, +illuminating the intricate textures and deep green hues within the wave’s body. A thick spray erupts from the breaking crest, +casting a misty veil that dances above the churning surface. As the perspective widens, the immense scale of the wave becomes +apparent, revealing the restless expanse of the ocean stretching beyond. The scene captures the ocean’s untamed beauty and +relentless force, with every droplet and ripple shimmering in the light. The dynamic motion and vivid colors evoke both awe and +respect for nature’s might. +""" +image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/helios/wave.jpg" + +output = pipeline( + prompt=prompt, + negative_prompt=negative_prompt, + image=load_image(image_path).resize((640, 384)), + num_frames=99, + pyramid_num_inference_steps_list=[20, 20, 20], + guidance_scale=5.0, + use_zero_init=True, + zero_steps=1, + generator=torch.Generator("cuda").manual_seed(42), +).frames[0] +export_to_video(output, "helios_pyramid_i2v_output.mp4", fps=24) + +# For Video-to-Video +prompt = """ +A bright yellow Lamborghini Huracn Tecnica speeds along a curving mountain road, surrounded by lush green trees +under a partly cloudy sky. The car's sleek design and vibrant color stand out against the natural backdrop, +emphasizing its dynamic movement. The road curves gently, with a guardrail visible on one side, adding depth to +the scene. The motion blur captures the sense of speed and energy, creating a thrilling and exhilarating atmosphere. +A front-facing shot from a slightly elevated angle, highlighting the car's aggressive stance and the surrounding greenery. +""" +video_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/helios/car.mp4" + +output = pipeline( + prompt=prompt, + negative_prompt=negative_prompt, + video=load_video(video_path), + num_frames=99, + pyramid_num_inference_steps_list=[20, 20, 20], + guidance_scale=5.0, + use_zero_init=True, + zero_steps=1, + generator=torch.Generator("cuda").manual_seed(42), +).frames[0] +export_to_video(output, "helios_pyramid_v2v_output.mp4", fps=24) +``` + + + + + +### Generation with Helios-Distilled + +The example below demonstrates how to use Helios-Distilled to generate video based on text, image or video. + + + + +```python +import torch +from diffusers import AutoModel, HeliosPyramidPipeline +from diffusers.utils import export_to_video, load_video, load_image + +vae = AutoModel.from_pretrained("BestWishYsh/Helios-Distilled", subfolder="vae", torch_dtype=torch.float32) + +pipeline = HeliosPyramidPipeline.from_pretrained( + "BestWishYsh/Helios-Distilled", + vae=vae, + torch_dtype=torch.bfloat16 +) +pipeline.to("cuda") + +negative_prompt = """ +Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, +low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, +misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards +""" + +# For Text-to-Video +prompt = """ +A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue +and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with +a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, +allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades +of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and +the vivid colors of its surroundings. A close-up shot with dynamic movement. +""" + +output = pipeline( + prompt=prompt, + negative_prompt=negative_prompt, + num_frames=240, + pyramid_num_inference_steps_list=[2, 2, 2], + guidance_scale=1.0, + is_amplify_first_chunk=True, + generator=torch.Generator("cuda").manual_seed(42), +).frames[0] +export_to_video(output, "helios_distilled_t2v_output.mp4", fps=24) + +# For Image-to-Video +prompt = """ +A towering emerald wave surges forward, its crest curling with raw power and energy. Sunlight glints off the translucent water, +illuminating the intricate textures and deep green hues within the wave’s body. A thick spray erupts from the breaking crest, +casting a misty veil that dances above the churning surface. As the perspective widens, the immense scale of the wave becomes +apparent, revealing the restless expanse of the ocean stretching beyond. The scene captures the ocean’s untamed beauty and +relentless force, with every droplet and ripple shimmering in the light. The dynamic motion and vivid colors evoke both awe and +respect for nature’s might. +""" +image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/helios/wave.jpg" + +output = pipeline( + prompt=prompt, + negative_prompt=negative_prompt, + image=load_image(image_path).resize((640, 384)), + num_frames=240, + pyramid_num_inference_steps_list=[2, 2, 2], + guidance_scale=1.0, + is_amplify_first_chunk=True, + generator=torch.Generator("cuda").manual_seed(42), +).frames[0] +export_to_video(output, "helios_distilled_i2v_output.mp4", fps=24) + +# For Video-to-Video +prompt = """ +A bright yellow Lamborghini Huracn Tecnica speeds along a curving mountain road, surrounded by lush green trees +under a partly cloudy sky. The car's sleek design and vibrant color stand out against the natural backdrop, +emphasizing its dynamic movement. The road curves gently, with a guardrail visible on one side, adding depth to +the scene. The motion blur captures the sense of speed and energy, creating a thrilling and exhilarating atmosphere. +A front-facing shot from a slightly elevated angle, highlighting the car's aggressive stance and the surrounding greenery. +""" +video_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/helios/car.mp4" + +output = pipeline( + prompt=prompt, + negative_prompt=negative_prompt, + video=load_video(video_path), + num_frames=240, + pyramid_num_inference_steps_list=[2, 2, 2], + guidance_scale=1.0, + is_amplify_first_chunk=True, + generator=torch.Generator("cuda").manual_seed(42), +).frames[0] +export_to_video(output, "helios_distilled_v2v_output.mp4", fps=24) +``` + + + + + +## HeliosPipeline + +[[autodoc]] HeliosPipeline + + - all + - __call__ + +## HeliosPyramidPipeline + +[[autodoc]] HeliosPyramidPipeline + + - all + - __call__ + +## HeliosPipelineOutput + +[[autodoc]] pipelines.helios.pipeline_output.HeliosPipelineOutput diff --git a/docs/source/en/api/pipelines/ltx2.md b/docs/source/en/api/pipelines/ltx2.md new file mode 100644 index 000000000000..85b0f9691891 --- /dev/null +++ b/docs/source/en/api/pipelines/ltx2.md @@ -0,0 +1,395 @@ + + +# LTX-2 + +
+ LoRA +
+ +LTX-2 is a DiT-based audio-video foundation model designed to generate synchronized video and audio within a single model. It brings together the core building blocks of modern video generation, with open weights and a focus on practical, local execution. + +You can find all the original LTX-Video checkpoints under the [Lightricks](https://huggingface.co/Lightricks) organization. + +The original codebase for LTX-2 can be found [here](https://github.com/Lightricks/LTX-2). + +## Two-stages Generation +Recommended pipeline to achieve production quality generation, this pipeline is composed of two stages: + +- Stage 1: Generate a video at the target resolution using diffusion sampling with classifier-free guidance (CFG). This stage produces a coherent low-noise video sequence that respects the text/image conditioning. +- Stage 2: Upsample the Stage 1 output by 2 and refine details using a distilled LoRA model to improve fidelity and visual quality. Stage 2 may apply lighter CFG to preserve the structure from Stage 1 while enhancing texture and sharpness. + +Sample usage of text-to-video two stages pipeline + +```py +import torch +from diffusers import FlowMatchEulerDiscreteScheduler +from diffusers.pipelines.ltx2 import LTX2Pipeline, LTX2LatentUpsamplePipeline +from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel +from diffusers.pipelines.ltx2.utils import STAGE_2_DISTILLED_SIGMA_VALUES +from diffusers.pipelines.ltx2.export_utils import encode_video + +device = "cuda:0" +width = 768 +height = 512 + +pipe = LTX2Pipeline.from_pretrained( + "Lightricks/LTX-2", torch_dtype=torch.bfloat16 +) +pipe.enable_sequential_cpu_offload(device=device) + +prompt = "A beautiful sunset over the ocean" +negative_prompt = "shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static." + +# Stage 1 default (non-distilled) inference +frame_rate = 24.0 +video_latent, audio_latent = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + width=width, + height=height, + num_frames=121, + frame_rate=frame_rate, + num_inference_steps=40, + sigmas=None, + guidance_scale=4.0, + output_type="latent", + return_dict=False, +) + +latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained( + "Lightricks/LTX-2", + subfolder="latent_upsampler", + torch_dtype=torch.bfloat16, +) +upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=latent_upsampler) +upsample_pipe.enable_model_cpu_offload(device=device) +upscaled_video_latent = upsample_pipe( + latents=video_latent, + output_type="latent", + return_dict=False, +)[0] + +# Load Stage 2 distilled LoRA +pipe.load_lora_weights( + "Lightricks/LTX-2", adapter_name="stage_2_distilled", weight_name="ltx-2-19b-distilled-lora-384.safetensors" +) +pipe.set_adapters("stage_2_distilled", 1.0) +# VAE tiling is usually necessary to avoid OOM error when VAE decoding +pipe.vae.enable_tiling() +# Change scheduler to use Stage 2 distilled sigmas as is +new_scheduler = FlowMatchEulerDiscreteScheduler.from_config( + pipe.scheduler.config, use_dynamic_shifting=False, shift_terminal=None +) +pipe.scheduler = new_scheduler +# Stage 2 inference with distilled LoRA and sigmas +video, audio = pipe( + latents=upscaled_video_latent, + audio_latents=audio_latent, + prompt=prompt, + negative_prompt=negative_prompt, + num_inference_steps=3, + noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0], # renoise with first sigma value https://github.com/Lightricks/LTX-2/blob/main/packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py#L218 + sigmas=STAGE_2_DISTILLED_SIGMA_VALUES, + guidance_scale=1.0, + output_type="np", + return_dict=False, +) + +encode_video( + video[0], + fps=frame_rate, + audio=audio[0].float().cpu(), + audio_sample_rate=pipe.vocoder.config.output_sampling_rate, + output_path="ltx2_lora_distilled_sample.mp4", +) +``` + +## Distilled checkpoint generation +Fastest two-stages generation pipeline using a distilled checkpoint. + +```py +import torch +from diffusers.pipelines.ltx2 import LTX2Pipeline, LTX2LatentUpsamplePipeline +from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel +from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES +from diffusers.pipelines.ltx2.export_utils import encode_video + +device = "cuda" +width = 768 +height = 512 +random_seed = 42 +generator = torch.Generator(device).manual_seed(random_seed) +model_path = "rootonchair/LTX-2-19b-distilled" + +pipe = LTX2Pipeline.from_pretrained( + model_path, torch_dtype=torch.bfloat16 +) +pipe.enable_sequential_cpu_offload(device=device) + +prompt = "A beautiful sunset over the ocean" +negative_prompt = "shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static." + +frame_rate = 24.0 +video_latent, audio_latent = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + width=width, + height=height, + num_frames=121, + frame_rate=frame_rate, + num_inference_steps=8, + sigmas=DISTILLED_SIGMA_VALUES, + guidance_scale=1.0, + generator=generator, + output_type="latent", + return_dict=False, +) + +latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained( + model_path, + subfolder="latent_upsampler", + torch_dtype=torch.bfloat16, +) +upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=latent_upsampler) +upsample_pipe.enable_model_cpu_offload(device=device) +upscaled_video_latent = upsample_pipe( + latents=video_latent, + output_type="latent", + return_dict=False, +)[0] + +video, audio = pipe( + latents=upscaled_video_latent, + audio_latents=audio_latent, + prompt=prompt, + negative_prompt=negative_prompt, + num_inference_steps=3, + noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0], # renoise with first sigma value https://github.com/Lightricks/LTX-2/blob/main/packages/ltx-pipelines/src/ltx_pipelines/distilled.py#L178 + sigmas=STAGE_2_DISTILLED_SIGMA_VALUES, + generator=generator, + guidance_scale=1.0, + output_type="np", + return_dict=False, +) + +encode_video( + video[0], + fps=frame_rate, + audio=audio[0].float().cpu(), + audio_sample_rate=pipe.vocoder.config.output_sampling_rate, + output_path="ltx2_distilled_sample.mp4", +) +``` + +## Condition Pipeline Generation + +You can use `LTX2ConditionPipeline` to specify image and/or video conditions at arbitrary latent indices. For example, we can specify both a first-frame and last-frame condition to perform first-last-frame-to-video (FLF2V) generation: + +```py +import torch +from diffusers import LTX2ConditionPipeline, LTX2LatentUpsamplePipeline +from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel +from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition +from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES +from diffusers.pipelines.ltx2.export_utils import encode_video +from diffusers.utils import load_image + +device = "cuda" +width = 768 +height = 512 +random_seed = 42 +generator = torch.Generator(device).manual_seed(random_seed) +model_path = "rootonchair/LTX-2-19b-distilled" + +pipe = LTX2ConditionPipeline.from_pretrained(model_path, torch_dtype=torch.bfloat16) +pipe.enable_sequential_cpu_offload(device=device) +pipe.vae.enable_tiling() + +prompt = ( + "CG animation style, a small blue bird takes off from the ground, flapping its wings. The bird's feathers are " + "delicate, with a unique pattern on its chest. The background shows a blue sky with white clouds under bright " + "sunshine. The camera follows the bird upward, capturing its flight and the vastness of the sky from a close-up, " + "low-angle perspective." +) + +first_image = load_image( + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_first_frame.png", +) +last_image = load_image( + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_last_frame.png", +) +first_cond = LTX2VideoCondition(frames=first_image, index=0, strength=1.0) +last_cond = LTX2VideoCondition(frames=last_image, index=-1, strength=1.0) +conditions = [first_cond, last_cond] + +frame_rate = 24.0 +video_latent, audio_latent = pipe( + conditions=conditions, + prompt=prompt, + width=width, + height=height, + num_frames=121, + frame_rate=frame_rate, + num_inference_steps=8, + sigmas=DISTILLED_SIGMA_VALUES, + guidance_scale=1.0, + generator=generator, + output_type="latent", + return_dict=False, +) + +latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained( + model_path, + subfolder="latent_upsampler", + torch_dtype=torch.bfloat16, +) +upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=latent_upsampler) +upsample_pipe.enable_model_cpu_offload(device=device) +upscaled_video_latent = upsample_pipe( + latents=video_latent, + output_type="latent", + return_dict=False, +)[0] + +video, audio = pipe( + latents=upscaled_video_latent, + audio_latents=audio_latent, + prompt=prompt, + width=width * 2, + height=height * 2, + num_inference_steps=3, + sigmas=STAGE_2_DISTILLED_SIGMA_VALUES, + generator=generator, + guidance_scale=1.0, + output_type="np", + return_dict=False, +) + +encode_video( + video[0], + fps=frame_rate, + audio=audio[0].float().cpu(), + audio_sample_rate=pipe.vocoder.config.output_sampling_rate, + output_path="ltx2_distilled_flf2v.mp4", +) +``` + +You can use both image and video conditions: + +```py +import torch +from diffusers import LTX2ConditionPipeline +from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition +from diffusers.pipelines.ltx2.export_utils import encode_video +from diffusers.utils import load_image, load_video + +device = "cuda" +width = 768 +height = 512 +random_seed = 42 +generator = torch.Generator(device).manual_seed(random_seed) +model_path = "rootonchair/LTX-2-19b-distilled" + +pipe = LTX2ConditionPipeline.from_pretrained(model_path, torch_dtype=torch.bfloat16) +pipe.enable_sequential_cpu_offload(device=device) +pipe.vae.enable_tiling() + +prompt = ( + "The video depicts a long, straight highway stretching into the distance, flanked by metal guardrails. The road is " + "divided into multiple lanes, with a few vehicles visible in the far distance. The surrounding landscape features " + "dry, grassy fields on one side and rolling hills on the other. The sky is mostly clear with a few scattered " + "clouds, suggesting a bright, sunny day. And then the camera switch to a winding mountain road covered in snow, " + "with a single vehicle traveling along it. The road is flanked by steep, rocky cliffs and sparse vegetation. The " + "landscape is characterized by rugged terrain and a river visible in the distance. The scene captures the " + "solitude and beauty of a winter drive through a mountainous region." +) +negative_prompt = ( + "blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, excessive noise, " + "grainy texture, poor lighting, flickering, motion blur, distorted proportions, unnatural skin tones, " + "deformed facial features, asymmetrical face, missing facial features, extra limbs, disfigured hands, " + "wrong hand count, artifacts around text, inconsistent perspective, camera shake, incorrect depth of " + "field, background too sharp, background clutter, distracting reflections, harsh shadows, inconsistent " + "lighting direction, color banding, cartoonish rendering, 3D CGI look, unrealistic materials, uncanny " + "valley effect, incorrect ethnicity, wrong gender, exaggerated expressions, wrong gaze direction, " + "mismatched lip sync, silent or muted audio, distorted voice, robotic voice, echo, background noise, " + "off-sync audio, incorrect dialogue, added dialogue, repetitive speech, jittery movement, awkward " + "pauses, incorrect timing, unnatural transitions, inconsistent framing, tilted camera, flat lighting, " + "inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts." +) + +cond_video = load_video( + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cosmos/cosmos-video2world-input-vid.mp4" +) +cond_image = load_image( + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cosmos/cosmos-video2world-input.jpg" +) +video_cond = LTX2VideoCondition(frames=cond_video, index=0, strength=1.0) +image_cond = LTX2VideoCondition(frames=cond_image, index=8, strength=1.0) +conditions = [video_cond, image_cond] + +frame_rate = 24.0 +video, audio = pipe( + conditions=conditions, + prompt=prompt, + negative_prompt=negative_prompt, + width=width, + height=height, + num_frames=121, + frame_rate=frame_rate, + num_inference_steps=40, + guidance_scale=4.0, + generator=generator, + output_type="np", + return_dict=False, +) + +encode_video( + video[0], + fps=frame_rate, + audio=audio[0].float().cpu(), + audio_sample_rate=pipe.vocoder.config.output_sampling_rate, + output_path="ltx2_cond_video.mp4", +) +``` + +Because the conditioning is done via latent frames, the 8 data space frames corresponding to the specified latent frame for an image condition will tend to be static. + +## LTX2Pipeline + +[[autodoc]] LTX2Pipeline + - all + - __call__ + +## LTX2ImageToVideoPipeline + +[[autodoc]] LTX2ImageToVideoPipeline + - all + - __call__ + +## LTX2ConditionPipeline + +[[autodoc]] LTX2ConditionPipeline + - all + - __call__ + +## LTX2LatentUpsamplePipeline + +[[autodoc]] LTX2LatentUpsamplePipeline + - all + - __call__ + +## LTX2PipelineOutput + +[[autodoc]] pipelines.ltx2.pipeline_output.LTX2PipelineOutput diff --git a/docs/source/en/api/pipelines/ltx_video.md b/docs/source/en/api/pipelines/ltx_video.md index 940144538a35..68658f41dabc 100644 --- a/docs/source/en/api/pipelines/ltx_video.md +++ b/docs/source/en/api/pipelines/ltx_video.md @@ -136,7 +136,7 @@ export_to_video(video, "output.mp4", fps=24) - The recommended dtype for the transformer, VAE, and text encoder is `torch.bfloat16`. The VAE and text encoder can also be `torch.float32` or `torch.float16`. - For guidance-distilled variants of LTX-Video, set `guidance_scale` to `1.0`. The `guidance_scale` for any other model should be set higher, like `5.0`, for good generation quality. - For timestep-aware VAE variants (LTX-Video 0.9.1 and above), set `decode_timestep` to `0.05` and `image_cond_noise_scale` to `0.025`. - - For variants that support interpolation between multiple conditioning images and videos (LTX-Video 0.9.5 and above), use similar images and videos for the best results. Divergence from the conditioning inputs may lead to abrupt transitionts in the generated video. + - For variants that support interpolation between multiple conditioning images and videos (LTX-Video 0.9.5 and above), use similar images and videos for the best results. Divergence from the conditioning inputs may lead to abrupt transitions in the generated video. - LTX-Video 0.9.7 includes a spatial latent upscaler and a 13B parameter transformer. During inference, a low resolution video is quickly generated first and then upscaled and refined. @@ -329,7 +329,7 @@ export_to_video(video, "output.mp4", fps=24)
Show example code - + ```python import torch from diffusers import LTXConditionPipeline, LTXLatentUpsamplePipeline @@ -474,6 +474,12 @@ export_to_video(video, "output.mp4", fps=24)
+## LTXI2VLongMultiPromptPipeline + +[[autodoc]] LTXI2VLongMultiPromptPipeline + - all + - __call__ + ## LTXPipeline [[autodoc]] LTXPipeline diff --git a/docs/source/en/api/pipelines/qwenimage.md b/docs/source/en/api/pipelines/qwenimage.md index b3dd3dd93618..c0994c8685d0 100644 --- a/docs/source/en/api/pipelines/qwenimage.md +++ b/docs/source/en/api/pipelines/qwenimage.md @@ -29,7 +29,7 @@ Qwen-Image comes in the following variants: | Qwen-Image-Edit Plus | [Qwen/Qwen-Image-Edit-2509](https://huggingface.co/Qwen/Qwen-Image-Edit-2509) | > [!TIP] -> [Caching](../../optimization/cache) may also speed up inference by storing and reusing intermediate outputs. +> See the [Caching](../../optimization/cache) guide to speed up inference by storing and reusing intermediate outputs. ## LoRA for faster inference @@ -95,7 +95,7 @@ image.save("qwen_fewsteps.png") With [`QwenImageEditPlusPipeline`], one can provide multiple images as input reference. -``` +```py import torch from PIL import Image from diffusers import QwenImageEditPlusPipeline @@ -108,12 +108,46 @@ pipe = QwenImageEditPlusPipeline.from_pretrained( image_1 = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/grumpy.jpg") image_2 = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/peng.png") image = pipe( - image=[image_1, image_2], - prompt='''put the penguin and the cat at a game show called "Qwen Edit Plus Games"''', + image=[image_1, image_2], + prompt='''put the penguin and the cat at a game show called "Qwen Edit Plus Games"''', num_inference_steps=50 ).images[0] ``` +## Performance + +### torch.compile + +Using `torch.compile` on the transformer provides ~2.4x speedup (A100 80GB: 4.70s → 1.93s): + +```python +import torch +from diffusers import QwenImagePipeline + +pipe = QwenImagePipeline.from_pretrained("Qwen/Qwen-Image", torch_dtype=torch.bfloat16).to("cuda") +pipe.transformer = torch.compile(pipe.transformer) + +# First call triggers compilation (~7s overhead) +# Subsequent calls run at ~2.4x faster +image = pipe("a cat", num_inference_steps=50).images[0] +``` + +### Batched Inference with Variable-Length Prompts + +When using classifier-free guidance (CFG) with prompts of different lengths, the pipeline properly handles padding through attention masking. This ensures padding tokens do not influence the generated output. + +```python +# CFG with different prompt lengths works correctly +image = pipe( + prompt="A cat", + negative_prompt="blurry, low quality, distorted", + true_cfg_scale=3.5, + num_inference_steps=50, +).images[0] +``` + +For detailed benchmark scripts and results, see [this gist](https://gist.github.com/cdutr/bea337e4680268168550292d7819dc2f). + ## QwenImagePipeline [[autodoc]] QwenImagePipeline @@ -156,6 +190,12 @@ image = pipe( - all - __call__ +## QwenImageLayeredPipeline + +[[autodoc]] QwenImageLayeredPipeline + - all + - __call__ + ## QwenImagePipelineOutput [[autodoc]] pipelines.qwenimage.pipeline_output.QwenImagePipelineOutput \ No newline at end of file diff --git a/docs/source/en/api/pipelines/skyreels_v2.md b/docs/source/en/api/pipelines/skyreels_v2.md index 6730f1551607..e1829bc409eb 100644 --- a/docs/source/en/api/pipelines/skyreels_v2.md +++ b/docs/source/en/api/pipelines/skyreels_v2.md @@ -37,7 +37,8 @@ The following SkyReels-V2 models are supported in Diffusers: - [SkyReels-V2 I2V 1.3B - 540P](https://huggingface.co/Skywork/SkyReels-V2-I2V-1.3B-540P-Diffusers) - [SkyReels-V2 I2V 14B - 540P](https://huggingface.co/Skywork/SkyReels-V2-I2V-14B-540P-Diffusers) - [SkyReels-V2 I2V 14B - 720P](https://huggingface.co/Skywork/SkyReels-V2-I2V-14B-720P-Diffusers) -- [SkyReels-V2 FLF2V 1.3B - 540P](https://huggingface.co/Skywork/SkyReels-V2-FLF2V-1.3B-540P-Diffusers) + +This model was contributed by [M. Tolga Cangöz](https://github.com/tolgacangoz). > [!TIP] > Click on the SkyReels-V2 models in the right sidebar for more examples of video generation. diff --git a/docs/source/en/api/pipelines/stable_diffusion/k_diffusion.md b/docs/source/en/api/pipelines/stable_diffusion/k_diffusion.md deleted file mode 100644 index 75f052b08f13..000000000000 --- a/docs/source/en/api/pipelines/stable_diffusion/k_diffusion.md +++ /dev/null @@ -1,30 +0,0 @@ - - -> [!WARNING] -> This pipeline is deprecated but it can still be used. However, we won't test the pipeline anymore and won't accept any changes to it. If you run into any issues, reinstall the last Diffusers version that supported this model. - -# K-Diffusion - -[k-diffusion](https://github.com/crowsonkb/k-diffusion) is a popular library created by [Katherine Crowson](https://github.com/crowsonkb/). We provide `StableDiffusionKDiffusionPipeline` and `StableDiffusionXLKDiffusionPipeline` that allow you to run Stable DIffusion with samplers from k-diffusion. - -Note that most the samplers from k-diffusion are implemented in Diffusers and we recommend using existing schedulers. You can find a mapping between k-diffusion samplers and schedulers in Diffusers [here](https://huggingface.co/docs/diffusers/api/schedulers/overview) - - -## StableDiffusionKDiffusionPipeline - -[[autodoc]] StableDiffusionKDiffusionPipeline - - -## StableDiffusionXLKDiffusionPipeline - -[[autodoc]] StableDiffusionXLKDiffusionPipeline \ No newline at end of file diff --git a/docs/source/en/api/pipelines/wan.md b/docs/source/en/api/pipelines/wan.md index 6aab6c5b33b9..d5fdbbfe0f95 100644 --- a/docs/source/en/api/pipelines/wan.md +++ b/docs/source/en/api/pipelines/wan.md @@ -250,9 +250,6 @@ The code snippets available in [this](https://github.com/huggingface/diffusers/p The general rule of thumb to keep in mind when preparing inputs for the VACE pipeline is that the input images, or frames of a video that you want to use for conditioning, should have a corresponding mask that is black in color. The black mask signifies that the model will not generate new content for that area, and only use those parts for conditioning the generation process. For parts/frames that should be generated by the model, the mask should be white in color. - - - ### Wan-Animate: Unified Character Animation and Replacement with Holistic Replication [Wan-Animate](https://huggingface.co/papers/2509.14055) by the Wan Team. diff --git a/docs/source/en/api/pipelines/z_image.md b/docs/source/en/api/pipelines/z_image.md index 5175f6b0fb6f..cf4c1aefb81f 100644 --- a/docs/source/en/api/pipelines/z_image.md +++ b/docs/source/en/api/pipelines/z_image.md @@ -53,6 +53,41 @@ image = pipe( image.save("zimage_img2img.png") ``` +## Inpainting + +Use [`ZImageInpaintPipeline`] to inpaint specific regions of an image based on a text prompt and mask. + +```python +import torch +import numpy as np +from PIL import Image +from diffusers import ZImageInpaintPipeline +from diffusers.utils import load_image + +pipe = ZImageInpaintPipeline.from_pretrained("Tongyi-MAI/Z-Image-Turbo", torch_dtype=torch.bfloat16) +pipe.to("cuda") + +url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg" +init_image = load_image(url).resize((1024, 1024)) + +# Create a mask (white = inpaint, black = preserve) +mask = np.zeros((1024, 1024), dtype=np.uint8) +mask[256:768, 256:768] = 255 # Inpaint center region +mask_image = Image.fromarray(mask) + +prompt = "A beautiful lake with mountains in the background" +image = pipe( + prompt, + image=init_image, + mask_image=mask_image, + strength=1.0, + num_inference_steps=9, + guidance_scale=0.0, + generator=torch.Generator("cuda").manual_seed(42), +).images[0] +image.save("zimage_inpaint.png") +``` + ## ZImagePipeline [[autodoc]] ZImagePipeline @@ -64,3 +99,9 @@ image.save("zimage_img2img.png") [[autodoc]] ZImageImg2ImgPipeline - all - __call__ + +## ZImageInpaintPipeline + +[[autodoc]] ZImageInpaintPipeline + - all + - __call__ diff --git a/docs/source/en/api/schedulers/helios.md b/docs/source/en/api/schedulers/helios.md new file mode 100644 index 000000000000..14c2be60bc89 --- /dev/null +++ b/docs/source/en/api/schedulers/helios.md @@ -0,0 +1,20 @@ + + +# HeliosScheduler + +`HeliosScheduler` is based on the pyramidal flow-matching sampling introduced in [Helios](https://huggingface.co/papers). + +## HeliosScheduler +[[autodoc]] HeliosScheduler + +scheduling_helios diff --git a/docs/source/en/api/schedulers/helios_dmd.md b/docs/source/en/api/schedulers/helios_dmd.md new file mode 100644 index 000000000000..4f075e8a7dfc --- /dev/null +++ b/docs/source/en/api/schedulers/helios_dmd.md @@ -0,0 +1,20 @@ + + +# HeliosDMDScheduler + +`HeliosDMDScheduler` is based on the pyramidal flow-matching sampling introduced in [Helios](https://huggingface.co/papers). + +## HeliosDMDScheduler +[[autodoc]] HeliosDMDScheduler + +scheduling_helios_dmd diff --git a/docs/source/en/hybrid_inference/api_reference.md b/docs/source/en/hybrid_inference/api_reference.md index 865aaba5ebb6..b538cb350481 100644 --- a/docs/source/en/hybrid_inference/api_reference.md +++ b/docs/source/en/hybrid_inference/api_reference.md @@ -1,9 +1,11 @@ -# Hybrid Inference API Reference +# Remote inference -## Remote Decode +Remote inference provides access to an [Inference Endpoint](https://huggingface.co/docs/inference-endpoints/index) to offload local generation requirements for decoding and encoding. + +## remote_decode [[autodoc]] utils.remote_utils.remote_decode -## Remote Encode +## remote_encode [[autodoc]] utils.remote_utils.remote_encode diff --git a/docs/source/en/hybrid_inference/overview.md b/docs/source/en/hybrid_inference/overview.md index 7ed1bbb88b3f..1384be9b7348 100644 --- a/docs/source/en/hybrid_inference/overview.md +++ b/docs/source/en/hybrid_inference/overview.md @@ -10,51 +10,296 @@ an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express o specific language governing permissions and limitations under the License. --> -# Hybrid Inference +# Remote inference -**Empowering local AI builders with Hybrid Inference** +> [!TIP] +> This is currently an experimental feature, and if you have any feedback, please feel free to leave it [here](https://github.com/huggingface/diffusers/issues/new?template=remote-vae-pilot-feedback.yml). +Remote inference offloads the decoding and encoding process to a remote endpoint to relax the memory requirements for local inference with large models. This feature is powered by [Inference Endpoints](https://huggingface.co/docs/inference-endpoints/index). Refer to the table below for the supported models and endpoint. -> [!TIP] -> Hybrid Inference is an [experimental feature](https://huggingface.co/blog/remote_vae). -> Feedback can be provided [here](https://github.com/huggingface/diffusers/issues/new?template=remote-vae-pilot-feedback.yml). +| Model | Endpoint | Checkpoint | Support | +|---|---|---|---| +| Stable Diffusion v1 | https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud | [stabilityai/sd-vae-ft-mse](https://huggingface.co/stabilityai/sd-vae-ft-mse) | encode/decode | +| Stable Diffusion XL | https://x2dmsqunjd6k9prw.us-east-1.aws.endpoints.huggingface.cloud | [madebyollin/sdxl-vae-fp16-fix](https://huggingface.co/madebyollin/sdxl-vae-fp16-fix) | encode/decode | +| Flux | https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud | [black-forest-labs/FLUX.1-schnell](https://huggingface.co/black-forest-labs/FLUX.1-schnell) | encode/decode | +| HunyuanVideo | https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud | [hunyuanvideo-community/HunyuanVideo](https://huggingface.co/hunyuanvideo-community/HunyuanVideo) | decode | + +This guide will show you how to encode and decode latents with remote inference. + +## Encoding + +Encoding converts images and videos into latent representations. Refer to the table below for the supported VAEs. + +Pass an image to [`~utils.remote_encode`] to encode it. The specific `scaling_factor` and `shift_factor` values for each model can be found in the [Remote inference](../hybrid_inference/api_reference) API reference. + +```py +import torch +from diffusers import FluxPipeline +from diffusers.utils import load_image +from diffusers.utils.remote_utils import remote_encode + +pipeline = FluxPipeline.from_pretrained( + "black-forest-labs/FLUX.1-schnell", + torch_dtype=torch.float16, + vae=None, + device_map="cuda" +) + +init_image = load_image( + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" +) +init_image = init_image.resize((768, 512)) + +init_latent = remote_encode( + endpoint="https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud", + image=init_image, + scaling_factor=0.3611, + shift_factor=0.1159 +) +``` + +## Decoding + +Decoding converts latent representations back into images or videos. Refer to the table below for the available and supported VAEs. + +Set the output type to `"latent"` in the pipeline and set the `vae` to `None`. Pass the latents to the [`~utils.remote_decode`] function. For Flux, the latents are packed so the `height` and `width` also need to be passed. The specific `scaling_factor` and `shift_factor` values for each model can be found in the [Remote inference](../hybrid_inference/api_reference) API reference. + + + + +```py +from diffusers import FluxPipeline + +pipeline = FluxPipeline.from_pretrained( + "black-forest-labs/FLUX.1-schnell", + torch_dtype=torch.bfloat16, + vae=None, + device_map="cuda" +) + +prompt = """ +A photorealistic Apollo-era photograph of a cat in a small astronaut suit with a bubble helmet, standing on the Moon and holding a flagpole planted in the dusty lunar soil. The flag shows a colorful paw-print emblem. Earth glows in the black sky above the stark gray surface, with sharp shadows and high-contrast lighting like vintage NASA photos. +""" + +latent = pipeline( + prompt=prompt, + guidance_scale=0.0, + num_inference_steps=4, + output_type="latent", +).images +image = remote_decode( + endpoint="https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud/", + tensor=latent, + height=1024, + width=1024, + scaling_factor=0.3611, + shift_factor=0.1159, +) +image.save("image.jpg") +``` + + + + +```py +import torch +from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel + +transformer = HunyuanVideoTransformer3DModel.from_pretrained( + "hunyuanvideo-community/HunyuanVideo", subfolder="transformer", torch_dtype=torch.bfloat16 +) +pipeline = HunyuanVideoPipeline.from_pretrained( + model_id, transformer=transformer, vae=None, torch_dtype=torch.float16, device_map="cuda" +) + +latent = pipeline( + prompt="A cat walks on the grass, realistic", + height=320, + width=512, + num_frames=61, + num_inference_steps=30, + output_type="latent", +).frames + +video = remote_decode( + endpoint="https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud/", + tensor=latent, + output_type="mp4", +) + +if isinstance(video, bytes): + with open("video.mp4", "wb") as f: + f.write(video) +``` + + + + +## Queuing + +Remote inference supports queuing to process multiple generation requests. While the current latent is being decoded, you can queue the next prompt. + +```py +import queue +import threading +from IPython.display import display +from diffusers import StableDiffusionXLPipeline + +def decode_worker(q: queue.Queue): + while True: + item = q.get() + if item is None: + break + image = remote_decode( + endpoint="https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud/", + tensor=item, + scaling_factor=0.13025, + ) + display(image) + q.task_done() + +q = queue.Queue() +thread = threading.Thread(target=decode_worker, args=(q,), daemon=True) +thread.start() + +def decode(latent: torch.Tensor): + q.put(latent) + +prompts = [ + "A grainy Apollo-era style photograph of a cat in a snug astronaut suit with a bubble helmet, standing on the lunar surface and gripping a flag with a paw-print emblem. The gray Moon landscape stretches behind it, Earth glowing vividly in the black sky, shadows crisp and high-contrast.", + "A vintage 1960s sci-fi pulp magazine cover illustration of a heroic cat astronaut planting a flag on the Moon. Bold, saturated colors, exaggerated space gear, playful typography floating in the background, Earth painted in bright blues and greens.", + "A hyper-detailed cinematic shot of a cat astronaut on the Moon holding a fluttering flag, fur visible through the helmet glass, lunar dust scattering under its feet. The vastness of space and Earth in the distance create an epic, awe-inspiring tone.", + "A colorful cartoon drawing of a happy cat wearing a chunky, oversized spacesuit, proudly holding a flag with a big paw print on it. The Moon’s surface is simplified with craters drawn like doodles, and Earth in the sky has a smiling face.", + "A monochrome 1969-style press photo of a “first cat on the Moon” moment. The cat, in a tiny astronaut suit, stands by a planted flag, with grainy textures, scratches, and a blurred Earth in the background, mimicking old archival space photos." +] + + +pipeline = StableDiffusionXLPipeline.from_pretrained( + "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0", + torch_dtype=torch.float16, + vae=None, + device_map="cuda" +) + +pipeline.unet = pipeline.unet.to(memory_format=torch.channels_last) +pipeline.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) + +_ = pipeline( + prompt=prompts[0], + output_type="latent", +) + +for prompt in prompts: + latent = pipeline( + prompt=prompt, + output_type="latent", + ).images + decode(latent) + +q.put(None) +thread.join() +``` + +## Benchmarks + +The tables demonstrate the memory requirements for encoding and decoding with Stable Diffusion v1.5 and SDXL on different GPUs. +For the majority of these GPUs, the memory usage dictates whether other models (text encoders, UNet/transformer) need to be offloaded or required tiled encoding. The latter two techniques increases inference time and impacts quality. +
Encoding - Stable Diffusion v1.5 -## Why use Hybrid Inference? +| GPU | Resolution | Time (seconds) | Memory (%) | Tiled Time (secs) | Tiled Memory (%) | +|:------------------------------|:-------------|-----------------:|-------------:|--------------------:|-------------------:| +| NVIDIA GeForce RTX 4090 | 512x512 | 0.015 | 3.51901 | 0.015 | 3.51901 | +| NVIDIA GeForce RTX 4090 | 256x256 | 0.004 | 1.3154 | 0.005 | 1.3154 | +| NVIDIA GeForce RTX 4090 | 2048x2048 | 0.402 | 47.1852 | 0.496 | 3.51901 | +| NVIDIA GeForce RTX 4090 | 1024x1024 | 0.078 | 12.2658 | 0.094 | 3.51901 | +| NVIDIA GeForce RTX 4080 SUPER | 512x512 | 0.023 | 5.30105 | 0.023 | 5.30105 | +| NVIDIA GeForce RTX 4080 SUPER | 256x256 | 0.006 | 1.98152 | 0.006 | 1.98152 | +| NVIDIA GeForce RTX 4080 SUPER | 2048x2048 | 0.574 | 71.08 | 0.656 | 5.30105 | +| NVIDIA GeForce RTX 4080 SUPER | 1024x1024 | 0.111 | 18.4772 | 0.14 | 5.30105 | +| NVIDIA GeForce RTX 3090 | 512x512 | 0.032 | 3.52782 | 0.032 | 3.52782 | +| NVIDIA GeForce RTX 3090 | 256x256 | 0.01 | 1.31869 | 0.009 | 1.31869 | +| NVIDIA GeForce RTX 3090 | 2048x2048 | 0.742 | 47.3033 | 0.954 | 3.52782 | +| NVIDIA GeForce RTX 3090 | 1024x1024 | 0.136 | 12.2965 | 0.207 | 3.52782 | +| NVIDIA GeForce RTX 3080 | 512x512 | 0.036 | 8.51761 | 0.036 | 8.51761 | +| NVIDIA GeForce RTX 3080 | 256x256 | 0.01 | 3.18387 | 0.01 | 3.18387 | +| NVIDIA GeForce RTX 3080 | 2048x2048 | 0.863 | 86.7424 | 1.191 | 8.51761 | +| NVIDIA GeForce RTX 3080 | 1024x1024 | 0.157 | 29.6888 | 0.227 | 8.51761 | +| NVIDIA GeForce RTX 3070 | 512x512 | 0.051 | 10.6941 | 0.051 | 10.6941 | +| NVIDIA GeForce RTX 3070 | 256x256 | 0.015 | 3.99743 | 0.015 | 3.99743 | +| NVIDIA GeForce RTX 3070 | 2048x2048 | 1.217 | 96.054 | 1.482 | 10.6941 | +| NVIDIA GeForce RTX 3070 | 1024x1024 | 0.223 | 37.2751 | 0.327 | 10.6941 | -Hybrid Inference offers a fast and simple way to offload local generation requirements. +
-- 🚀 **Reduced Requirements:** Access powerful models without expensive hardware. -- 💎 **Without Compromise:** Achieve the highest quality without sacrificing performance. -- 💰 **Cost Effective:** It's free! 🤑 -- 🎯 **Diverse Use Cases:** Fully compatible with Diffusers 🧨 and the wider community. -- 🔧 **Developer-Friendly:** Simple requests, fast responses. +
Encoding SDXL ---- +| GPU | Resolution | Time (seconds) | Memory Consumed (%) | Tiled Time (seconds) | Tiled Memory (%) | +|:------------------------------|:-------------|-----------------:|----------------------:|-----------------------:|-------------------:| +| NVIDIA GeForce RTX 4090 | 512x512 | 0.029 | 4.95707 | 0.029 | 4.95707 | +| NVIDIA GeForce RTX 4090 | 256x256 | 0.007 | 2.29666 | 0.007 | 2.29666 | +| NVIDIA GeForce RTX 4090 | 2048x2048 | 0.873 | 66.3452 | 0.863 | 15.5649 | +| NVIDIA GeForce RTX 4090 | 1024x1024 | 0.142 | 15.5479 | 0.143 | 15.5479 | +| NVIDIA GeForce RTX 4080 SUPER | 512x512 | 0.044 | 7.46735 | 0.044 | 7.46735 | +| NVIDIA GeForce RTX 4080 SUPER | 256x256 | 0.01 | 3.4597 | 0.01 | 3.4597 | +| NVIDIA GeForce RTX 4080 SUPER | 2048x2048 | 1.317 | 87.1615 | 1.291 | 23.447 | +| NVIDIA GeForce RTX 4080 SUPER | 1024x1024 | 0.213 | 23.4215 | 0.214 | 23.4215 | +| NVIDIA GeForce RTX 3090 | 512x512 | 0.058 | 5.65638 | 0.058 | 5.65638 | +| NVIDIA GeForce RTX 3090 | 256x256 | 0.016 | 2.45081 | 0.016 | 2.45081 | +| NVIDIA GeForce RTX 3090 | 2048x2048 | 1.755 | 77.8239 | 1.614 | 18.4193 | +| NVIDIA GeForce RTX 3090 | 1024x1024 | 0.265 | 18.4023 | 0.265 | 18.4023 | +| NVIDIA GeForce RTX 3080 | 512x512 | 0.064 | 13.6568 | 0.064 | 13.6568 | +| NVIDIA GeForce RTX 3080 | 256x256 | 0.018 | 5.91728 | 0.018 | 5.91728 | +| NVIDIA GeForce RTX 3080 | 2048x2048 | OOM | OOM | 1.866 | 44.4717 | +| NVIDIA GeForce RTX 3080 | 1024x1024 | 0.302 | 44.4308 | 0.302 | 44.4308 | +| NVIDIA GeForce RTX 3070 | 512x512 | 0.093 | 17.1465 | 0.093 | 17.1465 | +| NVIDIA GeForce RTX 3070 | 256x256 | 0.025 | 7.42931 | 0.026 | 7.42931 | +| NVIDIA GeForce RTX 3070 | 2048x2048 | OOM | OOM | 2.674 | 55.8355 | +| NVIDIA GeForce RTX 3070 | 1024x1024 | 0.443 | 55.7841 | 0.443 | 55.7841 | -## Available Models +
-* **VAE Decode 🖼️:** Quickly decode latent representations into high-quality images without compromising performance or workflow speed. -* **VAE Encode 🔢:** Efficiently encode images into latent representations for generation and training. -* **Text Encoders 📃 (coming soon):** Compute text embeddings for your prompts quickly and accurately, ensuring a smooth and high-quality workflow. +
Decoding - Stable Diffusion v1.5 ---- +| GPU | Resolution | Time (seconds) | Memory (%) | Tiled Time (secs) | Tiled Memory (%) | +| --- | --- | --- | --- | --- | --- | +| NVIDIA GeForce RTX 4090 | 512x512 | 0.031 | 5.60% | 0.031 (0%) | 5.60% | +| NVIDIA GeForce RTX 4090 | 1024x1024 | 0.148 | 20.00% | 0.301 (+103%) | 5.60% | +| NVIDIA GeForce RTX 4080 | 512x512 | 0.05 | 8.40% | 0.050 (0%) | 8.40% | +| NVIDIA GeForce RTX 4080 | 1024x1024 | 0.224 | 30.00% | 0.356 (+59%) | 8.40% | +| NVIDIA GeForce RTX 4070 Ti | 512x512 | 0.066 | 11.30% | 0.066 (0%) | 11.30% | +| NVIDIA GeForce RTX 4070 Ti | 1024x1024 | 0.284 | 40.50% | 0.454 (+60%) | 11.40% | +| NVIDIA GeForce RTX 3090 | 512x512 | 0.062 | 5.20% | 0.062 (0%) | 5.20% | +| NVIDIA GeForce RTX 3090 | 1024x1024 | 0.253 | 18.50% | 0.464 (+83%) | 5.20% | +| NVIDIA GeForce RTX 3080 | 512x512 | 0.07 | 12.80% | 0.070 (0%) | 12.80% | +| NVIDIA GeForce RTX 3080 | 1024x1024 | 0.286 | 45.30% | 0.466 (+63%) | 12.90% | +| NVIDIA GeForce RTX 3070 | 512x512 | 0.102 | 15.90% | 0.102 (0%) | 15.90% | +| NVIDIA GeForce RTX 3070 | 1024x1024 | 0.421 | 56.30% | 0.746 (+77%) | 16.00% | -## Integrations +
-* **[SD.Next](https://github.com/vladmandic/sdnext):** All-in-one UI with direct supports Hybrid Inference. -* **[ComfyUI-HFRemoteVae](https://github.com/kijai/ComfyUI-HFRemoteVae):** ComfyUI node for Hybrid Inference. +
Decoding SDXL -## Changelog +| GPU | Resolution | Time (seconds) | Memory Consumed (%) | Tiled Time (seconds) | Tiled Memory (%) | +| --- | --- | --- | --- | --- | --- | +| NVIDIA GeForce RTX 4090 | 512x512 | 0.057 | 10.00% | 0.057 (0%) | 10.00% | +| NVIDIA GeForce RTX 4090 | 1024x1024 | 0.256 | 35.50% | 0.257 (+0.4%) | 35.50% | +| NVIDIA GeForce RTX 4080 | 512x512 | 0.092 | 15.00% | 0.092 (0%) | 15.00% | +| NVIDIA GeForce RTX 4080 | 1024x1024 | 0.406 | 53.30% | 0.406 (0%) | 53.30% | +| NVIDIA GeForce RTX 4070 Ti | 512x512 | 0.121 | 20.20% | 0.120 (-0.8%) | 20.20% | +| NVIDIA GeForce RTX 4070 Ti | 1024x1024 | 0.519 | 72.00% | 0.519 (0%) | 72.00% | +| NVIDIA GeForce RTX 3090 | 512x512 | 0.107 | 10.50% | 0.107 (0%) | 10.50% | +| NVIDIA GeForce RTX 3090 | 1024x1024 | 0.459 | 38.00% | 0.460 (+0.2%) | 38.00% | +| NVIDIA GeForce RTX 3080 | 512x512 | 0.121 | 25.60% | 0.121 (0%) | 25.60% | +| NVIDIA GeForce RTX 3080 | 1024x1024 | 0.524 | 93.00% | 0.524 (0%) | 93.00% | +| NVIDIA GeForce RTX 3070 | 512x512 | 0.183 | 31.80% | 0.183 (0%) | 31.80% | +| NVIDIA GeForce RTX 3070 | 1024x1024 | 0.794 | 96.40% | 0.794 (0%) | 96.40% | -- March 10 2025: Added VAE encode -- March 2 2025: Initial release with VAE decoding +
-## Contents -The documentation is organized into three sections: +## Resources -* **VAE Decode** Learn the basics of how to use VAE Decode with Hybrid Inference. -* **VAE Encode** Learn the basics of how to use VAE Encode with Hybrid Inference. -* **API Reference** Dive into task-specific settings and parameters. +- Remote inference is also supported in [SD.Next](https://github.com/vladmandic/sdnext) and [ComfyUI-HFRemoteVae](https://github.com/kijai/ComfyUI-HFRemoteVae). +- Refer to the [Remote VAEs for decoding with Inference Endpoints](https://huggingface.co/blog/remote_vae) blog post to learn more. \ No newline at end of file diff --git a/docs/source/en/hybrid_inference/vae_decode.md b/docs/source/en/hybrid_inference/vae_decode.md deleted file mode 100644 index 1457090550c7..000000000000 --- a/docs/source/en/hybrid_inference/vae_decode.md +++ /dev/null @@ -1,345 +0,0 @@ -# Getting Started: VAE Decode with Hybrid Inference - -VAE decode is an essential component of diffusion models - turning latent representations into images or videos. - -## Memory - -These tables demonstrate the VRAM requirements for VAE decode with SD v1 and SD XL on different GPUs. - -For the majority of these GPUs the memory usage % dictates other models (text encoders, UNet/Transformer) must be offloaded, or tiled decoding has to be used which increases time taken and impacts quality. - -
SD v1.5 - -| GPU | Resolution | Time (seconds) | Memory (%) | Tiled Time (secs) | Tiled Memory (%) | -| --- | --- | --- | --- | --- | --- | -| NVIDIA GeForce RTX 4090 | 512x512 | 0.031 | 5.60% | 0.031 (0%) | 5.60% | -| NVIDIA GeForce RTX 4090 | 1024x1024 | 0.148 | 20.00% | 0.301 (+103%) | 5.60% | -| NVIDIA GeForce RTX 4080 | 512x512 | 0.05 | 8.40% | 0.050 (0%) | 8.40% | -| NVIDIA GeForce RTX 4080 | 1024x1024 | 0.224 | 30.00% | 0.356 (+59%) | 8.40% | -| NVIDIA GeForce RTX 4070 Ti | 512x512 | 0.066 | 11.30% | 0.066 (0%) | 11.30% | -| NVIDIA GeForce RTX 4070 Ti | 1024x1024 | 0.284 | 40.50% | 0.454 (+60%) | 11.40% | -| NVIDIA GeForce RTX 3090 | 512x512 | 0.062 | 5.20% | 0.062 (0%) | 5.20% | -| NVIDIA GeForce RTX 3090 | 1024x1024 | 0.253 | 18.50% | 0.464 (+83%) | 5.20% | -| NVIDIA GeForce RTX 3080 | 512x512 | 0.07 | 12.80% | 0.070 (0%) | 12.80% | -| NVIDIA GeForce RTX 3080 | 1024x1024 | 0.286 | 45.30% | 0.466 (+63%) | 12.90% | -| NVIDIA GeForce RTX 3070 | 512x512 | 0.102 | 15.90% | 0.102 (0%) | 15.90% | -| NVIDIA GeForce RTX 3070 | 1024x1024 | 0.421 | 56.30% | 0.746 (+77%) | 16.00% | - -
- -
SDXL - -| GPU | Resolution | Time (seconds) | Memory Consumed (%) | Tiled Time (seconds) | Tiled Memory (%) | -| --- | --- | --- | --- | --- | --- | -| NVIDIA GeForce RTX 4090 | 512x512 | 0.057 | 10.00% | 0.057 (0%) | 10.00% | -| NVIDIA GeForce RTX 4090 | 1024x1024 | 0.256 | 35.50% | 0.257 (+0.4%) | 35.50% | -| NVIDIA GeForce RTX 4080 | 512x512 | 0.092 | 15.00% | 0.092 (0%) | 15.00% | -| NVIDIA GeForce RTX 4080 | 1024x1024 | 0.406 | 53.30% | 0.406 (0%) | 53.30% | -| NVIDIA GeForce RTX 4070 Ti | 512x512 | 0.121 | 20.20% | 0.120 (-0.8%) | 20.20% | -| NVIDIA GeForce RTX 4070 Ti | 1024x1024 | 0.519 | 72.00% | 0.519 (0%) | 72.00% | -| NVIDIA GeForce RTX 3090 | 512x512 | 0.107 | 10.50% | 0.107 (0%) | 10.50% | -| NVIDIA GeForce RTX 3090 | 1024x1024 | 0.459 | 38.00% | 0.460 (+0.2%) | 38.00% | -| NVIDIA GeForce RTX 3080 | 512x512 | 0.121 | 25.60% | 0.121 (0%) | 25.60% | -| NVIDIA GeForce RTX 3080 | 1024x1024 | 0.524 | 93.00% | 0.524 (0%) | 93.00% | -| NVIDIA GeForce RTX 3070 | 512x512 | 0.183 | 31.80% | 0.183 (0%) | 31.80% | -| NVIDIA GeForce RTX 3070 | 1024x1024 | 0.794 | 96.40% | 0.794 (0%) | 96.40% | - -
- -## Available VAEs - -| | **Endpoint** | **Model** | -|:-:|:-----------:|:--------:| -| **Stable Diffusion v1** | [https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud](https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud) | [`stabilityai/sd-vae-ft-mse`](https://hf.co/stabilityai/sd-vae-ft-mse) | -| **Stable Diffusion XL** | [https://x2dmsqunjd6k9prw.us-east-1.aws.endpoints.huggingface.cloud](https://x2dmsqunjd6k9prw.us-east-1.aws.endpoints.huggingface.cloud) | [`madebyollin/sdxl-vae-fp16-fix`](https://hf.co/madebyollin/sdxl-vae-fp16-fix) | -| **Flux** | [https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud](https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud) | [`black-forest-labs/FLUX.1-schnell`](https://hf.co/black-forest-labs/FLUX.1-schnell) | -| **HunyuanVideo** | [https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud](https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud) | [`hunyuanvideo-community/HunyuanVideo`](https://hf.co/hunyuanvideo-community/HunyuanVideo) | - - -> [!TIP] -> Model support can be requested [here](https://github.com/huggingface/diffusers/issues/new?template=remote-vae-pilot-feedback.yml). - - -## Code - -> [!TIP] -> Install `diffusers` from `main` to run the code: `pip install git+https://github.com/huggingface/diffusers@main` - - -A helper method simplifies interacting with Hybrid Inference. - -```python -from diffusers.utils.remote_utils import remote_decode -``` - -### Basic example - -Here, we show how to use the remote VAE on random tensors. - -
Code - -```python -image = remote_decode( - endpoint="https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud/", - tensor=torch.randn([1, 4, 64, 64], dtype=torch.float16), - scaling_factor=0.18215, -) -``` - -
- -
- -
- -Usage for Flux is slightly different. Flux latents are packed so we need to send the `height` and `width`. - -
Code - -```python -image = remote_decode( - endpoint="https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud/", - tensor=torch.randn([1, 4096, 64], dtype=torch.float16), - height=1024, - width=1024, - scaling_factor=0.3611, - shift_factor=0.1159, -) -``` - -
- -
- -
- -Finally, an example for HunyuanVideo. - -
Code - -```python -video = remote_decode( - endpoint="https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud/", - tensor=torch.randn([1, 16, 3, 40, 64], dtype=torch.float16), - output_type="mp4", -) -with open("video.mp4", "wb") as f: - f.write(video) -``` - -
- -
- -
- - -### Generation - -But we want to use the VAE on an actual pipeline to get an actual image, not random noise. The example below shows how to do it with SD v1.5. - -
Code - -```python -from diffusers import StableDiffusionPipeline - -pipe = StableDiffusionPipeline.from_pretrained( - "stable-diffusion-v1-5/stable-diffusion-v1-5", - torch_dtype=torch.float16, - variant="fp16", - vae=None, -).to("cuda") - -prompt = "Strawberry ice cream, in a stylish modern glass, coconut, splashing milk cream and honey, in a gradient purple background, fluid motion, dynamic movement, cinematic lighting, Mysterious" - -latent = pipe( - prompt=prompt, - output_type="latent", -).images -image = remote_decode( - endpoint="https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud/", - tensor=latent, - scaling_factor=0.18215, -) -image.save("test.jpg") -``` - -
- -
- -
- -Here’s another example with Flux. - -
Code - -```python -from diffusers import FluxPipeline - -pipe = FluxPipeline.from_pretrained( - "black-forest-labs/FLUX.1-schnell", - torch_dtype=torch.bfloat16, - vae=None, -).to("cuda") - -prompt = "Strawberry ice cream, in a stylish modern glass, coconut, splashing milk cream and honey, in a gradient purple background, fluid motion, dynamic movement, cinematic lighting, Mysterious" - -latent = pipe( - prompt=prompt, - guidance_scale=0.0, - num_inference_steps=4, - output_type="latent", -).images -image = remote_decode( - endpoint="https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud/", - tensor=latent, - height=1024, - width=1024, - scaling_factor=0.3611, - shift_factor=0.1159, -) -image.save("test.jpg") -``` - -
- -
- -
- -Here’s an example with HunyuanVideo. - -
Code - -```python -from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel - -model_id = "hunyuanvideo-community/HunyuanVideo" -transformer = HunyuanVideoTransformer3DModel.from_pretrained( - model_id, subfolder="transformer", torch_dtype=torch.bfloat16 -) -pipe = HunyuanVideoPipeline.from_pretrained( - model_id, transformer=transformer, vae=None, torch_dtype=torch.float16 -).to("cuda") - -latent = pipe( - prompt="A cat walks on the grass, realistic", - height=320, - width=512, - num_frames=61, - num_inference_steps=30, - output_type="latent", -).frames - -video = remote_decode( - endpoint="https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud/", - tensor=latent, - output_type="mp4", -) - -if isinstance(video, bytes): - with open("video.mp4", "wb") as f: - f.write(video) -``` - -
- -
- -
- - -### Queueing - -One of the great benefits of using a remote VAE is that we can queue multiple generation requests. While the current latent is being processed for decoding, we can already queue another one. This helps improve concurrency. - - -
Code - -```python -import queue -import threading -from IPython.display import display -from diffusers import StableDiffusionPipeline - -def decode_worker(q: queue.Queue): - while True: - item = q.get() - if item is None: - break - image = remote_decode( - endpoint="https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud/", - tensor=item, - scaling_factor=0.18215, - ) - display(image) - q.task_done() - -q = queue.Queue() -thread = threading.Thread(target=decode_worker, args=(q,), daemon=True) -thread.start() - -def decode(latent: torch.Tensor): - q.put(latent) - -prompts = [ - "Blueberry ice cream, in a stylish modern glass , ice cubes, nuts, mint leaves, splashing milk cream, in a gradient purple background, fluid motion, dynamic movement, cinematic lighting, Mysterious", - "Lemonade in a glass, mint leaves, in an aqua and white background, flowers, ice cubes, halo, fluid motion, dynamic movement, soft lighting, digital painting, rule of thirds composition, Art by Greg rutkowski, Coby whitmore", - "Comic book art, beautiful, vintage, pastel neon colors, extremely detailed pupils, delicate features, light on face, slight smile, Artgerm, Mary Blair, Edmund Dulac, long dark locks, bangs, glowing, fashionable style, fairytale ambience, hot pink.", - "Masterpiece, vanilla cone ice cream garnished with chocolate syrup, crushed nuts, choco flakes, in a brown background, gold, cinematic lighting, Art by WLOP", - "A bowl of milk, falling cornflakes, berries, blueberries, in a white background, soft lighting, intricate details, rule of thirds, octane render, volumetric lighting", - "Cold Coffee with cream, crushed almonds, in a glass, choco flakes, ice cubes, wet, in a wooden background, cinematic lighting, hyper realistic painting, art by Carne Griffiths, octane render, volumetric lighting, fluid motion, dynamic movement, muted colors,", -] - -pipe = StableDiffusionPipeline.from_pretrained( - "Lykon/dreamshaper-8", - torch_dtype=torch.float16, - vae=None, -).to("cuda") - -pipe.unet = pipe.unet.to(memory_format=torch.channels_last) -pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) - -_ = pipe( - prompt=prompts[0], - output_type="latent", -) - -for prompt in prompts: - latent = pipe( - prompt=prompt, - output_type="latent", - ).images - decode(latent) - -q.put(None) -thread.join() -``` - -
- - -
- -
- -## Integrations - -* **[SD.Next](https://github.com/vladmandic/sdnext):** All-in-one UI with direct supports Hybrid Inference. -* **[ComfyUI-HFRemoteVae](https://github.com/kijai/ComfyUI-HFRemoteVae):** ComfyUI node for Hybrid Inference. diff --git a/docs/source/en/hybrid_inference/vae_encode.md b/docs/source/en/hybrid_inference/vae_encode.md deleted file mode 100644 index dd285fa25c03..000000000000 --- a/docs/source/en/hybrid_inference/vae_encode.md +++ /dev/null @@ -1,183 +0,0 @@ -# Getting Started: VAE Encode with Hybrid Inference - -VAE encode is used for training, image-to-image and image-to-video - turning into images or videos into latent representations. - -## Memory - -These tables demonstrate the VRAM requirements for VAE encode with SD v1 and SD XL on different GPUs. - -For the majority of these GPUs the memory usage % dictates other models (text encoders, UNet/Transformer) must be offloaded, or tiled encoding has to be used which increases time taken and impacts quality. - -
SD v1.5 - -| GPU | Resolution | Time (seconds) | Memory (%) | Tiled Time (secs) | Tiled Memory (%) | -|:------------------------------|:-------------|-----------------:|-------------:|--------------------:|-------------------:| -| NVIDIA GeForce RTX 4090 | 512x512 | 0.015 | 3.51901 | 0.015 | 3.51901 | -| NVIDIA GeForce RTX 4090 | 256x256 | 0.004 | 1.3154 | 0.005 | 1.3154 | -| NVIDIA GeForce RTX 4090 | 2048x2048 | 0.402 | 47.1852 | 0.496 | 3.51901 | -| NVIDIA GeForce RTX 4090 | 1024x1024 | 0.078 | 12.2658 | 0.094 | 3.51901 | -| NVIDIA GeForce RTX 4080 SUPER | 512x512 | 0.023 | 5.30105 | 0.023 | 5.30105 | -| NVIDIA GeForce RTX 4080 SUPER | 256x256 | 0.006 | 1.98152 | 0.006 | 1.98152 | -| NVIDIA GeForce RTX 4080 SUPER | 2048x2048 | 0.574 | 71.08 | 0.656 | 5.30105 | -| NVIDIA GeForce RTX 4080 SUPER | 1024x1024 | 0.111 | 18.4772 | 0.14 | 5.30105 | -| NVIDIA GeForce RTX 3090 | 512x512 | 0.032 | 3.52782 | 0.032 | 3.52782 | -| NVIDIA GeForce RTX 3090 | 256x256 | 0.01 | 1.31869 | 0.009 | 1.31869 | -| NVIDIA GeForce RTX 3090 | 2048x2048 | 0.742 | 47.3033 | 0.954 | 3.52782 | -| NVIDIA GeForce RTX 3090 | 1024x1024 | 0.136 | 12.2965 | 0.207 | 3.52782 | -| NVIDIA GeForce RTX 3080 | 512x512 | 0.036 | 8.51761 | 0.036 | 8.51761 | -| NVIDIA GeForce RTX 3080 | 256x256 | 0.01 | 3.18387 | 0.01 | 3.18387 | -| NVIDIA GeForce RTX 3080 | 2048x2048 | 0.863 | 86.7424 | 1.191 | 8.51761 | -| NVIDIA GeForce RTX 3080 | 1024x1024 | 0.157 | 29.6888 | 0.227 | 8.51761 | -| NVIDIA GeForce RTX 3070 | 512x512 | 0.051 | 10.6941 | 0.051 | 10.6941 | -| NVIDIA GeForce RTX 3070 | 256x256 | 0.015 | 3.99743 | 0.015 | 3.99743 | -| NVIDIA GeForce RTX 3070 | 2048x2048 | 1.217 | 96.054 | 1.482 | 10.6941 | -| NVIDIA GeForce RTX 3070 | 1024x1024 | 0.223 | 37.2751 | 0.327 | 10.6941 | - - -
- -
SDXL - -| GPU | Resolution | Time (seconds) | Memory Consumed (%) | Tiled Time (seconds) | Tiled Memory (%) | -|:------------------------------|:-------------|-----------------:|----------------------:|-----------------------:|-------------------:| -| NVIDIA GeForce RTX 4090 | 512x512 | 0.029 | 4.95707 | 0.029 | 4.95707 | -| NVIDIA GeForce RTX 4090 | 256x256 | 0.007 | 2.29666 | 0.007 | 2.29666 | -| NVIDIA GeForce RTX 4090 | 2048x2048 | 0.873 | 66.3452 | 0.863 | 15.5649 | -| NVIDIA GeForce RTX 4090 | 1024x1024 | 0.142 | 15.5479 | 0.143 | 15.5479 | -| NVIDIA GeForce RTX 4080 SUPER | 512x512 | 0.044 | 7.46735 | 0.044 | 7.46735 | -| NVIDIA GeForce RTX 4080 SUPER | 256x256 | 0.01 | 3.4597 | 0.01 | 3.4597 | -| NVIDIA GeForce RTX 4080 SUPER | 2048x2048 | 1.317 | 87.1615 | 1.291 | 23.447 | -| NVIDIA GeForce RTX 4080 SUPER | 1024x1024 | 0.213 | 23.4215 | 0.214 | 23.4215 | -| NVIDIA GeForce RTX 3090 | 512x512 | 0.058 | 5.65638 | 0.058 | 5.65638 | -| NVIDIA GeForce RTX 3090 | 256x256 | 0.016 | 2.45081 | 0.016 | 2.45081 | -| NVIDIA GeForce RTX 3090 | 2048x2048 | 1.755 | 77.8239 | 1.614 | 18.4193 | -| NVIDIA GeForce RTX 3090 | 1024x1024 | 0.265 | 18.4023 | 0.265 | 18.4023 | -| NVIDIA GeForce RTX 3080 | 512x512 | 0.064 | 13.6568 | 0.064 | 13.6568 | -| NVIDIA GeForce RTX 3080 | 256x256 | 0.018 | 5.91728 | 0.018 | 5.91728 | -| NVIDIA GeForce RTX 3080 | 2048x2048 | OOM | OOM | 1.866 | 44.4717 | -| NVIDIA GeForce RTX 3080 | 1024x1024 | 0.302 | 44.4308 | 0.302 | 44.4308 | -| NVIDIA GeForce RTX 3070 | 512x512 | 0.093 | 17.1465 | 0.093 | 17.1465 | -| NVIDIA GeForce RTX 3070 | 256x256 | 0.025 | 7.42931 | 0.026 | 7.42931 | -| NVIDIA GeForce RTX 3070 | 2048x2048 | OOM | OOM | 2.674 | 55.8355 | -| NVIDIA GeForce RTX 3070 | 1024x1024 | 0.443 | 55.7841 | 0.443 | 55.7841 | - -
- -## Available VAEs - -| | **Endpoint** | **Model** | -|:-:|:-----------:|:--------:| -| **Stable Diffusion v1** | [https://qc6479g0aac6qwy9.us-east-1.aws.endpoints.huggingface.cloud](https://qc6479g0aac6qwy9.us-east-1.aws.endpoints.huggingface.cloud) | [`stabilityai/sd-vae-ft-mse`](https://hf.co/stabilityai/sd-vae-ft-mse) | -| **Stable Diffusion XL** | [https://xjqqhmyn62rog84g.us-east-1.aws.endpoints.huggingface.cloud](https://xjqqhmyn62rog84g.us-east-1.aws.endpoints.huggingface.cloud) | [`madebyollin/sdxl-vae-fp16-fix`](https://hf.co/madebyollin/sdxl-vae-fp16-fix) | -| **Flux** | [https://ptccx55jz97f9zgo.us-east-1.aws.endpoints.huggingface.cloud](https://ptccx55jz97f9zgo.us-east-1.aws.endpoints.huggingface.cloud) | [`black-forest-labs/FLUX.1-schnell`](https://hf.co/black-forest-labs/FLUX.1-schnell) | - - -> [!TIP] -> Model support can be requested [here](https://github.com/huggingface/diffusers/issues/new?template=remote-vae-pilot-feedback.yml). - - -## Code - -> [!TIP] -> Install `diffusers` from `main` to run the code: `pip install git+https://github.com/huggingface/diffusers@main` - - -A helper method simplifies interacting with Hybrid Inference. - -```python -from diffusers.utils.remote_utils import remote_encode -``` - -### Basic example - -Let's encode an image, then decode it to demonstrate. - -
- -
- -
Code - -```python -from diffusers.utils import load_image -from diffusers.utils.remote_utils import remote_decode - -image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg?download=true") - -latent = remote_encode( - endpoint="https://ptccx55jz97f9zgo.us-east-1.aws.endpoints.huggingface.cloud/", - scaling_factor=0.3611, - shift_factor=0.1159, -) - -decoded = remote_decode( - endpoint="https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud/", - tensor=latent, - scaling_factor=0.3611, - shift_factor=0.1159, -) -``` - -
- -
- -
- - -### Generation - -Now let's look at a generation example, we'll encode the image, generate then remotely decode too! - -
Code - -```python -import torch -from diffusers import StableDiffusionImg2ImgPipeline -from diffusers.utils import load_image -from diffusers.utils.remote_utils import remote_decode, remote_encode - -pipe = StableDiffusionImg2ImgPipeline.from_pretrained( - "stable-diffusion-v1-5/stable-diffusion-v1-5", - torch_dtype=torch.float16, - variant="fp16", - vae=None, -).to("cuda") - -init_image = load_image( - "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg" -) -init_image = init_image.resize((768, 512)) - -init_latent = remote_encode( - endpoint="https://qc6479g0aac6qwy9.us-east-1.aws.endpoints.huggingface.cloud/", - image=init_image, - scaling_factor=0.18215, -) - -prompt = "A fantasy landscape, trending on artstation" -latent = pipe( - prompt=prompt, - image=init_latent, - strength=0.75, - output_type="latent", -).images - -image = remote_decode( - endpoint="https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud/", - tensor=latent, - scaling_factor=0.18215, -) -image.save("fantasy_landscape.jpg") -``` - -
- -
- -
- -## Integrations - -* **[SD.Next](https://github.com/vladmandic/sdnext):** All-in-one UI with direct supports Hybrid Inference. -* **[ComfyUI-HFRemoteVae](https://github.com/kijai/ComfyUI-HFRemoteVae):** ComfyUI node for Hybrid Inference. diff --git a/docs/source/en/modular_diffusers/auto_pipeline_blocks.md b/docs/source/en/modular_diffusers/auto_pipeline_blocks.md index 2d4d82c735bd..1bcf1d691036 100644 --- a/docs/source/en/modular_diffusers/auto_pipeline_blocks.md +++ b/docs/source/en/modular_diffusers/auto_pipeline_blocks.md @@ -121,7 +121,7 @@ from diffusers.modular_pipelines import AutoPipelineBlocks class AutoImageBlocks(AutoPipelineBlocks): # List of sub-block classes to choose from - block_classes = [block_inpaint_cls, block_i2i_cls, block_t2i_cls] + block_classes = [InpaintBlock, ImageToImageBlock, TextToImageBlock] # Names for each block in the same order block_names = ["inpaint", "img2img", "text2img"] # Trigger inputs that determine which block to run @@ -129,8 +129,8 @@ class AutoImageBlocks(AutoPipelineBlocks): # - "image" triggers img2img workflow (but only if mask is not provided) # - if none of above, runs the text2img workflow (default) block_trigger_inputs = ["mask", "image", None] - # Description is extremely important for AutoPipelineBlocks + @property def description(self): return ( "Pipeline generates images given different types of conditions!\n" @@ -141,7 +141,7 @@ class AutoImageBlocks(AutoPipelineBlocks): ) ``` -It is **very** important to include a `description` to avoid any confusion over how to run a block and what inputs are required. While [`~modular_pipelines.AutoPipelineBlocks`] are convenient, it's conditional logic may be difficult to figure out if it isn't properly explained. +It is **very** important to include a `description` to avoid any confusion over how to run a block and what inputs are required. While [`~modular_pipelines.AutoPipelineBlocks`] are convenient, its conditional logic may be difficult to figure out if it isn't properly explained. Create an instance of `AutoImageBlocks`. @@ -152,5 +152,74 @@ auto_blocks = AutoImageBlocks() For more complex compositions, such as nested [`~modular_pipelines.AutoPipelineBlocks`] blocks when they're used as sub-blocks in larger pipelines, use the [`~modular_pipelines.SequentialPipelineBlocks.get_execution_blocks`] method to extract the a block that is actually run based on your input. ```py -auto_blocks.get_execution_blocks("mask") +auto_blocks.get_execution_blocks(mask=True) +``` + +## ConditionalPipelineBlocks + +[`~modular_pipelines.AutoPipelineBlocks`] is a special case of [`~modular_pipelines.ConditionalPipelineBlocks`]. While [`~modular_pipelines.AutoPipelineBlocks`] selects blocks based on whether a trigger input is provided or not, [`~modular_pipelines.ConditionalPipelineBlocks`] is able to select a block based on custom selection logic provided in the `select_block` method. + +Here is the same example written using [`~modular_pipelines.ConditionalPipelineBlocks`] directly: + +```py +from diffusers.modular_pipelines import ConditionalPipelineBlocks + +class AutoImageBlocks(ConditionalPipelineBlocks): + block_classes = [InpaintBlock, ImageToImageBlock, TextToImageBlock] + block_names = ["inpaint", "img2img", "text2img"] + block_trigger_inputs = ["mask", "image"] + default_block_name = "text2img" + + @property + def description(self): + return ( + "Pipeline generates images given different types of conditions!\n" + + "This is an auto pipeline block that works for text2img, img2img and inpainting tasks.\n" + + " - inpaint workflow is run when `mask` is provided.\n" + + " - img2img workflow is run when `image` is provided (but only when `mask` is not provided).\n" + + " - text2img workflow is run when neither `image` nor `mask` is provided.\n" + ) + + def select_block(self, mask=None, image=None) -> str | None: + if mask is not None: + return "inpaint" + if image is not None: + return "img2img" + return None # falls back to default_block_name ("text2img") +``` + +The inputs listed in `block_trigger_inputs` are passed as keyword arguments to `select_block()`. When `select_block` returns `None`, it falls back to `default_block_name`. If `default_block_name` is also `None`, the entire conditional block is skipped — this is useful for optional processing steps that should only run when specific inputs are provided. + +## Workflows + +Pipelines that contain conditional blocks ([`~modular_pipelines.AutoPipelineBlocks`] or [`~modular_pipelines.ConditionalPipelineBlocks]`) can support multiple workflows — for example, our SDXL modular pipeline supports a dozen workflows all in one pipeline. But this also means it can be confusing for users to know what workflows are supported and how to run them. For pipeline builders, it's useful to be able to extract only the blocks relevant to a specific workflow. + +We recommend defining a `_workflow_map` to give each workflow a name and explicitly list the inputs it requires. + +```py +from diffusers.modular_pipelines import SequentialPipelineBlocks + +class MyPipelineBlocks(SequentialPipelineBlocks): + block_classes = [TextEncoderBlock, AutoImageBlocks, DecodeBlock] + block_names = ["text_encoder", "auto_image", "decode"] + + _workflow_map = { + "text2image": {"prompt": True}, + "image2image": {"image": True, "prompt": True}, + "inpaint": {"mask": True, "image": True, "prompt": True}, + } +``` + +All of our built-in modular pipelines come with pre-defined workflows. The `available_workflows` property lists all supported workflows: + +```py +pipeline_blocks = MyPipelineBlocks() +pipeline_blocks.available_workflows +# ['text2image', 'image2image', 'inpaint'] +``` + +Retrieve a specific workflow with `get_workflow` to inspect and debug a specific block that executes the workflow. + +```py +pipeline_blocks.get_workflow("inpaint") ``` \ No newline at end of file diff --git a/docs/source/en/modular_diffusers/components_manager.md b/docs/source/en/modular_diffusers/components_manager.md index af53411b9533..426739347f27 100644 --- a/docs/source/en/modular_diffusers/components_manager.md +++ b/docs/source/en/modular_diffusers/components_manager.md @@ -12,179 +12,85 @@ specific language governing permissions and limitations under the License. # ComponentsManager -The [`ComponentsManager`] is a model registry and management system for Modular Diffusers. It adds and tracks models, stores useful metadata (model size, device placement, adapters), prevents duplicate model instances, and supports offloading. +The [`ComponentsManager`] is a model registry and management system for Modular Diffusers. It adds and tracks models, stores useful metadata (model size, device placement, adapters), and supports offloading. This guide will show you how to use [`ComponentsManager`] to manage components and device memory. -## Add a component +## Connect to a pipeline -The [`ComponentsManager`] should be created alongside a [`ModularPipeline`] in either [`~ModularPipeline.from_pretrained`] or [`~ModularPipelineBlocks.init_pipeline`]. +Create a [`ComponentsManager`] and pass it to a [`ModularPipeline`] with either [`~ModularPipeline.from_pretrained`] or [`~ModularPipelineBlocks.init_pipeline`]. -> [!TIP] -> The `collection` parameter is optional but makes it easier to organize and manage components. ```py from diffusers import ModularPipeline, ComponentsManager +import torch -comp = ComponentsManager() -pipe = ModularPipeline.from_pretrained("YiYiXu/modular-demo-auto", components_manager=comp, collection="test1") +manager = ComponentsManager() +pipe = ModularPipeline.from_pretrained("Tongyi-MAI/Z-Image-Turbo", components_manager=manager) +pipe.load_components(torch_dtype=torch.bfloat16) ``` ```py -from diffusers import ComponentsManager -from diffusers.modular_pipelines import SequentialPipelineBlocks -from diffusers.modular_pipelines.stable_diffusion_xl import TEXT2IMAGE_BLOCKS - -t2i_blocks = SequentialPipelineBlocks.from_blocks_dict(TEXT2IMAGE_BLOCKS) - -modular_repo_id = "YiYiXu/modular-loader-t2i-0704" -components = ComponentsManager() -t2i_pipeline = t2i_blocks.init_pipeline(modular_repo_id, components_manager=components) +from diffusers import ModularPipelineBlocks, ComponentsManager +import torch +manager = ComponentsManager() +blocks = ModularPipelineBlocks.from_pretrained("diffusers/Florence2-image-Annotator", trust_remote_code=True) +pipe= blocks.init_pipeline(components_manager=manager) +pipe.load_components(torch_dtype=torch.bfloat16) ``` -Components are only loaded and registered when using [`~ModularPipeline.load_components`] or [`~ModularPipeline.load_components`]. The example below uses [`~ModularPipeline.load_components`] to create a second pipeline that reuses all the components from the first one, and assigns it to a different collection - -```py -pipe.load_components() -pipe2 = ModularPipeline.from_pretrained("YiYiXu/modular-demo-auto", components_manager=comp, collection="test2") -``` - -Use the [`~ModularPipeline.null_component_names`] property to identify any components that need to be loaded, retrieve them with [`~ComponentsManager.get_components_by_names`], and then call [`~ModularPipeline.update_components`] to add the missing components. - -```py -pipe2.null_component_names -['text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'image_encoder', 'unet', 'vae', 'scheduler', 'controlnet'] - -comp_dict = comp.get_components_by_names(names=pipe2.null_component_names) -pipe2.update_components(**comp_dict) -``` - -To add individual components, use the [`~ComponentsManager.add`] method. This registers a component with a unique id. - -```py -from diffusers import AutoModel - -text_encoder = AutoModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder") -component_id = comp.add("text_encoder", text_encoder) -comp -``` - -Use [`~ComponentsManager.remove`] to remove a component using their id. - -```py -comp.remove("text_encoder_139917733042864") -``` - -## Retrieve a component - -The [`ComponentsManager`] provides several methods to retrieve registered components. - -### get_one - -The [`~ComponentsManager.get_one`] method returns a single component and supports pattern matching for the `name` parameter. If multiple components match, [`~ComponentsManager.get_one`] returns an error. - -| Pattern | Example | Description | -|-------------|----------------------------------|-------------------------------------------| -| exact | `comp.get_one(name="unet")` | exact name match | -| wildcard | `comp.get_one(name="unet*")` | names starting with "unet" | -| exclusion | `comp.get_one(name="!unet")` | exclude components named "unet" | -| or | `comp.get_one(name="unet|vae")` | name is "unet" or "vae" | - -[`~ComponentsManager.get_one`] also filters components by the `collection` argument or `load_id` argument. - -```py -comp.get_one(name="unet", collection="sdxl") -``` - -### get_components_by_names - -The [`~ComponentsManager.get_components_by_names`] method accepts a list of names and returns a dictionary mapping names to components. This is especially useful with [`ModularPipeline`] since they provide lists of required component names and the returned dictionary can be passed directly to [`~ModularPipeline.update_components`]. - -```py -component_dict = comp.get_components_by_names(names=["text_encoder", "unet", "vae"]) -{"text_encoder": component1, "unet": component2, "vae": component3} -``` - -## Duplicate detection - -It is recommended to load model components with [`ComponentSpec`] to assign components with a unique id that encodes their loading parameters. This allows [`ComponentsManager`] to automatically detect and prevent duplicate model instances even when different objects represent the same underlying checkpoint. - -```py -from diffusers import ComponentSpec, ComponentsManager -from transformers import CLIPTextModel - -comp = ComponentsManager() - -# Create ComponentSpec for the first text encoder -spec = ComponentSpec(name="text_encoder", repo="stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder", type_hint=AutoModel) -# Create ComponentSpec for a duplicate text encoder (it is same checkpoint, from the same repo/subfolder) -spec_duplicated = ComponentSpec(name="text_encoder_duplicated", repo="stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder", type_hint=CLIPTextModel) - -# Load and add both components - the manager will detect they're the same model -comp.add("text_encoder", spec.load()) -comp.add("text_encoder_duplicated", spec_duplicated.load()) -``` - -This returns a warning with instructions for removing the duplicate. +Components loaded by the pipeline are automatically registered in the manager. You can inspect them right away. -```py -ComponentsManager: adding component 'text_encoder_duplicated_139917580682672', but it has duplicate load_id 'stabilityai/stable-diffusion-xl-base-1.0|text_encoder|null|null' with existing components: text_encoder_139918506246832. To remove a duplicate, call `components_manager.remove('')`. -'text_encoder_duplicated_139917580682672' -``` - -You could also add a component without using [`ComponentSpec`] and duplicate detection still works in most cases even if you're adding the same component under a different name. - -However, [`ComponentManager`] can't detect duplicates when you load the same component into different objects. In this case, you should load a model with [`ComponentSpec`]. - -```py -text_encoder_2 = AutoModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder") -comp.add("text_encoder", text_encoder_2) -'text_encoder_139917732983664' -``` +## Inspect components -## Collections +Print the [`ComponentsManager`] to see all registered components, including their class, device placement, dtype, memory size, and load ID. -Collections are labels assigned to components for better organization and management. Add a component to a collection with the `collection` argument in [`~ComponentsManager.add`]. - -Only one component per name is allowed in each collection. Adding a second component with the same name automatically removes the first component. +The output below corresponds to the `from_pretrained` example above. ```py -from diffusers import ComponentSpec, ComponentsManager - -comp = ComponentsManager() -# Create ComponentSpec for the first UNet -spec = ComponentSpec(name="unet", repo="stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet", type_hint=AutoModel) -# Create ComponentSpec for a different UNet -spec2 = ComponentSpec(name="unet", repo="RunDiffusion/Juggernaut-XL-v9", subfolder="unet", type_hint=AutoModel, variant="fp16") - -# Add both UNets to the same collection - the second one will replace the first -comp.add("unet", spec.load(), collection="sdxl") -comp.add("unet", spec2.load(), collection="sdxl") +Components: +============================================================================================================================= +Models: +----------------------------------------------------------------------------------------------------------------------------- +Name_ID | Class | Device: act(exec) | Dtype | Size (GB) | Load ID +----------------------------------------------------------------------------------------------------------------------------- +text_encoder_140458257514752 | Qwen3Model | cpu | torch.bfloat16 | 7.49 | Tongyi-MAI/Z-Image-Turbo|text_encoder|null|null +vae_140458257515376 | AutoencoderKL | cpu | torch.bfloat16 | 0.16 | Tongyi-MAI/Z-Image-Turbo|vae|null|null +transformer_140458257515616 | ZImageTransformer2DModel | cpu | torch.bfloat16 | 11.46 | Tongyi-MAI/Z-Image-Turbo|transformer|null|null +----------------------------------------------------------------------------------------------------------------------------- + +Other Components: +----------------------------------------------------------------------------------------------------------------------------- +ID | Class | Collection +----------------------------------------------------------------------------------------------------------------------------- +scheduler_140461023555264 | FlowMatchEulerDiscreteScheduler | N/A +tokenizer_140458256346432 | Qwen2Tokenizer | N/A +----------------------------------------------------------------------------------------------------------------------------- ``` -This makes it convenient to work with node-based systems because you can: - -- Mark all models as loaded from one node with the `collection` label. -- Automatically replace models when new checkpoints are loaded under the same name. -- Batch delete all models in a collection when a node is removed. +The table shows models (with device, dtype, and memory info) separately from other components like schedulers and tokenizers. If any models have LoRA adapters, IP-Adapters, or quantization applied, that information is displayed in an additional section at the bottom. ## Offloading The [`~ComponentsManager.enable_auto_cpu_offload`] method is a global offloading strategy that works across all models regardless of which pipeline is using them. Once enabled, you don't need to worry about device placement if you add or remove components. ```py -comp.enable_auto_cpu_offload(device="cuda") +manager.enable_auto_cpu_offload(device="cuda") ``` All models begin on the CPU and [`ComponentsManager`] moves them to the appropriate device right before they're needed, and moves other models back to the CPU when GPU memory is low. -You can set your own rules for which models to offload first. +Call [`~ComponentsManager.disable_auto_cpu_offload`] to disable offloading. + +```py +manager.disable_auto_cpu_offload() +``` diff --git a/docs/source/en/modular_diffusers/custom_blocks.md b/docs/source/en/modular_diffusers/custom_blocks.md index 1c311582264e..66e1de172b34 100644 --- a/docs/source/en/modular_diffusers/custom_blocks.md +++ b/docs/source/en/modular_diffusers/custom_blocks.md @@ -16,7 +16,7 @@ specific language governing permissions and limitations under the License. [ModularPipelineBlocks](./pipeline_block) are the fundamental building blocks of a [`ModularPipeline`]. You can create custom blocks by defining their inputs, outputs, and computation logic. This guide demonstrates how to create and use a custom block. > [!TIP] -> Explore the [Modular Diffusers Custom Blocks](https://huggingface.co/collections/diffusers/modular-diffusers-custom-blocks) collection for official custom modular blocks like Nano Banana. +> Explore the [Modular Diffusers Custom Blocks](https://huggingface.co/collections/diffusers/modular-diffusers-custom-blocks) collection for official custom blocks. ## Project Structure @@ -31,54 +31,58 @@ Your custom block project should use the following structure: - `block.py` contains the custom block implementation - `modular_config.json` contains the metadata needed to load the block -## Example: Florence 2 Inpainting Block +## Quick Start with Template -In this example we will create a custom block that uses the [Florence 2](https://huggingface.co/docs/transformers/model_doc/florence2) model to process an input image and generate a mask for inpainting. +The fastest way to create a custom block is to start from our template. The template provides a pre-configured project structure with `block.py` and `modular_config.json` files, plus commented examples showing how to define components, inputs, outputs, and the `__call__` method—so you can focus on your custom logic instead of boilerplate setup. -The first step is to define the components that the block will use. In this case, we will need to use the `Florence2ForConditionalGeneration` model and its corresponding processor `AutoProcessor`. When defining components, we must specify the name of the component within our pipeline, model class via `type_hint`, and provide a `pretrained_model_name_or_path` for the component if we intend to load the model weights from a specific repository on the Hub. +### Download the template -```py -# Inside block.py -from diffusers.modular_pipelines import ( - ModularPipelineBlocks, - ComponentSpec, +```python +from diffusers import ModularPipelineBlocks + +model_id = "diffusers/custom-block-template" +local_dir = model_id.split("/")[-1] + +blocks = ModularPipelineBlocks.from_pretrained( + model_id, + trust_remote_code=True, + local_dir=local_dir ) -from transformers import AutoProcessor, Florence2ForConditionalGeneration +``` +This saves the template files to `custom-block-template/` locally or you could use `local_dir` to save to a specific location. -class Florence2ImageAnnotatorBlock(ModularPipelineBlocks): +### Edit locally - @property - def expected_components(self): - return [ - ComponentSpec( - name="image_annotator", - type_hint=Florence2ForConditionalGeneration, - pretrained_model_name_or_path="florence-community/Florence-2-base-ft", - ), - ComponentSpec( - name="image_annotator_processor", - type_hint=AutoProcessor, - pretrained_model_name_or_path="florence-community/Florence-2-base-ft", - ), - ] +Open `block.py` and implement your custom block. The template includes commented examples showing how to define each property. See the [Florence-2 example](#example-florence-2-image-annotator) below for a complete implementation. + +### Test your block + +```python +from diffusers import ModularPipelineBlocks + +blocks = ModularPipelineBlocks.from_pretrained(local_dir, trust_remote_code=True) +pipeline = blocks.init_pipeline() +output = pipeline(...) # your inputs here ``` -Next, we define the inputs and outputs of the block. The inputs include the image to be annotated, the annotation task, and the annotation prompt. The outputs include the generated mask image and annotations. +### Upload to the Hub -```py -from typing import List, Union -from PIL import Image, ImageDraw -import torch -import numpy as np - -from diffusers.modular_pipelines import ( - PipelineState, - ModularPipelineBlocks, - InputParam, - ComponentSpec, - OutputParam, -) +```python +pipeline.save_pretrained(local_dir, repo_id="your-username/your-block-name", push_to_hub=True) +``` + +## Example: Florence-2 Image Annotator + +This example creates a custom block with [Florence-2](https://huggingface.co/docs/transformers/model_doc/florence2) to process an input image and generate a mask for inpainting. + +### Define components + +Define the components the block needs, `Florence2ForConditionalGeneration` and its processor. When defining components, specify the `name` (how you'll access it in code), `type_hint` (the model class), and `pretrained_model_name_or_path` (where to load weights from). + +```python +# Inside block.py +from diffusers.modular_pipelines import ModularPipelineBlocks, ComponentSpec from transformers import AutoProcessor, Florence2ForConditionalGeneration @@ -98,122 +102,21 @@ class Florence2ImageAnnotatorBlock(ModularPipelineBlocks): pretrained_model_name_or_path="florence-community/Florence-2-base-ft", ), ] - - @property - def inputs(self) -> List[InputParam]: - return [ - InputParam( - "image", - type_hint=Union[Image.Image, List[Image.Image]], - required=True, - description="Image(s) to annotate", - ), - InputParam( - "annotation_task", - type_hint=Union[str, List[str]], - required=True, - default="", - description="""Annotation Task to perform on the image. - Supported Tasks: - - - - - - - - - - - """, - ), - InputParam( - "annotation_prompt", - type_hint=Union[str, List[str]], - required=True, - description="""Annotation Prompt to provide more context to the task. - Can be used to detect or segment out specific elements in the image - """, - ), - InputParam( - "annotation_output_type", - type_hint=str, - required=True, - default="mask_image", - description="""Output type from annotation predictions. Availabe options are - mask_image: - -black and white mask image for the given image based on the task type - mask_overlay: - - mask overlayed on the original image - bounding_box: - - bounding boxes drawn on the original image - """, - ), - InputParam( - "annotation_overlay", - type_hint=bool, - required=True, - default=False, - description="", - ), - ] - - @property - def intermediate_outputs(self) -> List[OutputParam]: - return [ - OutputParam( - "mask_image", - type_hint=Image, - description="Inpainting Mask for input Image(s)", - ), - OutputParam( - "annotations", - type_hint=dict, - description="Annotations Predictions for input Image(s)", - ), - OutputParam( - "image", - type_hint=Image, - description="Annotated input Image(s)", - ), - ] - ``` -Now we implement the `__call__` method, which contains the logic for processing the input image and generating the mask. +### Define inputs and outputs -```py +Inputs include the image, annotation task, and prompt. Outputs include the generated mask and annotations. + +```python from typing import List, Union -from PIL import Image, ImageDraw -import torch -import numpy as np - -from diffusers.modular_pipelines import ( - PipelineState, - ModularPipelineBlocks, - InputParam, - ComponentSpec, - OutputParam, -) -from transformers import AutoProcessor, Florence2ForConditionalGeneration +from PIL import Image +from diffusers.modular_pipelines import InputParam, OutputParam class Florence2ImageAnnotatorBlock(ModularPipelineBlocks): - @property - def expected_components(self): - return [ - ComponentSpec( - name="image_annotator", - type_hint=Florence2ForConditionalGeneration, - pretrained_model_name_or_path="florence-community/Florence-2-base-ft", - ), - ComponentSpec( - name="image_annotator_processor", - type_hint=AutoProcessor, - pretrained_model_name_or_path="florence-community/Florence-2-base-ft", - ), - ] + # ... expected_components from above ... @property def inputs(self) -> List[InputParam]: @@ -226,51 +129,21 @@ class Florence2ImageAnnotatorBlock(ModularPipelineBlocks): ), InputParam( "annotation_task", - type_hint=Union[str, List[str]], - required=True, + type_hint=str, default="", - description="""Annotation Task to perform on the image. - Supported Tasks: - - - - - - - - - - - """, + description="Annotation task to perform (e.g., , , )", ), InputParam( "annotation_prompt", - type_hint=Union[str, List[str]], + type_hint=str, required=True, - description="""Annotation Prompt to provide more context to the task. - Can be used to detect or segment out specific elements in the image - """, + description="Prompt to provide context for the annotation task", ), InputParam( "annotation_output_type", type_hint=str, - required=True, default="mask_image", - description="""Output type from annotation predictions. Availabe options are - mask_image: - -black and white mask image for the given image based on the task type - mask_overlay: - - mask overlayed on the original image - bounding_box: - - bounding boxes drawn on the original image - """, - ), - InputParam( - "annotation_overlay", - type_hint=bool, - required=True, - default=False, - description="", + description="Output type: 'mask_image', 'mask_overlay', or 'bounding_box'", ), ] @@ -279,109 +152,45 @@ class Florence2ImageAnnotatorBlock(ModularPipelineBlocks): return [ OutputParam( "mask_image", - type_hint=Image, - description="Inpainting Mask for input Image(s)", + type_hint=Image.Image, + description="Inpainting mask for the input image", ), OutputParam( "annotations", type_hint=dict, - description="Annotations Predictions for input Image(s)", + description="Raw annotation predictions", ), OutputParam( "image", - type_hint=Image, - description="Annotated input Image(s)", + type_hint=Image.Image, + description="Annotated image", ), ] +``` - def get_annotations(self, components, images, prompts, task): - task_prompts = [task + prompt for prompt in prompts] +### Implement the `__call__` method - inputs = components.image_annotator_processor( - text=task_prompts, images=images, return_tensors="pt" - ).to(components.image_annotator.device, components.image_annotator.dtype) +The `__call__` method contains the block's logic. Access inputs via `block_state`, run your computation, and set outputs back to `block_state`. - generated_ids = components.image_annotator.generate( - input_ids=inputs["input_ids"], - pixel_values=inputs["pixel_values"], - max_new_tokens=1024, - early_stopping=False, - do_sample=False, - num_beams=3, - ) - annotations = components.image_annotator_processor.batch_decode( - generated_ids, skip_special_tokens=False - ) - outputs = [] - for image, annotation in zip(images, annotations): - outputs.append( - components.image_annotator_processor.post_process_generation( - annotation, task=task, image_size=(image.width, image.height) - ) - ) - return outputs - - def prepare_mask(self, images, annotations, overlay=False, fill="white"): - masks = [] - for image, annotation in zip(images, annotations): - mask_image = image.copy() if overlay else Image.new("L", image.size, 0) - draw = ImageDraw.Draw(mask_image) - - for _, _annotation in annotation.items(): - if "polygons" in _annotation: - for polygon in _annotation["polygons"]: - polygon = np.array(polygon).reshape(-1, 2) - if len(polygon) < 3: - continue - polygon = polygon.reshape(-1).tolist() - draw.polygon(polygon, fill=fill) - - elif "bbox" in _annotation: - bbox = _annotation["bbox"] - draw.rectangle(bbox, fill="white") - - masks.append(mask_image) - - return masks - - def prepare_bounding_boxes(self, images, annotations): - outputs = [] - for image, annotation in zip(images, annotations): - image_copy = image.copy() - draw = ImageDraw.Draw(image_copy) - for _, _annotation in annotation.items(): - bbox = _annotation["bbox"] - label = _annotation["label"] - - draw.rectangle(bbox, outline="red", width=3) - draw.text((bbox[0], bbox[1] - 20), label, fill="red") - - outputs.append(image_copy) - - return outputs - - def prepare_inputs(self, images, prompts): - prompts = prompts or "" - - if isinstance(images, Image.Image): - images = [images] - if isinstance(prompts, str): - prompts = [prompts] - - if len(images) != len(prompts): - raise ValueError("Number of images and annotation prompts must match.") - - return images, prompts +```python +import torch +from diffusers.modular_pipelines import PipelineState + + +class Florence2ImageAnnotatorBlock(ModularPipelineBlocks): + + # ... expected_components, inputs, intermediate_outputs from above ... @torch.no_grad() def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) + images, annotation_task_prompt = self.prepare_inputs( block_state.image, block_state.annotation_prompt ) task = block_state.annotation_task fill = block_state.fill - + annotations = self.get_annotations( components, images, annotation_task_prompt, task ) @@ -400,67 +209,69 @@ class Florence2ImageAnnotatorBlock(ModularPipelineBlocks): self.set_block_state(state, block_state) return components, state - -``` - -Once we have defined our custom block, we can save it to the Hub, using either the CLI or the [`push_to_hub`] method. This will make it easy to share and reuse our custom block with other pipelines. - - - - -```shell -# In the folder with the `block.py` file, run: -diffusers-cli custom_block -``` - -Then upload the block to the Hub: - -```shell -hf upload . . + + # Helper methods for mask/bounding box generation... ``` - - -```py -from block import Florence2ImageAnnotatorBlock -block = Florence2ImageAnnotatorBlock() -block.push_to_hub("") -``` - - - +> [!TIP] +> See the complete implementation at [diffusers/Florence2-image-Annotator](https://huggingface.co/diffusers/Florence2-image-Annotator). ## Using Custom Blocks -Load the custom block with [`~ModularPipelineBlocks.from_pretrained`] and set `trust_remote_code=True`. +Load a custom block with [`~ModularPipeline.from_pretrained`] and set `trust_remote_code=True`. ```py import torch -from diffusers.modular_pipelines import ModularPipelineBlocks, SequentialPipelineBlocks -from diffusers.modular_pipelines.stable_diffusion_xl import INPAINT_BLOCKS +from diffusers import ModularPipeline from diffusers.utils import load_image -# Fetch the Florence2 image annotator block that will create our mask -image_annotator_block = ModularPipelineBlocks.from_pretrained("diffusers/florence-2-custom-block", trust_remote_code=True) +# Load the Florence-2 annotator pipeline +image_annotator = ModularPipeline.from_pretrained( + "diffusers/Florence2-image-Annotator", + trust_remote_code=True +) -my_blocks = INPAINT_BLOCKS.copy() -# insert the annotation block before the image encoding step -my_blocks.insert("image_annotator", image_annotator_block, 1) +# Check the docstring to see inputs/outputs +print(image_annotator.blocks.doc) +``` -# Create our initial set of inpainting blocks -blocks = SequentialPipelineBlocks.from_blocks_dict(my_blocks) +Use the block to generate a mask: -repo_id = "diffusers/modular-stable-diffusion-xl-base-1.0" -pipe = blocks.init_pipeline(repo_id) -pipe.load_components(torch_dtype=torch.float16, device_map="cuda", trust_remote_code=True) +```python +image_annotator.load_components(torch_dtype=torch.bfloat16) +image_annotator.to("cuda") -image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true") +image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg") image = image.resize((1024, 1024)) - prompt = ["A red car"] annotation_task = "" annotation_prompt = ["the car"] +mask_image = image_annotator_node( + prompt=prompt, + image=image, + annotation_task=annotation_task, + annotation_prompt=annotation_prompt, + annotation_output_type="mask_image", +).images +mask_image[0].save("car-mask.png") +``` + +Compose it with other blocks to create a new pipeline: + +```python +# Get the annotator block +annotator_block = image_annotator.blocks + +# Get an inpainting workflow and insert the annotator at the beginning +inpaint_blocks = ModularPipeline.from_pretrained("Qwen/Qwen-Image").blocks.get_workflow("inpainting") +inpaint_blocks.sub_blocks.insert("image_annotator", annotator_block, 0) + +# Initialize the combined pipeline +pipe = inpaint_blocks.init_pipeline() +pipe.load_components(torch_dtype=torch.float16, device="cuda") + +# Now the pipeline automatically generates masks from prompts output = pipe( prompt=prompt, image=image, @@ -475,18 +286,95 @@ output = pipe( output[0].save("florence-inpainting.png") ``` -## Editing Custom Blocks +## Editing custom blocks + +Edit custom blocks by downloading it locally. This is the same workflow as the [Quick Start with Template](#quick-start-with-template), but starting from an existing block instead of the template. + +Use the `local_dir` argument to download a custom block to a specific folder: + +```python +from diffusers import ModularPipelineBlocks -By default, custom blocks are saved in your cache directory. Use the `local_dir` argument to download and edit a custom block in a specific folder. +# Download to a local folder for editing +annotator_block = ModularPipelineBlocks.from_pretrained( + "diffusers/Florence2-image-Annotator", + trust_remote_code=True, + local_dir="./my-florence-block" +) +``` + +Any changes made to the block files in this folder will be reflected when you load the block again. When you're ready to share your changes, upload to a new repository: + +```python +pipeline = annotator_block.init_pipeline() +pipeline.save_pretrained("./my-florence-block", repo_id="your-username/my-custom-florence", push_to_hub=True) +``` + +## Next Steps + + + + +This guide covered creating a single custom block. Learn how to compose multiple blocks together: + +- [SequentialPipelineBlocks](./sequential_pipeline_blocks): Chain blocks to execute in sequence +- [ConditionalPipelineBlocks](./auto_pipeline_blocks): Create conditional blocks that select different execution paths +- [LoopSequentialPipelineBlocks](./loop_sequential_pipeline_blocks): Define an iterative workflows like the denoising loop + + + + +Make your custom block work with Mellon's visual interface. See the [Mellon Custom Blocks](./mellon) guide. + + + + +Browse the [Modular Diffusers Custom Blocks](https://huggingface.co/collections/diffusers/modular-diffusers-custom-blocks) collection for inspiration and ready-to-use blocks. + + + + +## Dependencies + +Declaring package dependencies in custom blocks prevents runtime import errors later on. Diffusers validates the dependencies and returns a warning if a package is missing or incompatible. + +Set a `_requirements` attribute in your block class, mapping package names to version specifiers. ```py -import torch -from diffusers.modular_pipelines import ModularPipelineBlocks, SequentialPipelineBlocks -from diffusers.modular_pipelines.stable_diffusion_xl import INPAINT_BLOCKS -from diffusers.utils import load_image +from diffusers.modular_pipelines import PipelineBlock -# Fetch the Florence2 image annotator block that will create our mask -image_annotator_block = ModularPipelineBlocks.from_pretrained("diffusers/florence-2-custom-block", trust_remote_code=True, local_dir="/my-local-folder") +class MyCustomBlock(PipelineBlock): + _requirements = { + "transformers": ">=4.44.0", + "sentencepiece": ">=0.2.0" + } ``` -Any changes made to the block files in this folder will be reflected when you load the block again. +When there are blocks with different requirements, Diffusers merges their requirements. + +```py +from diffusers.modular_pipelines import SequentialPipelineBlocks + +class BlockA(PipelineBlock): + _requirements = {"transformers": ">=4.44.0"} + # ... + +class BlockB(PipelineBlock): + _requirements = {"sentencepiece": ">=0.2.0"} + # ... + +pipe = SequentialPipelineBlocks.from_blocks_dict({ + "block_a": BlockA, + "block_b": BlockB, +}) +``` + +When this block is saved with [`~ModularPipeline.save_pretrained`], the requirements are saved to the `modular_config.json` file. When this block is loaded, Diffusers checks each requirement against the current environment. If there is a mismatch or a package isn't found, Diffusers returns the following warning. + +```md +# missing package +xyz-package was specified in the requirements but wasn't found in the current environment. + +# version mismatch +xyz requirement 'specific-version' is not satisfied by the installed version 'actual-version'. Things might work unexpected. +``` diff --git a/docs/source/en/modular_diffusers/guiders.md b/docs/source/en/modular_diffusers/guiders.md index 6abe4fad2736..ffe039f41556 100644 --- a/docs/source/en/modular_diffusers/guiders.md +++ b/docs/source/en/modular_diffusers/guiders.md @@ -89,10 +89,8 @@ t2i_pipeline.guider ## Changing guider parameters -The guider parameters can be adjusted with either the [`~ComponentSpec.create`] method or with [`~ModularPipeline.update_components`]. The example below changes the `guidance_scale` value. +The guider parameters can be adjusted with the [`~ComponentSpec.create`] method and [`~ModularPipeline.update_components`]. The example below changes the `guidance_scale` value. - - ```py guider_spec = t2i_pipeline.get_component_spec("guider") @@ -100,18 +98,6 @@ guider = guider_spec.create(guidance_scale=10) t2i_pipeline.update_components(guider=guider) ``` - - - -```py -guider_spec = t2i_pipeline.get_component_spec("guider") -guider_spec.config["guidance_scale"] = 10 -t2i_pipeline.update_components(guider=guider_spec) -``` - - - - ## Uploading custom guiders Call the [`~utils.PushToHubMixin.push_to_hub`] method on a custom guider to share it to the Hub. diff --git a/docs/source/en/modular_diffusers/loop_sequential_pipeline_blocks.md b/docs/source/en/modular_diffusers/loop_sequential_pipeline_blocks.md index a80309de19a6..74a868922799 100644 --- a/docs/source/en/modular_diffusers/loop_sequential_pipeline_blocks.md +++ b/docs/source/en/modular_diffusers/loop_sequential_pipeline_blocks.md @@ -53,7 +53,7 @@ The loop wrapper can pass additional arguments, like current iteration index, to A loop block is a [`~modular_pipelines.ModularPipelineBlocks`], but the `__call__` method behaves differently. -- It recieves the iteration variable from the loop wrapper. +- It receives the iteration variable from the loop wrapper. - It works directly with the [`~modular_pipelines.BlockState`] instead of the [`~modular_pipelines.PipelineState`]. - It doesn't require retrieving or updating the [`~modular_pipelines.BlockState`]. diff --git a/docs/source/en/modular_diffusers/mellon.md b/docs/source/en/modular_diffusers/mellon.md new file mode 100644 index 000000000000..808e62ad7966 --- /dev/null +++ b/docs/source/en/modular_diffusers/mellon.md @@ -0,0 +1,270 @@ + + + +## Using Custom Blocks with Mellon + +[Mellon](https://github.com/cubiq/Mellon) is a visual workflow interface that integrates with Modular Diffusers and is designed for node-based workflows. + +> [!WARNING] +> Mellon is in early development and not ready for production use yet. Consider this a sneak peek of how the integration works! + + +Custom blocks work in Mellon out of the box - just need to add a `mellon_pipeline_config.json` to your repository. This config file tells Mellon how to render your block's parameters as UI components. + +Here's what it looks like in action with the [Gemini Prompt Expander](https://huggingface.co/diffusers/gemini-prompt-expander-mellon) block: + +![Mellon custom block demo](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/modular_demo_dynamic.gif) + +To use a modular diffusers custom block in Mellon: +1. Drag a **Dynamic Block Node** from the ModularDiffusers section +2. Enter the `repo_id` (e.g., `diffusers/gemini-prompt-expander-mellon`) +3. Click **Load Custom Block** +4. The node transforms to show your block's inputs and outputs + +Now let's walk through how to create this config for your own custom block. + +## Steps to create a Mellon config + +1. **Specify Mellon types for your parameters** - Each `InputParam`/`OutputParam` needs a type that tells Mellon what UI component to render (e.g., `"textbox"`, `"dropdown"`, `"image"`). +2. **Generate `mellon_pipeline_config.json`** - Use our utility to generate a config template and push it to your Hub repository. +3. **(Optional) Manually adjust the config** - Fine-tune the generated config for your specific needs. + +## Specify Mellon types for parameters + +Mellon types determine how each parameter renders in the UI. If you don't specify a type for a parameter, it will default to `"custom"`, which renders as a simple connection dot. You can always adjust this later in the generated config. + + +| Type | Input/Output | Description | +|------|--------------|-------------| +| `image` | Both | Image (PIL Image) | +| `video` | Both | Video | +| `text` | Both | Text display | +| `textbox` | Input | Text input | +| `dropdown` | Input | Dropdown selection menu | +| `slider` | Input | Slider for numeric values | +| `number` | Input | Numeric input | +| `checkbox` | Input | Boolean toggle | + +For parameters that need more configuration (like dropdowns with options, or sliders with min/max values), pass a `MellonParam` instance directly instead of a string. You can use one of the class methods below, or create a fully custom one with `MellonParam(name, label, type, ...)`. + +| Method | Description | +|--------|-------------| +| `MellonParam.Input.image(name)` | Image input | +| `MellonParam.Input.textbox(name, default)` | Text input as textarea | +| `MellonParam.Input.dropdown(name, options, default)` | Dropdown selection | +| `MellonParam.Input.slider(name, default, min, max, step)` | Slider for numeric values | +| `MellonParam.Input.number(name, default, min, max, step)` | Numeric input (no slider) | +| `MellonParam.Input.seed(name, default)` | Seed input with randomize button | +| `MellonParam.Input.checkbox(name, default)` | Boolean checkbox | +| `MellonParam.Input.model(name)` | Model input for diffusers components | +| `MellonParam.Output.image(name)` | Image output | +| `MellonParam.Output.video(name)` | Video output | +| `MellonParam.Output.text(name)` | Text output | +| `MellonParam.Output.model(name)` | Model output for diffusers components | + +Choose one of the methods below to specify a Mellon type. + +### Using `metadata` in block definitions + +If you're defining a custom block from scratch, add `metadata={"mellon": ""}` directly to your `InputParam` and `OutputParam` definitions. If you're editing an existing custom block from the Hub, see [Editing custom blocks](./custom_blocks#editing-custom-blocks) for how to download it locally. + +```python +class GeminiPromptExpander(ModularPipelineBlocks): + + @property + def inputs(self) -> List[InputParam]: + return [ + InputParam( + "prompt", + type_hint=str, + required=True, + description="Prompt to use", + metadata={"mellon": "textbox"}, # Text input + ) + ] + + @property + def intermediate_outputs(self) -> List[OutputParam]: + return [ + OutputParam( + "prompt", + type_hint=str, + description="Expanded prompt by the LLM", + metadata={"mellon": "text"}, # Text output + ), + OutputParam( + "old_prompt", + type_hint=str, + description="Old prompt provided by the user", + # No metadata - we don't want to render this in UI + ) + ] +``` + +For full control over UI configuration, pass a `MellonParam` instance directly: +```python +from diffusers.modular_pipelines.mellon_node_utils import MellonParam + +InputParam( + "mode", + type_hint=str, + default="balanced", + metadata={"mellon": MellonParam.Input.dropdown("mode", options=["fast", "balanced", "quality"])}, +) +``` + +### Using `input_types` and `output_types` when Generating Config + +If you're working with an existing pipeline or prefer to keep your block definitions clean, specify types when generating the config using the `input_types/output_types` argument: +```python +from diffusers.modular_pipelines.mellon_node_utils import MellonPipelineConfig + +mellon_config = MellonPipelineConfig.from_custom_block( + blocks, + input_types={"prompt": "textbox"}, + output_types={"prompt": "text"} +) +``` + +> [!NOTE] +> When both `metadata` and `input_types`/`output_types` are specified, the arguments overrides `metadata`. + +## Generate and push the Mellon config + +After adding metadata to your block, generate the default Mellon configuration template and push it to the Hub: + +```python +from diffusers import ModularPipelineBlocks +from diffusers.modular_pipelines.mellon_node_utils import MellonPipelineConfig + +# load your custom blocks from your local dir +blocks = ModularPipelineBlocks.from_pretrained("/path/local/folder", trust_remote_code=True) + +# Generate the default config template +mellon_config = MellonPipelineConfig.from_custom_block(blocks) +# push the default template to `repo_id`, you will need to pass the same local folder path so that it will save the config locally first +mellon_config.save( + local_dir="/path/local/folder", + repo_id= repo_id, + push_to_hub=True +) +``` + +This creates a `mellon_pipeline_config.json` file in your repository. + +## Review and adjust the config + +The generated template is a starting point - you may want to adjust it for your needs. Let's walk through the generated config for the Gemini Prompt Expander: + +```json +{ + "label": "Gemini Prompt Expander", + "default_repo": "", + "default_dtype": "", + "node_params": { + "custom": { + "params": { + "prompt": { + "label": "Prompt", + "type": "string", + "display": "textarea", + "default": "" + }, + "out_prompt": { + "label": "Prompt", + "type": "string", + "display": "output" + }, + "old_prompt": { + "label": "Old Prompt", + "type": "custom", + "display": "output" + }, + "doc": { + "label": "Doc", + "type": "string", + "display": "output" + } + }, + "input_names": ["prompt"], + "model_input_names": [], + "output_names": ["out_prompt", "old_prompt", "doc"], + "block_name": "custom", + "node_type": "custom" + } + } +} +``` + +### Understanding the Structure + +The `params` dict defines how each UI element renders. The `input_names`, `model_input_names`, and `output_names` lists map these UI elements to the underlying [`ModularPipelineBlocks`]'s I/O interface: + +| Mellon Config | ModularPipelineBlocks | +|---------------|----------------------| +| `input_names` | `inputs` property | +| `model_input_names` | `expected_components` property | +| `output_names` | `intermediate_outputs` property | + +In this example: `prompt` is the only input. There are no model components, and outputs include `out_prompt`, `old_prompt`, and `doc`. + +Now let's look at the `params` dict: + +- **`prompt`**: An input parameter with `display: "textarea"` (renders as a text input box), `label: "Prompt"` (shown in the UI), and `default: ""` (starts empty). The `type: "string"` field is important in Mellon because it determines which nodes can connect together - only matching types can be linked with "noodles". + +- **`out_prompt`**: The expanded prompt output. The `out_` prefix was automatically added because the input and output share the same name (`prompt`), avoiding naming conflicts in the config. It has `display: "output"` which renders as an output socket. + +- **`old_prompt`**: Has `type: "custom"` because we didn't specify metadata. This renders as a simple dot in the UI. Since we don't actually want to expose this in the UI, we can remove it. + +- **`doc`**: The documentation output, automatically added to all custom blocks. + +### Making Adjustments + +Remove `old_prompt` from both `params` and `output_names` because you won't need to use it. + +```json +{ + "label": "Gemini Prompt Expander", + "default_repo": "", + "default_dtype": "", + "node_params": { + "custom": { + "params": { + "prompt": { + "label": "Prompt", + "type": "string", + "display": "textarea", + "default": "" + }, + "out_prompt": { + "label": "Prompt", + "type": "string", + "display": "output" + }, + "doc": { + "label": "Doc", + "type": "string", + "display": "output" + } + }, + "input_names": ["prompt"], + "model_input_names": [], + "output_names": ["out_prompt", "doc"], + "block_name": "custom", + "node_type": "custom" + } + } +} +``` + +See the final config at [diffusers/gemini-prompt-expander-mellon](https://huggingface.co/diffusers/gemini-prompt-expander-mellon). \ No newline at end of file diff --git a/docs/source/en/modular_diffusers/modular_diffusers_states.md b/docs/source/en/modular_diffusers/modular_diffusers_states.md index eb55b524e491..657b088fe485 100644 --- a/docs/source/en/modular_diffusers/modular_diffusers_states.md +++ b/docs/source/en/modular_diffusers/modular_diffusers_states.md @@ -25,9 +25,7 @@ This guide explains how states work and how they connect blocks. The [`~modular_pipelines.PipelineState`] is a global state container for all blocks. It maintains the complete runtime state of the pipeline and provides a structured way for blocks to read from and write to shared data. -There are two dict's in [`~modular_pipelines.PipelineState`] for structuring data. - -- The `values` dict is a **mutable** state containing a copy of user provided input values and intermediate output values generated by blocks. If a block modifies an `input`, it will be reflected in the `values` dict after calling `set_block_state`. +[`~modular_pipelines.PipelineState`] stores all data in a `values` dict, which is a **mutable** state containing user provided input values and intermediate output values generated by blocks. If a block modifies an `input`, it will be reflected in the `values` dict after calling `set_block_state`. ```py PipelineState( diff --git a/docs/source/en/modular_diffusers/modular_pipeline.md b/docs/source/en/modular_diffusers/modular_pipeline.md index 34cd8f72b5b7..e28e13ed5655 100644 --- a/docs/source/en/modular_diffusers/modular_pipeline.md +++ b/docs/source/en/modular_diffusers/modular_pipeline.md @@ -12,27 +12,28 @@ specific language governing permissions and limitations under the License. # ModularPipeline -[`ModularPipeline`] converts [`~modular_pipelines.ModularPipelineBlocks`]'s into an executable pipeline that loads models and performs the computation steps defined in the block. It is the main interface for running a pipeline and it is very similar to the [`DiffusionPipeline`] API. +[`ModularPipeline`] converts [`~modular_pipelines.ModularPipelineBlocks`] into an executable pipeline that loads models and performs the computation steps defined in the blocks. It is the main interface for running a pipeline and the API is very similar to [`DiffusionPipeline`] but with a few key differences. -The main difference is to include an expected `output` argument in the pipeline. +- **Loading is lazy.** With [`DiffusionPipeline`], [`~DiffusionPipeline.from_pretrained`] creates the pipeline and loads all models at the same time. With [`ModularPipeline`], creating and loading are two separate steps: [`~ModularPipeline.from_pretrained`] reads the configuration and knows where to load each component from, but doesn't actually load the model weights. You load the models later with [`~ModularPipeline.load_components`], which is where you pass loading arguments like `torch_dtype` and `quantization_config`. + +- **Two ways to create a pipeline.** You can use [`~ModularPipeline.from_pretrained`] with an existing diffusers model repository — it automatically maps to the default pipeline blocks and then converts to a [`ModularPipeline`] with no extra setup. You can check the [modular_pipelines_directory](https://github.com/huggingface/diffusers/tree/main/src/diffusers/modular_pipelines) to see which models are currently supported. You can also assemble your own pipeline from [`ModularPipelineBlocks`] and convert it with the [`~ModularPipelineBlocks.init_pipeline`] method (see [Creating a pipeline](#creating-a-pipeline) for more details). + +- **Running the pipeline is the same.** Once loaded, you call the pipeline with the same arguments you're used to. A single [`ModularPipeline`] can support multiple workflows (text-to-image, image-to-image, inpainting, etc.) when the pipeline blocks use [`AutoPipelineBlocks`](./auto_pipeline_blocks) to automatically select the workflow based on your inputs. + +Below are complete examples for text-to-image, image-to-image, and inpainting with SDXL. ```py import torch -from diffusers.modular_pipelines import SequentialPipelineBlocks -from diffusers.modular_pipelines.stable_diffusion_xl import TEXT2IMAGE_BLOCKS - -blocks = SequentialPipelineBlocks.from_blocks_dict(TEXT2IMAGE_BLOCKS) - -modular_repo_id = "YiYiXu/modular-loader-t2i-0704" -pipeline = blocks.init_pipeline(modular_repo_id) +from diffusers import ModularPipeline +pipeline = ModularPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0") pipeline.load_components(torch_dtype=torch.float16) pipeline.to("cuda") -image = pipeline(prompt="Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", output="images")[0] +image = pipeline(prompt="Astronaut in a jungle, cold color palette, muted colors, detailed, 8k").images[0] image.save("modular_t2i_out.png") ``` @@ -41,21 +42,17 @@ image.save("modular_t2i_out.png") ```py import torch -from diffusers.modular_pipelines import SequentialPipelineBlocks -from diffusers.modular_pipelines.stable_diffusion_xl import IMAGE2IMAGE_BLOCKS - -blocks = SequentialPipelineBlocks.from_blocks_dict(IMAGE2IMAGE_BLOCKS) - -modular_repo_id = "YiYiXu/modular-loader-t2i-0704" -pipeline = blocks.init_pipeline(modular_repo_id) +from diffusers import ModularPipeline +from diffusers.utils import load_image +pipeline = ModularPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0") pipeline.load_components(torch_dtype=torch.float16) pipeline.to("cuda") url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/sdxl-text2img.png" init_image = load_image(url) prompt = "a dog catching a frisbee in the jungle" -image = pipeline(prompt=prompt, image=init_image, strength=0.8, output="images")[0] +image = pipeline(prompt=prompt, image=init_image, strength=0.8).images[0] image.save("modular_i2i_out.png") ``` @@ -64,15 +61,10 @@ image.save("modular_i2i_out.png") ```py import torch -from diffusers.modular_pipelines import SequentialPipelineBlocks -from diffusers.modular_pipelines.stable_diffusion_xl import INPAINT_BLOCKS +from diffusers import ModularPipeline from diffusers.utils import load_image -blocks = SequentialPipelineBlocks.from_blocks_dict(INPAINT_BLOCKS) - -modular_repo_id = "YiYiXu/modular-loader-t2i-0704" -pipeline = blocks.init_pipeline(modular_repo_id) - +pipeline = ModularPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0") pipeline.load_components(torch_dtype=torch.float16) pipeline.to("cuda") @@ -83,276 +75,353 @@ init_image = load_image(img_url) mask_image = load_image(mask_url) prompt = "A deep sea diver floating" -image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, strength=0.85, output="images")[0] -image.save("moduar_inpaint_out.png") +image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, strength=0.85).images[0] +image.save("modular_inpaint_out.png") ``` -This guide will show you how to create a [`ModularPipeline`] and manage the components in it. +This guide will show you how to create a [`ModularPipeline`], manage its components, and run the pipeline. -## Adding blocks +## Creating a pipeline -Blocks are [`InsertableDict`] objects that can be inserted at specific positions, providing a flexible way to mix-and-match blocks. +There are two ways to create a [`ModularPipeline`]. Assemble and create a pipeline from [`ModularPipelineBlocks`] with [`~ModularPipelineBlocks.init_pipeline`], or load an existing pipeline with [`~ModularPipeline.from_pretrained`]. -Use [`~modular_pipelines.modular_pipeline_utils.InsertableDict.insert`] on either the block class or `sub_blocks` attribute to add a block. +You can also initialize a [`ComponentsManager`](./components_manager) to handle device placement and memory management. If you don't need automatic offloading, you can skip this and move the pipeline to your device manually with `pipeline.to("cuda")`. + +> [!TIP] +> Refer to the [ComponentsManager](./components_manager) doc for more details about how it can help manage components across different workflows. + +### init_pipeline + +[`~ModularPipelineBlocks.init_pipeline`] converts any [`ModularPipelineBlocks`] into a [`ModularPipeline`]. + +Let's define a minimal block to see how it works: ```py -# BLOCKS is dict of block classes, you need to add class to it -BLOCKS.insert("block_name", BlockClass, index) -# sub_blocks attribute contains instance, add a block instance to the attribute -t2i_blocks.sub_blocks.insert("block_name", block_instance, index) +from transformers import CLIPTextModel +from diffusers.modular_pipelines import ( + ComponentSpec, + ModularPipelineBlocks, + PipelineState, +) + +class MyBlock(ModularPipelineBlocks): + @property + def expected_components(self): + return [ + ComponentSpec( + name="text_encoder", + type_hint=CLIPTextModel, + pretrained_model_name_or_path="openai/clip-vit-large-patch14", + ), + ] + + def __call__(self, components, state: PipelineState) -> PipelineState: + return components, state ``` -Use [`~modular_pipelines.modular_pipeline_utils.InsertableDict.pop`] on either the block class or `sub_blocks` attribute to remove a block. +Call [`~ModularPipelineBlocks.init_pipeline`] to convert it into a pipeline. The `blocks` attribute on the pipeline is the blocks it was created from — it determines the expected inputs, outputs, and computation logic. ```py -# remove a block class from preset -BLOCKS.pop("text_encoder") -# split out a block instance on its own -text_encoder_block = t2i_blocks.sub_blocks.pop("text_encoder") +block = MyBlock() +pipe = block.init_pipeline() +pipe.blocks ``` -Swap blocks by setting the existing block to the new block. +``` +MyBlock { + "_class_name": "MyBlock", + "_diffusers_version": "0.37.0.dev0" +} +``` + +> [!WARNING] +> Blocks are mutable — you can freely add, remove, or swap blocks before creating a pipeline. However, once a pipeline is created, modifying `pipeline.blocks` won't affect the pipeline because it returns a copy. If you want a different block structure, create a new pipeline after modifying the blocks. + +When you call [`~ModularPipelineBlocks.init_pipeline`] without a repository, it uses the `pretrained_model_name_or_path` defined in the block's [`ComponentSpec`] to determine where to load each component from. Printing the pipeline shows the component loading configuration. ```py -# Replace block class in preset -BLOCKS["prepare_latents"] = CustomPrepareLatents -# Replace in sub_blocks attribute using an block instance -t2i_blocks.sub_blocks["prepare_latents"] = CustomPrepareLatents() +pipe +ModularPipeline { + "_blocks_class_name": "MyBlock", + "_class_name": "ModularPipeline", + "_diffusers_version": "0.37.0.dev0", + "text_encoder": [ + null, + null, + { + "pretrained_model_name_or_path": "openai/clip-vit-large-patch14", + "revision": null, + "subfolder": "", + "type_hint": [ + "transformers", + "CLIPTextModel" + ], + "variant": null + } + ] +} ``` -## Creating a pipeline +If you pass a repository to [`~ModularPipelineBlocks.init_pipeline`], it overrides the loading path by matching your block's components against the pipeline config in that repository (`model_index.json` or `modular_model_index.json`). -There are two ways to create a [`ModularPipeline`]. Assemble and create a pipeline from [`ModularPipelineBlocks`] or load an existing pipeline with [`~ModularPipeline.from_pretrained`]. +In the example below, the `pretrained_model_name_or_path` will be updated to `"stabilityai/stable-diffusion-xl-base-1.0"`. -You should also initialize a [`ComponentsManager`] to handle device placement and memory and component management. +```py +pipe = block.init_pipeline("stabilityai/stable-diffusion-xl-base-1.0") +pipe +ModularPipeline { + "_blocks_class_name": "MyBlock", + "_class_name": "ModularPipeline", + "_diffusers_version": "0.37.0.dev0", + "text_encoder": [ + null, + null, + { + "pretrained_model_name_or_path": "stabilityai/stable-diffusion-xl-base-1.0", + "revision": null, + "subfolder": "text_encoder", + "type_hint": [ + "transformers", + "CLIPTextModel" + ], + "variant": null + } + ] +} +``` -> [!TIP] -> Refer to the [ComponentsManager](./components_manager) doc for more details about how it can help manage components across different workflows. +If a component in your block doesn't exist in the repository, it remains `null` and is skipped during [`~ModularPipeline.load_components`]. - - +### from_pretrained -Use the [`~ModularPipelineBlocks.init_pipeline`] method to create a [`ModularPipeline`] from the component and configuration specifications. This method loads the *specifications* from a `modular_model_index.json` file, but it doesn't load the *models* yet. +[`~ModularPipeline.from_pretrained`] is a convenient way to create a [`ModularPipeline`] without defining blocks yourself. -```py -from diffusers import ComponentsManager -from diffusers.modular_pipelines import SequentialPipelineBlocks -from diffusers.modular_pipelines.stable_diffusion_xl import TEXT2IMAGE_BLOCKS +It works with three types of repositories. + +**A regular diffusers repository.** Pass any supported model repository and it automatically maps to the default pipeline blocks. Currently supported models include SDXL, Wan, Qwen, Z-Image, Flux, and Flux2. -t2i_blocks = SequentialPipelineBlocks.from_blocks_dict(TEXT2IMAGE_BLOCKS) +```py +from diffusers import ModularPipeline, ComponentsManager -modular_repo_id = "YiYiXu/modular-loader-t2i-0704" components = ComponentsManager() -t2i_pipeline = t2i_blocks.init_pipeline(modular_repo_id, components_manager=components) +pipeline = ModularPipeline.from_pretrained( + "stabilityai/stable-diffusion-xl-base-1.0", components_manager=components +) ``` - - - -The [`~ModularPipeline.from_pretrained`] method creates a [`ModularPipeline`] from a modular repository on the Hub. +**A modular repository.** These repositories contain a `modular_model_index.json` that specifies where to load each component from — the components can come from different repositories and the modular repository itself may not contain any model weights. For example, [diffusers/flux2-bnb-4bit-modular](https://huggingface.co/diffusers/flux2-bnb-4bit-modular) loads a quantized transformer from one repository and the remaining components from another. See [Modular repository](#modular-repository) for more details on the format. ```py from diffusers import ModularPipeline, ComponentsManager components = ComponentsManager() -pipeline = ModularPipeline.from_pretrained("YiYiXu/modular-loader-t2i-0704", components_manager=components) +pipeline = ModularPipeline.from_pretrained( + "diffusers/flux2-bnb-4bit-modular", components_manager=components +) ``` -Add the `trust_remote_code` argument to load a custom [`ModularPipeline`]. +**A modular repository with custom code.** Some repositories include custom pipeline blocks alongside the loading configuration. Add `trust_remote_code=True` to load them. See [Custom blocks](./custom_blocks) for how to create your own. ```py from diffusers import ModularPipeline, ComponentsManager components = ComponentsManager() -modular_repo_id = "YiYiXu/modular-diffdiff-0704" -diffdiff_pipeline = ModularPipeline.from_pretrained(modular_repo_id, trust_remote_code=True, components_manager=components) +pipeline = ModularPipeline.from_pretrained( + "diffusers/Florence2-image-Annotator", trust_remote_code=True, components_manager=components +) ``` - - - ## Loading components -A [`ModularPipeline`] doesn't automatically instantiate with components. It only loads the configuration and component specifications. You can load all components with [`~ModularPipeline.load_components`] or only load specific components with [`~ModularPipeline.load_components`]. +A [`ModularPipeline`] doesn't automatically instantiate with components. It only loads the configuration and component specifications. You can load components with [`~ModularPipeline.load_components`]. - - +This will load all the components that have a valid loading spec. ```py import torch -t2i_pipeline.load_components(torch_dtype=torch.float16) -t2i_pipeline.to("cuda") +pipeline.load_components(torch_dtype=torch.float16) ``` - - - -The example below only loads the UNet and VAE. +You can also load specific components by name. The example below only loads the `text_encoder`. ```py -import torch - -t2i_pipeline.load_components(names=["unet", "vae"], torch_dtype=torch.float16) +pipeline.load_components(names=["text_encoder"], torch_dtype=torch.float16) ``` - - - -Print the pipeline to inspect the loaded pretrained components. +After loading, printing the pipeline shows which components are loaded — the first two fields change from `null` to the component's library and class. ```py -t2i_pipeline +pipeline ``` -This should match the `modular_model_index.json` file from the modular repository a pipeline is initialized from. If a pipeline doesn't need a component, it won't be included even if it exists in the modular repository. - -To modify where components are loaded from, edit the `modular_model_index.json` file in the repository and change it to your desired loading path. The example below loads a UNet from a different repository. - -```json -# original -"unet": [ - null, null, - { - "repo": "stabilityai/stable-diffusion-xl-base-1.0", - "subfolder": "unet", - "variant": "fp16" - } +``` +# text_encoder is loaded - shows library and class +"text_encoder": [ + "transformers", + "CLIPTextModel", + { ... } ] -# modified +# unet is not loaded yet - still null "unet": [ - null, null, - { - "repo": "RunDiffusion/Juggernaut-XL-v9", - "subfolder": "unet", - "variant": "fp16" - } + null, + null, + { ... } ] ``` -### Component loading status - -The pipeline properties below provide more information about which components are loaded. - -Use `component_names` to return all expected components. +Loading keyword arguments like `torch_dtype`, `variant`, `revision`, and `quantization_config` are passed through to `from_pretrained()` for each component. You can pass a single value to apply to all components, or a dict to set per-component values. ```py -t2i_pipeline.component_names -['text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'guider', 'scheduler', 'unet', 'vae', 'image_processor'] +# apply bfloat16 to all components +pipeline.load_components(torch_dtype=torch.bfloat16) + +# different dtypes per component +pipeline.load_components(torch_dtype={"transformer": torch.bfloat16, "default": torch.float32}) ``` -Use `null_component_names` to return components that aren't loaded yet. Load these components with [`~ModularPipeline.from_pretrained`]. +[`~ModularPipeline.load_components`] only loads components that haven't been loaded yet and have a valid loading spec. This means if you've already set a component on the pipeline, calling [`~ModularPipeline.load_components`] again won't reload it. -```py -t2i_pipeline.null_component_names -['text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'scheduler'] -``` +## Updating components -Use `pretrained_component_names` to return components that will be loaded from pretrained models. +[`~ModularPipeline.update_components`] replaces a component on the pipeline with a new one. When a component is updated, the loading specifications are also updated in the pipeline config and [`~ModularPipeline.load_components`] will skip it on subsequent calls. -```py -t2i_pipeline.pretrained_component_names -['text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'scheduler', 'unet', 'vae'] -``` +### From AutoModel -Use `config_component_names` to return components that are created with the default config (not loaded from a modular repository). Components from a config aren't included because they are already initialized during pipeline creation. This is why they aren't listed in `null_component_names`. +You can pass a model object loaded with `AutoModel.from_pretrained()`. Models loaded this way are automatically tagged with their loading information. ```py -t2i_pipeline.config_component_names -['guider', 'image_processor'] +from diffusers import AutoModel + +unet = AutoModel.from_pretrained( + "RunDiffusion/Juggernaut-XL-v9", subfolder="unet", variant="fp16", torch_dtype=torch.float16 +) +pipeline.update_components(unet=unet) ``` -## Updating components +### From ComponentSpec -Components may be updated depending on whether it is a *pretrained component* or a *config component*. +Use [`~ModularPipeline.get_component_spec`] to get a copy of the current component specification, modify it, and load a new component. -> [!WARNING] -> A component may change from pretrained to config when updating a component. The component type is initially defined in a block's `expected_components` field. +```py +unet_spec = pipeline.get_component_spec("unet") + +# modify to load from a different repository +unet_spec.pretrained_model_name_or_path = "RunDiffusion/Juggernaut-XL-v9" -A pretrained component is updated with [`ComponentSpec`] whereas a config component is updated by eihter passing the object directly or with [`ComponentSpec`]. +# load and update +unet = unet_spec.load(torch_dtype=torch.float16) +pipeline.update_components(unet=unet) +``` -The [`ComponentSpec`] shows `default_creation_method="from_pretrained"` for a pretrained component shows `default_creation_method="from_config` for a config component. +You can also create a [`ComponentSpec`] from scratch. -To update a pretrained component, create a [`ComponentSpec`] with the name of the component and where to load it from. Use the [`~ComponentSpec.load`] method to load the component. +Not all components are loaded from pretrained weights — some are created from a config (listed under `pipeline.config_component_names`). For these, use [`~ComponentSpec.create`] instead of [`~ComponentSpec.load`]. ```py -from diffusers import ComponentSpec, UNet2DConditionModel - -unet_spec = ComponentSpec(name="unet",type_hint=UNet2DConditionModel, repo="stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet", variant="fp16") -unet = unet_spec.load(torch_dtype=torch.float16) +guider_spec = pipeline.get_component_spec("guider") +guider_spec.config = {"guidance_scale": 5.0} +guider = guider_spec.create() +pipeline.update_components(guider=guider) ``` -The [`~ModularPipeline.update_components`] method replaces the component with a new one. +Or simply pass the object directly. ```py -t2i_pipeline.update_components(unet=unet2) +from diffusers.guiders import ClassifierFreeGuidance + +guider = ClassifierFreeGuidance(guidance_scale=5.0) +pipeline.update_components(guider=guider) ``` -When a component is updated, the loading specifications are also updated in the pipeline config. +See the [Guiders](./guiders) guide for more details on available guiders and how to configure them. -### Component extraction and modification +## Splitting a pipeline into stages -When you use [`~ComponentSpec.load`], the new component maintains its loading specifications. This makes it possible to extract the specification and recreate the component. +Since blocks are composable, you can take a pipeline apart and reconstruct it into separate pipelines for each stage. The example below shows how we can separate the text encoder block from the rest of the pipeline, so you can encode the prompt independently and pass the embeddings to the main pipeline. ```py -spec = ComponentSpec.from_component("unet", unet2) -spec -ComponentSpec(name='unet', type_hint=, description=None, config=None, repo='stabilityai/stable-diffusion-xl-base-1.0', subfolder='unet', variant='fp16', revision=None, default_creation_method='from_pretrained') -unet2_recreated = spec.load(torch_dtype=torch.float16) -``` +from diffusers import ModularPipeline, ComponentsManager +import torch -The [`~ModularPipeline.get_component_spec`] method gets a copy of the current component specification to modify or update. +device = "cuda" +dtype = torch.bfloat16 +repo_id = "black-forest-labs/FLUX.2-klein-4B" -```py -unet_spec = t2i_pipeline.get_component_spec("unet") -unet_spec -ComponentSpec( - name='unet', - type_hint=, - pretrained_model_name_or_path='RunDiffusion/Juggernaut-XL-v9', - subfolder='unet', - variant='fp16', - default_creation_method='from_pretrained' -) +# get the blocks and separate out the text encoder +blocks = ModularPipeline.from_pretrained(repo_id).blocks +text_block = blocks.sub_blocks.pop("text_encoder") -# modify to load from a different repository -unet_spec.pretrained_model_name_or_path = "stabilityai/stable-diffusion-xl-base-1.0" +# use ComponentsManager to handle offloading across multiple pipelines +manager = ComponentsManager() +manager.enable_auto_cpu_offload(device=device) -# load component with modified spec -unet = unet_spec.load(torch_dtype=torch.float16) +# create separate pipelines for each stage +text_encoder_pipeline = text_block.init_pipeline(repo_id, components_manager=manager) +pipeline = blocks.init_pipeline(repo_id, components_manager=manager) + +# encode text +text_encoder_pipeline.load_components(torch_dtype=dtype) +text_embeddings = text_encoder_pipeline(prompt="a cat").get_by_kwargs("denoiser_input_fields") + +# denoise and decode +pipeline.load_components(torch_dtype=dtype) +output = pipeline( + **text_embeddings, + num_inference_steps=4, +).images[0] ``` +[`ComponentsManager`] handles memory across multiple pipelines. Unlike the offloading strategies in [`DiffusionPipeline`] that follow a fixed order, [`ComponentsManager`] makes offloading decisions dynamically each time a model forward pass runs, based on the current memory situation. This means it works regardless of how many pipelines you create or what order you run them in. See the [ComponentsManager](./components_manager) guide for more details. + +If pipeline stages share components (e.g., the same VAE used for encoding and decoding), you can use [`~ModularPipeline.update_components`] to pass an already-loaded component to another pipeline instead of loading it again. + ## Modular repository A repository is required if the pipeline blocks use *pretrained components*. The repository supplies loading specifications and metadata. -[`ModularPipeline`] specifically requires *modular repositories* (see [example repository](https://huggingface.co/YiYiXu/modular-diffdiff)) which are more flexible than a typical repository. It contains a `modular_model_index.json` file containing the following 3 elements. +[`ModularPipeline`] works with regular diffusers repositories out of the box. However, you can also create a *modular repository* for more flexibility. A modular repository contains a `modular_model_index.json` file containing the following 3 elements. -- `library` and `class` shows which library the component was loaded from and it's class. If `null`, the component hasn't been loaded yet. +- `library` and `class` shows which library the component was loaded from and its class. If `null`, the component hasn't been loaded yet. - `loading_specs_dict` contains the information required to load the component such as the repository and subfolder it is loaded from. -Unlike standard repositories, a modular repository can fetch components from different repositories based on the `loading_specs_dict`. Components don't need to exist in the same repository. +The key advantage of a modular repository is that components can be loaded from different repositories. For example, [diffusers/flux2-bnb-4bit-modular](https://huggingface.co/diffusers/flux2-bnb-4bit-modular) loads a quantized transformer from `diffusers/FLUX.2-dev-bnb-4bit` while loading the remaining components from `black-forest-labs/FLUX.2-dev`. + +To convert a regular diffusers repository into a modular one, create the pipeline using the regular repository, and then push to the Hub. The saved repository will contain a `modular_model_index.json` with all the loading specifications. -A modular repository may contain custom code for loading a [`ModularPipeline`]. This allows you to use specialized blocks that aren't native to Diffusers. +```py +from diffusers import ModularPipeline + +# load from a regular repo +pipeline = ModularPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0") + +# push as a modular repository +pipeline.save_pretrained("local/path", repo_id="my-username/sdxl-modular", push_to_hub=True) +``` + +A modular repository can also include custom pipeline blocks as Python code. This allows you to share specialized blocks that aren't native to Diffusers. For example, [diffusers/Florence2-image-Annotator](https://huggingface.co/diffusers/Florence2-image-Annotator) contains custom blocks alongside the loading configuration: ``` -modular-diffdiff-0704/ +Florence2-image-Annotator/ ├── block.py # Custom pipeline blocks implementation ├── config.json # Pipeline configuration and auto_map +├── mellon_config.json # UI configuration for Mellon └── modular_model_index.json # Component loading specifications ``` -The [config.json](https://huggingface.co/YiYiXu/modular-diffdiff-0704/blob/main/config.json) file contains an `auto_map` key that points to where a custom block is defined in `block.py`. +The `config.json` file contains an `auto_map` key that tells [`ModularPipeline`] where to find the custom blocks: ```json { - "_class_name": "DiffDiffBlocks", + "_class_name": "Florence2AnnotatorBlocks", "auto_map": { - "ModularPipelineBlocks": "block.DiffDiffBlocks" + "ModularPipelineBlocks": "block.Florence2AnnotatorBlocks" } } ``` + +Load custom code repositories with `trust_remote_code=True` as shown in [from_pretrained](#from_pretrained). See [Custom blocks](./custom_blocks) for how to create and share your own. \ No newline at end of file diff --git a/docs/source/en/modular_diffusers/overview.md b/docs/source/en/modular_diffusers/overview.md index 7d07c4b73434..83975200d664 100644 --- a/docs/source/en/modular_diffusers/overview.md +++ b/docs/source/en/modular_diffusers/overview.md @@ -24,7 +24,7 @@ The Modular Diffusers docs are organized as shown below. ## Quickstart -- A [quickstart](./quickstart) demonstrating how to implement an example workflow with Modular Diffusers. +- The [quickstart](./quickstart) shows you how to run a modular pipeline, understand its structure, and customize it by modifying the blocks that compose it. ## ModularPipelineBlocks @@ -33,9 +33,14 @@ The Modular Diffusers docs are organized as shown below. - [SequentialPipelineBlocks](./sequential_pipeline_blocks) is a type of block that chains multiple blocks so they run one after another, passing data along the chain. This guide shows you how to create [`~modular_pipelines.SequentialPipelineBlocks`] and how they connect and work together. - [LoopSequentialPipelineBlocks](./loop_sequential_pipeline_blocks) is a type of block that runs a series of blocks in a loop. This guide shows you how to create [`~modular_pipelines.LoopSequentialPipelineBlocks`]. - [AutoPipelineBlocks](./auto_pipeline_blocks) is a type of block that automatically chooses which blocks to run based on the input. This guide shows you how to create [`~modular_pipelines.AutoPipelineBlocks`]. +- [Building Custom Blocks](./custom_blocks) shows you how to create your own custom blocks and share them on the Hub. ## ModularPipeline - [ModularPipeline](./modular_pipeline) shows you how to create and convert pipeline blocks into an executable [`ModularPipeline`]. - [ComponentsManager](./components_manager) shows you how to manage and reuse components across multiple pipelines. -- [Guiders](./guiders) shows you how to use different guidance methods in the pipeline. \ No newline at end of file +- [Guiders](./guiders) shows you how to use different guidance methods in the pipeline. + +## Mellon Integration + +- [Using Custom Blocks with Mellon](./mellon) shows you how to make your custom blocks work with [Mellon](https://github.com/cubiq/Mellon), a visual node-based interface for building workflows. \ No newline at end of file diff --git a/docs/source/en/modular_diffusers/pipeline_block.md b/docs/source/en/modular_diffusers/pipeline_block.md index 06c115e1fb52..612736f3d307 100644 --- a/docs/source/en/modular_diffusers/pipeline_block.md +++ b/docs/source/en/modular_diffusers/pipeline_block.md @@ -25,81 +25,151 @@ This guide will show you how to create a [`~modular_pipelines.ModularPipelineBlo A [`~modular_pipelines.ModularPipelineBlocks`] requires `inputs`, and `intermediate_outputs`. -- `inputs` are values provided by a user and retrieved from the [`~modular_pipelines.PipelineState`]. This is useful because some workflows resize an image, but the original image is still required. The [`~modular_pipelines.PipelineState`] maintains the original image. +- `inputs` are values a block reads from the [`~modular_pipelines.PipelineState`] to perform its computation. These can be values provided by a user (like a prompt or image) or values produced by a previous block (like encoded `image_latents`). Use `InputParam` to define `inputs`. - ```py - from diffusers.modular_pipelines import InputParam - - user_inputs = [ - InputParam(name="image", type_hint="PIL.Image", description="raw input image to process") - ] - ``` +```py +class ImageEncodeStep(ModularPipelineBlocks): + ... + + @property + def inputs(self): + return [ + InputParam(name="image", type_hint="PIL.Image", required=True, description="raw input image to process"), + ] + ... +``` - `intermediate_outputs` are new values created by a block and added to the [`~modular_pipelines.PipelineState`]. The `intermediate_outputs` are available as `inputs` for subsequent blocks or available as the final output from running the pipeline. Use `OutputParam` to define `intermediate_outputs`. - ```py - from diffusers.modular_pipelines import OutputParam +```py +class ImageEncodeStep(ModularPipelineBlocks): + ... - user_intermediate_outputs = [ - OutputParam(name="image_latents", description="latents representing the image") - ] - ``` + @property + def intermediate_outputs(self): + return [ + OutputParam(name="image_latents", description="latents representing the image"), + ] + + ... +``` The intermediate inputs and outputs share data to connect blocks. They are accessible at any point, allowing you to track the workflow's progress. +## Components and configs + +The components and pipeline-level configs a block needs are specified in [`ComponentSpec`] and [`~modular_pipelines.ConfigSpec`]. + +- [`ComponentSpec`] contains the expected components used by a block. You need the `name` of the component and ideally a `type_hint` that specifies exactly what the component is. +- [`~modular_pipelines.ConfigSpec`] contains pipeline-level settings that control behavior across all blocks. + +```py +class ImageEncodeStep(ModularPipelineBlocks): + ... + + @property + def expected_components(self): + return [ + ComponentSpec(name="vae", type_hint=AutoencoderKL), + ] + + @property + def expected_configs(self): + return [ + ConfigSpec("force_zeros_for_empty_prompt", True), + ] + + ... +``` + +When the blocks are converted into a pipeline, the components become available to the block as the first argument in `__call__`. + ## Computation logic The computation a block performs is defined in the `__call__` method and it follows a specific structure. -1. Retrieve the [`~modular_pipelines.BlockState`] to get a local view of the `inputs` +1. Retrieve the [`~modular_pipelines.BlockState`] to get a local view of the `inputs`. 2. Implement the computation logic on the `inputs`. 3. Update [`~modular_pipelines.PipelineState`] to push changes from the local [`~modular_pipelines.BlockState`] back to the global [`~modular_pipelines.PipelineState`]. 4. Return the components and state which becomes available to the next block. ```py -def __call__(self, components, state): - # Get a local view of the state variables this block needs - block_state = self.get_block_state(state) +class ImageEncodeStep(ModularPipelineBlocks): + + def __call__(self, components, state): + # Get a local view of the state variables this block needs + block_state = self.get_block_state(state) - # Your computation logic here - # block_state contains all your inputs - # Access them like: block_state.image, block_state.processed_image + # Your computation logic here + # block_state contains all your inputs + # Access them like: block_state.image, block_state.processed_image - # Update the pipeline state with your updated block_states - self.set_block_state(state, block_state) - return components, state + # Update the pipeline state with your updated block_states + self.set_block_state(state, block_state) + return components, state ``` -### Components and configs +## Putting it all together -The components and pipeline-level configs a block needs are specified in [`ComponentSpec`] and [`~modular_pipelines.ConfigSpec`]. +Here is the complete block with all the pieces connected. -- [`ComponentSpec`] contains the expected components used by a block. You need the `name` of the component and ideally a `type_hint` that specifies exactly what the component is. -- [`~modular_pipelines.ConfigSpec`] contains pipeline-level settings that control behavior across all blocks. +```py +from diffusers import ComponentSpec, AutoencoderKL +from diffusers.modular_pipelines import InputParam, ModularPipelineBlocks, OutputParam + + +class ImageEncodeStep(ModularPipelineBlocks): + + @property + def description(self): + return "Encode an image into latent space." + + @property + def expected_components(self): + return [ + ComponentSpec(name="vae", type_hint=AutoencoderKL), + ] + + @property + def inputs(self): + return [ + InputParam(name="image", type_hint="PIL.Image", required=True, description="raw input image to process"), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam(name="image_latents", type_hint="torch.Tensor", description="latents representing the image"), + ] + + def __call__(self, components, state): + block_state = self.get_block_state(state) + block_state.image_latents = components.vae.encode(block_state.image) + self.set_block_state(state, block_state) + return components, state +``` + +Every block has a `doc` property that is automatically generated from the properties you defined above. It provides a summary of the block's description, components, inputs, and outputs. ```py -from diffusers import ComponentSpec, ConfigSpec +block = ImageEncoderStep() +print(block.doc) +class ImageEncodeStep -expected_components = [ - ComponentSpec(name="unet", type_hint=UNet2DConditionModel), - ComponentSpec(name="scheduler", type_hint=EulerDiscreteScheduler) -] + Encode an image into latent space. -expected_config = [ - ConfigSpec("force_zeros_for_empty_prompt", True) -] -``` + Components: + vae (`AutoencoderKL`) -When the blocks are converted into a pipeline, the components become available to the block as the first argument in `__call__`. + Inputs: + image (`PIL.Image`): + raw input image to process -```py -def __call__(self, components, state): - # Access components using dot notation - unet = components.unet - vae = components.vae - scheduler = components.scheduler -``` + Outputs: + image_latents (`torch.Tensor`): + latents representing the image +``` \ No newline at end of file diff --git a/docs/source/en/modular_diffusers/quickstart.md b/docs/source/en/modular_diffusers/quickstart.md index 32d14d84e243..884495e107b4 100644 --- a/docs/source/en/modular_diffusers/quickstart.md +++ b/docs/source/en/modular_diffusers/quickstart.md @@ -12,333 +12,286 @@ specific language governing permissions and limitations under the License. # Quickstart -Modular Diffusers is a framework for quickly building flexible and customizable pipelines. At the core of Modular Diffusers are [`ModularPipelineBlocks`] that can be combined with other blocks to adapt to new workflows. The blocks are converted into a [`ModularPipeline`], a friendly user-facing interface developers can use. +Modular Diffusers is a framework for quickly building flexible and customizable pipelines. These pipelines can go beyond what standard `DiffusionPipeline`s can do. At the core of Modular Diffusers are [`ModularPipelineBlocks`] that can be combined with other blocks to adapt to new workflows. The blocks are converted into a [`ModularPipeline`], a friendly user-facing interface for running generation tasks. -This doc will show you how to implement a [Differential Diffusion](https://differential-diffusion.github.io/) pipeline with the modular framework. +This guide shows you how to run a modular pipeline, understand its structure, and customize it by modifying the blocks that compose it. -## ModularPipelineBlocks +## Run a pipeline -[`ModularPipelineBlocks`] are *definitions* that specify the components, inputs, outputs, and computation logic for a single step in a pipeline. There are four types of blocks. +[`ModularPipeline`] is the main interface for loading, running, and managing modular pipelines. +```py +import torch +from diffusers import ModularPipeline, ComponentsManager -- [`ModularPipelineBlocks`] is the most basic block for a single step. -- [`SequentialPipelineBlocks`] is a multi-block that composes other blocks linearly. The outputs of one block are the inputs to the next block. -- [`LoopSequentialPipelineBlocks`] is a multi-block that runs iteratively and is designed for iterative workflows. -- [`AutoPipelineBlocks`] is a collection of blocks for different workflows and it selects which block to run based on the input. It is designed to conveniently package multiple workflows into a single pipeline. +# Use ComponentsManager to enable auto CPU offloading for memory efficiency +manager = ComponentsManager() +manager.enable_auto_cpu_offload(device="cuda:0") -[Differential Diffusion](https://differential-diffusion.github.io/) is an image-to-image workflow. Start with the `IMAGE2IMAGE_BLOCKS` preset, a collection of `ModularPipelineBlocks` for image-to-image generation. +pipe = ModularPipeline.from_pretrained("Qwen/Qwen-Image", components_manager=manager) +pipe.load_components(torch_dtype=torch.bfloat16) -```py -from diffusers.modular_pipelines.stable_diffusion_xl import IMAGE2IMAGE_BLOCKS -IMAGE2IMAGE_BLOCKS = InsertableDict([ - ("text_encoder", StableDiffusionXLTextEncoderStep), - ("image_encoder", StableDiffusionXLVaeEncoderStep), - ("input", StableDiffusionXLInputStep), - ("set_timesteps", StableDiffusionXLImg2ImgSetTimestepsStep), - ("prepare_latents", StableDiffusionXLImg2ImgPrepareLatentsStep), - ("prepare_add_cond", StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep), - ("denoise", StableDiffusionXLDenoiseStep), - ("decode", StableDiffusionXLDecodeStep) -]) +image = pipe( + prompt="cat wizard with red hat, gandalf, lord of the rings, detailed, fantasy, cute, adorable, Pixar, Disney", +).images[0] +image ``` -## Pipeline and block states - -Modular Diffusers uses *state* to communicate data between blocks. There are two types of states. +[`~ModularPipeline.from_pretrained`] uses lazy loading - it reads the configuration to learn where to load each component from, but doesn't actually load the model weights until you call [`~ModularPipeline.load_components`]. This gives you control over when and how components are loaded. -- [`PipelineState`] is a global state that can be used to track all inputs and outputs across all blocks. -- [`BlockState`] is a local view of relevant variables from [`PipelineState`] for an individual block. +> [!TIP] +> `ComponentsManager` with `enable_auto_cpu_offload` automatically moves models between CPU and GPU as needed, reducing memory usage for large models like Qwen-Image. Learn more in the [ComponentsManager](./components_manager) guide. +> +> If you don't need offloading, remove the `components_manager` argument and move the pipeline to your device manually with `to("cuda")`. -## Customizing blocks +Learn more about creating and loading pipelines in the [Creating a pipeline](https://huggingface.co/docs/diffusers/modular_diffusers/modular_pipeline#creating-a-pipeline) and [Loading components](https://huggingface.co/docs/diffusers/modular_diffusers/modular_pipeline#loading-components) guides. -[Differential Diffusion](https://differential-diffusion.github.io/) differs from standard image-to-image in its `prepare_latents` and `denoise` blocks. All the other blocks can be reused, but you'll need to modify these two. +## Understand the structure -Create placeholder `ModularPipelineBlocks` for `prepare_latents` and `denoise` by copying and modifying the existing ones. - -Print the `denoise` block to see that it is composed of [`LoopSequentialPipelineBlocks`] with three sub-blocks, `before_denoiser`, `denoiser`, and `after_denoiser`. Only the `before_denoiser` sub-block needs to be modified to prepare the latent input for the denoiser based on the change map. +A [`ModularPipeline`] has two parts: a **definition** (the blocks) and a **state** (the loaded components and configs). +Print the pipeline to see its state — the components and their loading status and configuration. ```py -denoise_blocks = IMAGE2IMAGE_BLOCKS["denoise"]() -print(denoise_blocks) +print(pipe) +``` +``` +QwenImageModularPipeline { + "_blocks_class_name": "QwenImageAutoBlocks", + "_class_name": "QwenImageModularPipeline", + "_diffusers_version": "0.37.0.dev0", + "transformer": [ + "diffusers", + "QwenImageTransformer2DModel", + { + "pretrained_model_name_or_path": "Qwen/Qwen-Image", + "revision": null, + "subfolder": "transformer", + "type_hint": [ + "diffusers", + "QwenImageTransformer2DModel" + ], + "variant": null + } + ], + ... +} ``` -Replace the `StableDiffusionXLLoopBeforeDenoiser` sub-block with the new `SDXLDiffDiffLoopBeforeDenoiser` block. - +Access the definition through `pipe.blocks` — this is the [`~modular_pipelines.ModularPipelineBlocks`] that defines the pipeline's workflows, inputs, outputs, and computation logic. ```py -# Copy existing blocks as placeholders -class SDXLDiffDiffPrepareLatentsStep(ModularPipelineBlocks): - """Copied from StableDiffusionXLImg2ImgPrepareLatentsStep - will modify later""" - # ... same implementation as StableDiffusionXLImg2ImgPrepareLatentsStep - -class SDXLDiffDiffDenoiseStep(StableDiffusionXLDenoiseLoopWrapper): - block_classes = [SDXLDiffDiffLoopBeforeDenoiser, StableDiffusionXLLoopDenoiser, StableDiffusionXLLoopAfterDenoiser] - block_names = ["before_denoiser", "denoiser", "after_denoiser"] +print(pipe.blocks) ``` - -### prepare_latents - -The `prepare_latents` block requires the following changes. - -- a processor to process the change map -- a new `inputs` to accept the user-provided change map, `timestep` for precomputing all the latents and `num_inference_steps` to create the mask for updating the image regions -- update the computation in the `__call__` method for processing the change map and creating the masks, and storing it in the [`BlockState`] - -```diff -class SDXLDiffDiffPrepareLatentsStep(ModularPipelineBlocks): - @property - def expected_components(self) -> List[ComponentSpec]: - return [ - ComponentSpec("vae", AutoencoderKL), - ComponentSpec("scheduler", EulerDiscreteScheduler), -+ ComponentSpec("mask_processor", VaeImageProcessor, config=FrozenDict({"do_normalize": False, "do_convert_grayscale": True})) - ] - @property - def inputs(self) -> List[Tuple[str, Any]]: - return [ - InputParam("generator"), -+ InputParam("diffdiff_map", required=True), -- InputParam("latent_timestep", required=True, type_hint=torch.Tensor), -+ InputParam("timesteps", type_hint=torch.Tensor), -+ InputParam("num_inference_steps", type_hint=int), - ] - - @property - def intermediate_outputs(self) -> List[OutputParam]: - return [ -+ OutputParam("original_latents", type_hint=torch.Tensor), -+ OutputParam("diffdiff_masks", type_hint=torch.Tensor), - ] - def __call__(self, components, state: PipelineState): - # ... existing logic ... -+ # Process change map and create masks -+ diffdiff_map = components.mask_processor.preprocess(block_state.diffdiff_map, height=latent_height, width=latent_width) -+ thresholds = torch.arange(block_state.num_inference_steps, dtype=diffdiff_map.dtype) / block_state.num_inference_steps -+ block_state.diffdiff_masks = diffdiff_map > (thresholds + (block_state.denoising_start or 0)) -+ block_state.original_latents = block_state.latents ``` - -### denoise - -The `before_denoiser` sub-block requires the following changes. - -- a new `inputs` to accept a `denoising_start` parameter, `original_latents` and `diffdiff_masks` from the `prepare_latents` block -- update the computation in the `__call__` method for applying Differential Diffusion - -```diff -class SDXLDiffDiffLoopBeforeDenoiser(ModularPipelineBlocks): - @property - def description(self) -> str: - return ( - "Step within the denoising loop for differential diffusion that prepare the latent input for the denoiser" - ) - - @property - def inputs(self) -> List[str]: - return [ - InputParam("latents", required=True, type_hint=torch.Tensor), -+ InputParam("denoising_start"), -+ InputParam("original_latents", type_hint=torch.Tensor), -+ InputParam("diffdiff_masks", type_hint=torch.Tensor), - ] - - def __call__(self, components, block_state, i, t): -+ # Apply differential diffusion logic -+ if i == 0 and block_state.denoising_start is None: -+ block_state.latents = block_state.original_latents[:1] -+ else: -+ block_state.mask = block_state.diffdiff_masks[i].unsqueeze(0).unsqueeze(1) -+ block_state.latents = block_state.original_latents[i] * block_state.mask + block_state.latents * (1 - block_state.mask) - - # ... rest of existing logic ... +QwenImageAutoBlocks( + Class: SequentialPipelineBlocks + + Description: Auto Modular pipeline for text-to-image, image-to-image, inpainting, and controlnet tasks using QwenImage. + + Supported workflows: + - `text2image`: requires `prompt` + - `image2image`: requires `prompt`, `image` + - `inpainting`: requires `prompt`, `mask_image`, `image` + - `controlnet_text2image`: requires `prompt`, `control_image` + ... + + Components: + text_encoder (`Qwen2_5_VLForConditionalGeneration`) + vae (`AutoencoderKLQwenImage`) + transformer (`QwenImageTransformer2DModel`) + ... + + Sub-Blocks: + [0] text_encoder (QwenImageAutoTextEncoderStep) + [1] vae_encoder (QwenImageAutoVaeEncoderStep) + [2] controlnet_vae_encoder (QwenImageOptionalControlNetVaeEncoderStep) + [3] denoise (QwenImageAutoCoreDenoiseStep) + [4] decode (QwenImageAutoDecodeStep) +) ``` -## Assembling the blocks - -You should have all the blocks you need at this point to create a [`ModularPipeline`]. +The output returns: +- The supported workflows (text2image, image2image, inpainting, etc.) +- The Sub-Blocks it's composed of (text_encoder, vae_encoder, denoise, decode) -Copy the existing `IMAGE2IMAGE_BLOCKS` preset and for the `set_timesteps` block, use the `set_timesteps` from the `TEXT2IMAGE_BLOCKS` because Differential Diffusion doesn't require a `strength` parameter. +### Workflows -Set the `prepare_latents` and `denoise` blocks to the `SDXLDiffDiffPrepareLatentsStep` and `SDXLDiffDiffDenoiseStep` blocks you just modified. - -Call [`SequentialPipelineBlocks.from_blocks_dict`] on the blocks to create a `SequentialPipelineBlocks`. +This pipeline supports multiple workflows and adapts its behavior based on the inputs you provide. For example, if you pass `image` to the pipeline, it runs an image-to-image workflow instead of text-to-image. Learn more about how this works under the hood in the [AutoPipelineBlocks](https://huggingface.co/docs/diffusers/modular_diffusers/auto_pipeline_blocks) guide. ```py -DIFFDIFF_BLOCKS = IMAGE2IMAGE_BLOCKS.copy() -DIFFDIFF_BLOCKS["set_timesteps"] = TEXT2IMAGE_BLOCKS["set_timesteps"] -DIFFDIFF_BLOCKS["prepare_latents"] = SDXLDiffDiffPrepareLatentsStep -DIFFDIFF_BLOCKS["denoise"] = SDXLDiffDiffDenoiseStep - -dd_blocks = SequentialPipelineBlocks.from_blocks_dict(DIFFDIFF_BLOCKS) -print(dd_blocks) -``` +from diffusers.utils import load_image -## ModularPipeline +input_image = load_image("https://github.com/Trgtuan10/Image_storage/blob/main/cute_cat.png?raw=true") -Convert the [`SequentialPipelineBlocks`] into a [`ModularPipeline`] with the [`ModularPipeline.init_pipeline`] method. This initializes the expected components to load from a `modular_model_index.json` file. Explicitly load the components by calling [`ModularPipeline.load_components`]. - -It is a good idea to initialize the [`ComponentManager`] with the pipeline to help manage the different components. Once you call [`~ModularPipeline.load_components`], the components are registered to the [`ComponentManager`] and can be shared between workflows. The example below uses the `collection` argument to assign the components a `"diffdiff"` label for better organization. +image = pipe( + prompt="cat wizard with red hat, gandalf, lord of the rings, detailed, fantasy, cute, adorable, Pixar, Disney", + image=input_image, +).images[0] +``` +Use `get_workflow()` to extract the blocks for a specific workflow. Pass the workflow name (e.g., `"image2image"`, `"inpainting"`, `"controlnet_text2image"`) to get only the blocks relevant to that workflow. This is useful when you want to customize or debug a specific workflow. You can check `pipe.blocks.available_workflows` to see all available workflows. ```py -from diffusers.modular_pipelines import ComponentsManager - -components = ComponentManager() - -dd_pipeline = dd_blocks.init_pipeline("YiYiXu/modular-demo-auto", components_manager=components, collection="diffdiff") -dd_pipeline.load_componenets(torch_dtype=torch.float16) -dd_pipeline.to("cuda") +img2img_blocks = pipe.blocks.get_workflow("image2image") ``` -## Adding workflows - -Other workflows can be added to the [`ModularPipeline`] to support additional features without rewriting the entire pipeline from scratch. - -This section demonstrates how to add an IP-Adapter or ControlNet. - -### IP-Adapter -Stable Diffusion XL already has a preset IP-Adapter block that you can use and doesn't require any changes to the existing Differential Diffusion pipeline. +### Sub-blocks -```py -from diffusers.modular_pipelines.stable_diffusion_xl.encoders import StableDiffusionXLAutoIPAdapterStep +Blocks can contain other blocks. `pipe.blocks` gives you the top-level block definition (here, `QwenImageAutoBlocks`), while `sub_blocks` lets you access the smaller blocks inside it. -ip_adapter_block = StableDiffusionXLAutoIPAdapterStep() -``` +`QwenImageAutoBlocks` is composed of: `text_encoder`, `vae_encoder`, `controlnet_vae_encoder`, `denoise`, and `decode`. -Use the [`sub_blocks.insert`] method to insert it into the [`ModularPipeline`]. The example below inserts the `ip_adapter_block` at position `0`. Print the pipeline to see that the `ip_adapter_block` is added and it requires an `ip_adapter_image`. This also added two components to the pipeline, the `image_encoder` and `feature_extractor`. +These sub-blocks run one after another and data flows linearly from one block to the next — each block's `intermediate_outputs` become available as `inputs` to the next block. This is how [`SequentialPipelineBlocks`](./sequential_pipeline_blocks) work. +You can access them through the `sub_blocks` property. The `doc` property is useful for seeing the full documentation of any block, including its inputs, outputs, and components. ```py -dd_blocks.sub_blocks.insert("ip_adapter", ip_adapter_block, 0) +vae_encoder_block = pipe.blocks.sub_blocks["vae_encoder"] +print(vae_encoder_block.doc) ``` -Call [`~ModularPipeline.init_pipeline`] to initialize a [`ModularPipeline`] and use [`~ModularPipeline.load_components`] to load the model components. Load and set the IP-Adapter to run the pipeline. - +This block can be converted to a pipeline so that it can run on its own with [`~ModularPipelineBlocks.init_pipeline`]. ```py -dd_pipeline = dd_blocks.init_pipeline("YiYiXu/modular-demo-auto", collection="diffdiff") -dd_pipeline.load_components(torch_dtype=torch.float16) -dd_pipeline.loader.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter_sdxl.bin") -dd_pipeline.loader.set_ip_adapter_scale(0.6) -dd_pipeline = dd_pipeline.to(device) - -ip_adapter_image = load_image("https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/diffdiff_orange.jpeg") -image = load_image("https://huggingface.co/datasets/OzzyGT/testing-resources/resolve/main/differential/20240329211129_4024911930.png?download=true") -mask = load_image("https://huggingface.co/datasets/OzzyGT/testing-resources/resolve/main/differential/gradient_mask.png?download=true") +vae_encoder_pipe = vae_encoder_block.init_pipeline() -prompt = "a green pear" -negative_prompt = "blurry" -generator = torch.Generator(device=device).manual_seed(42) +# Reuse the VAE we already loaded, we can reuse it with update_components() method +vae_encoder_pipe.update_components(vae=pipe.vae) -image = dd_pipeline( - prompt=prompt, - negative_prompt=negative_prompt, - num_inference_steps=25, - generator=generator, - ip_adapter_image=ip_adapter_image, - diffdiff_map=mask, - image=image, - output="images" -)[0] +# Run just this block +image_latents = vae_encoder_pipe(image=input_image).image_latents +print(image_latents.shape) ``` -### ControlNet +It reuses the VAE from our original pipeline instead of reloading it, keeping memory usage efficient. Learn more in the [Loading components](https://huggingface.co/docs/diffusers/modular_diffusers/modular_pipeline#loading-components) guide. + +Since blocks are composable, you can modify the pipeline's definition by adding, removing, or swapping blocks to create new workflows. In the next section, we'll add a canny edge detection block to a ControlNet pipeline, so you can pass a regular image instead of a pre-processed canny edge map. -Stable Diffusion XL already has a preset ControlNet block that can readily be used. +## Compose new workflows +Let's add a canny edge detection block to a ControlNet pipeline. First, load a pre-built canny block from the Hub (see [Building Custom Blocks](https://huggingface.co/docs/diffusers/modular_diffusers/custom_blocks) to create your own). ```py -from diffusers.modular_pipelines.stable_diffusion_xl.modular_blocks import StableDiffusionXLAutoControlNetInputStep +from diffusers.modular_pipelines import ModularPipelineBlocks -control_input_block = StableDiffusionXLAutoControlNetInputStep() +# Load a canny block from the Hub +canny_block = ModularPipelineBlocks.from_pretrained( + "diffusers-internal-dev/canny-filtering", + trust_remote_code=True, +) + +print(canny_block.doc) +``` +``` +class CannyBlock + + Inputs: + image (`Union[Image, ndarray]`): + Image to compute canny filter on + low_threshold (`int`, *optional*, defaults to 50): + Low threshold for the canny filter. + high_threshold (`int`, *optional*, defaults to 200): + High threshold for the canny filter. + ... + + Outputs: + control_image (`PIL.Image`): + Canny map for input image ``` -However, it requires modifying the `denoise` block because that's where the ControlNet injects the control information into the UNet. +Use `get_workflow` to extract the ControlNet workflow from [`QwenImageAutoBlocks`]. +```py +# Get the controlnet workflow that we want to work with +blocks = pipe.blocks.get_workflow("controlnet_text2image") +print(blocks.doc) +``` +``` +class SequentialPipelineBlocks + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + control_image (`Image`): + Control image for ControlNet conditioning. + ... +``` -Modify the `denoise` block by replacing the `StableDiffusionXLLoopDenoiser` sub-block with the `StableDiffusionXLControlNetLoopDenoiser`. +The extracted workflow is a [`SequentialPipelineBlocks`](./sequential_pipeline_blocks) and it currently requires `control_image` as input. Insert the canny block at the beginning so the pipeline accepts a regular image instead. ```py -class SDXLDiffDiffControlNetDenoiseStep(StableDiffusionXLDenoiseLoopWrapper): - block_classes = [SDXLDiffDiffLoopBeforeDenoiser, StableDiffusionXLControlNetLoopDenoiser, StableDiffusionXLDenoiseLoopAfterDenoiser] - block_names = ["before_denoiser", "denoiser", "after_denoiser"] +# Insert canny at the beginning +blocks.sub_blocks.insert("canny", canny_block, 0) -controlnet_denoise_block = SDXLDiffDiffControlNetDenoiseStep() +# Check the updated structure: CannyBlock is now listed as first sub-block +print(blocks) +# Check the updated doc +print(blocks.doc) +``` +``` +class SequentialPipelineBlocks + + Inputs: + image (`Union[Image, ndarray]`): + Image to compute canny filter on + low_threshold (`int`, *optional*, defaults to 50): + Low threshold for the canny filter. + high_threshold (`int`, *optional*, defaults to 200): + High threshold for the canny filter. + prompt (`str`): + The prompt or prompts to guide image generation. + ... ``` -Insert the `controlnet_input` block and replace the `denoise` block with the new `controlnet_denoise_block`. Initialize a [`ModularPipeline`] and [`~ModularPipeline.load_components`] into it. +Now the pipeline takes `image` as input instead of `control_image`. Because blocks in a sequence share data automatically, the canny block's output (`control_image`) flows to the denoise block that needs it, and the canny block's input (`image`) becomes a pipeline input since no earlier block provides it. +Create a pipeline from the modified blocks and load a ControlNet model. The ControlNet isn't part of the original model repository, so load it separately and add it with [`~ModularPipeline.update_components`]. ```py -dd_blocks.sub_blocks.insert("controlnet_input", control_input_block, 7) -dd_blocks.sub_blocks["denoise"] = controlnet_denoise_block - -dd_pipeline = dd_blocks.init_pipeline("YiYiXu/modular-demo-auto", collection="diffdiff") -dd_pipeline.load_components(torch_dtype=torch.float16) -dd_pipeline = dd_pipeline.to(device) - -control_image = load_image("https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/diffdiff_tomato_canny.jpeg") -image = load_image("https://huggingface.co/datasets/OzzyGT/testing-resources/resolve/main/differential/20240329211129_4024911930.png?download=true") -mask = load_image("https://huggingface.co/datasets/OzzyGT/testing-resources/resolve/main/differential/gradient_mask.png?download=true") +pipeline = blocks.init_pipeline("Qwen/Qwen-Image", components_manager=manager) -prompt = "a green pear" -negative_prompt = "blurry" -generator = torch.Generator(device=device).manual_seed(42) +pipeline.load_components(torch_dtype=torch.bfloat16) -image = dd_pipeline( - prompt=prompt, - negative_prompt=negative_prompt, - num_inference_steps=25, - generator=generator, - control_image=control_image, - controlnet_conditioning_scale=0.5, - diffdiff_map=mask, - image=image, - output="images" -)[0] +# Load the ControlNet model +controlnet_spec = pipeline.get_component_spec("controlnet") +controlnet_spec.pretrained_model_name_or_path = "InstantX/Qwen-Image-ControlNet-Union" +controlnet = controlnet_spec.load(torch_dtype=torch.bfloat16) +pipeline.update_components(controlnet=controlnet) ``` -### AutoPipelineBlocks - -The Differential Diffusion, IP-Adapter, and ControlNet workflows can be bundled into a single [`ModularPipeline`] by using [`AutoPipelineBlocks`]. This allows automatically selecting which sub-blocks to run based on the inputs like `control_image` or `ip_adapter_image`. If none of these inputs are passed, then it defaults to the Differential Diffusion. +Now run the pipeline - the canny block preprocesses the image for ControlNet. +```py +from diffusers.utils import load_image -Use `block_trigger_inputs` to only run the `SDXLDiffDiffControlNetDenoiseStep` block if a `control_image` input is provided. Otherwise, the `SDXLDiffDiffDenoiseStep` is used. +prompt = "cat wizard with red hat, gandalf, lord of the rings, detailed, fantasy, cute, adorable, Pixar, Disney" +image = load_image("https://github.com/Trgtuan10/Image_storage/blob/main/cute_cat.png?raw=true") -```py -class SDXLDiffDiffAutoDenoiseStep(AutoPipelineBlocks): - block_classes = [SDXLDiffDiffControlNetDenoiseStep, SDXLDiffDiffDenoiseStep] - block_names = ["controlnet_denoise", "denoise"] - block_trigger_inputs = ["controlnet_cond", None] +output = pipeline( + prompt=prompt, + image=image, +).images[0] +output ``` -Add the `ip_adapter` and `controlnet_input` blocks. +## Next steps -```py -DIFFDIFF_AUTO_BLOCKS = IMAGE2IMAGE_BLOCKS.copy() -DIFFDIFF_AUTO_BLOCKS["prepare_latents"] = SDXLDiffDiffPrepareLatentsStep -DIFFDIFF_AUTO_BLOCKS["set_timesteps"] = TEXT2IMAGE_BLOCKS["set_timesteps"] -DIFFDIFF_AUTO_BLOCKS["denoise"] = SDXLDiffDiffAutoDenoiseStep -DIFFDIFF_AUTO_BLOCKS.insert("ip_adapter", StableDiffusionXLAutoIPAdapterStep, 0) -DIFFDIFF_AUTO_BLOCKS.insert("controlnet_input",StableDiffusionXLControlNetAutoInput, 7) -``` + + -Call [`SequentialPipelineBlocks.from_blocks_dict`] to create a [`SequentialPipelineBlocks`] and create a [`ModularPipeline`] and load in the model components to run. +Understand the core building blocks of Modular Diffusers: -```py -dd_auto_blocks = SequentialPipelineBlocks.from_blocks_dict(DIFFDIFF_AUTO_BLOCKS) -dd_pipeline = dd_auto_blocks.init_pipeline("YiYiXu/modular-demo-auto", collection="diffdiff") -dd_pipeline.load_components(torch_dtype=torch.float16) -``` +- [ModularPipelineBlocks](./pipeline_block): The basic unit for defining a step in a pipeline. +- [SequentialPipelineBlocks](./sequential_pipeline_blocks): Chain blocks to run in sequence. +- [AutoPipelineBlocks](./auto_pipeline_blocks): Create pipelines that support multiple workflows. +- [States](./modular_diffusers_states): How data is shared between blocks. -## Share + + -Add your [`ModularPipeline`] to the Hub with [`~ModularPipeline.save_pretrained`] and set `push_to_hub` argument to `True`. +Learn how to create your own blocks with custom logic in the [Building Custom Blocks](./custom_blocks) guide. -```py -dd_pipeline.save_pretrained("YiYiXu/test_modular_doc", push_to_hub=True) -``` + + -Other users can load the [`ModularPipeline`] with [`~ModularPipeline.from_pretrained`]. +Use [`ComponentsManager`](./components_manager) to share models across multiple pipelines and manage memory efficiently. -```py -import torch -from diffusers.modular_pipelines import ModularPipeline, ComponentsManager + + -components = ComponentsManager() +Connect modular pipelines to [Mellon](https://github.com/cubiq/Mellon), a visual node-based interface for building workflows. Custom blocks built with Modular Diffusers work out of the box with Mellon - no UI code required. Read more in the Mellon guide. -diffdiff_pipeline = ModularPipeline.from_pretrained("YiYiXu/modular-diffdiff-0704", trust_remote_code=True, components_manager=components, collection="diffdiff") -diffdiff_pipeline.load_components(torch_dtype=torch.float16) -``` + + \ No newline at end of file diff --git a/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md b/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md index f1549a26b86f..1bd67e17b8bf 100644 --- a/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md +++ b/docs/source/en/modular_diffusers/sequential_pipeline_blocks.md @@ -91,23 +91,42 @@ class ImageEncoderBlock(ModularPipelineBlocks): -Connect the two blocks by defining an [`InsertableDict`] to map the block names to the block instances. Blocks are executed in the order they're registered in `blocks_dict`. - -Use [`~modular_pipelines.SequentialPipelineBlocks.from_blocks_dict`] to create a [`~modular_pipelines.SequentialPipelineBlocks`]. +Connect the two blocks by defining a [`~modular_pipelines.SequentialPipelineBlocks`]. List the block instances in `block_classes` and their corresponding names in `block_names`. The blocks are executed in the order they appear in `block_classes`, and data flows from one block to the next through [`~modular_pipelines.PipelineState`]. ```py -from diffusers.modular_pipelines import SequentialPipelineBlocks, InsertableDict - -blocks_dict = InsertableDict() -blocks_dict["input"] = input_block -blocks_dict["image_encoder"] = image_encoder_block +class ImageProcessingStep(SequentialPipelineBlocks): + """ + # auto_docstring + """ + model_name = "my_model" + block_classes = [InputBlock(), ImageEncoderBlock()] + block_names = ["input", "image_encoder"] -blocks = SequentialPipelineBlocks.from_blocks_dict(blocks_dict) + @property + def description(self): + return ( + "Process text prompts and images for the pipeline. It:\n" + " - Determines the batch size from the prompts.\n" + " - Encodes the image into latent space." + ) ``` -Inspect the sub-blocks in [`~modular_pipelines.SequentialPipelineBlocks`] by calling `blocks`, and for more details about the inputs and outputs, access the `docs` attribute. +When you create a [`~modular_pipelines.SequentialPipelineBlocks`], properties like `inputs`, `intermediate_outputs`, and `expected_components` are automatically aggregated from the sub-blocks, so there is no need to define them again. + +There are a few properties you should set: + +- `description`: We recommend adding a description for the assembled block to explain what the combined step does. +- `model_name`: This is automatically derived from the sub-blocks but isn't always correct, so you may need to override it. +- `outputs`: By default this is the same as `intermediate_outputs`, but you can manually set it to control which values appear in the doc. This is useful for showing only the final outputs instead of all intermediate values. + +These properties, together with the aggregated `inputs`, `intermediate_outputs`, and `expected_components`, are used to automatically generate the `doc` property. + + +Print the `ImageProcessingStep` block to inspect its sub-blocks, and use `doc` for a full summary of the block's inputs, outputs, and components. + ```py +blocks = ImageProcessingStep() print(blocks) print(blocks.doc) -``` +``` \ No newline at end of file diff --git a/docs/source/en/optimization/cache.md b/docs/source/en/optimization/cache.md index 6397c7d4cd2e..4eccd70cb304 100644 --- a/docs/source/en/optimization/cache.md +++ b/docs/source/en/optimization/cache.md @@ -68,6 +68,20 @@ config = FasterCacheConfig( pipeline.transformer.enable_cache(config) ``` +## FirstBlockCache + +[FirstBlock Cache](https://huggingface.co/docs/diffusers/main/en/api/cache#diffusers.FirstBlockCacheConfig) checks how much the early layers of the denoiser changes from one timestep to the next. If the change is small, the model skips the expensive later layers and reuses the previous output. + +```py +import torch +from diffusers import DiffusionPipeline +from diffusers.hooks import apply_first_block_cache, FirstBlockCacheConfig + +pipeline = DiffusionPipeline.from_pretrained( + "Qwen/Qwen-Image", torch_dtype=torch.bfloat16 +) +apply_first_block_cache(pipeline.transformer, FirstBlockCacheConfig(threshold=0.2)) +``` ## TaylorSeer Cache [TaylorSeer Cache](https://huggingface.co/papers/2403.06923) accelerates diffusion inference by using Taylor series expansions to approximate and cache intermediate activations across denoising steps. The method predicts future outputs based on past computations, reusing them at specified intervals to reduce redundant calculations. @@ -87,8 +101,7 @@ from diffusers import FluxPipeline, TaylorSeerCacheConfig pipe = FluxPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, -) -pipe.to("cuda") +).to("cuda") config = TaylorSeerCacheConfig( cache_interval=5, @@ -97,4 +110,58 @@ config = TaylorSeerCacheConfig( taylor_factors_dtype=torch.bfloat16, ) pipe.transformer.enable_cache(config) -``` \ No newline at end of file +``` + +## MagCache + +[MagCache](https://github.com/Zehong-Ma/MagCache) accelerates inference by skipping transformer blocks based on the magnitude of the residual update. It observes that the magnitude of updates (Output - Input) decays predictably over the diffusion process. By accumulating an "error budget" based on pre-computed magnitude ratios, it dynamically decides when to skip computation and reuse the previous residual. + +MagCache relies on **Magnitude Ratios** (`mag_ratios`), which describe this decay curve. These ratios are specific to the model checkpoint and scheduler. + +### Usage + +To use MagCache, you typically follow a two-step process: **Calibration** and **Inference**. + +1. **Calibration**: Run inference once with `calibrate=True`. The hook will measure the residual magnitudes and print the calculated ratios to the console. +2. **Inference**: Pass these ratios to `MagCacheConfig` to enable acceleration. + +```python +import torch +from diffusers import FluxPipeline, MagCacheConfig + +pipe = FluxPipeline.from_pretrained( + "black-forest-labs/FLUX.1-schnell", + torch_dtype=torch.bfloat16 +).to("cuda") + +# 1. Calibration Step +# Run full inference to measure model behavior. +calib_config = MagCacheConfig(calibrate=True, num_inference_steps=4) +pipe.transformer.enable_cache(calib_config) + +# Run a prompt to trigger calibration +pipe("A cat playing chess", num_inference_steps=4) +# Logs will print something like: "MagCache Calibration Results: [1.0, 1.37, 0.97, 0.87]" + +# 2. Inference Step +# Apply the specific ratios obtained from calibration for optimized speed. +# Note: For Flux models, you can also import defaults: +# from diffusers.hooks.mag_cache import FLUX_MAG_RATIOS +mag_config = MagCacheConfig( + mag_ratios=[1.0, 1.37, 0.97, 0.87], + num_inference_steps=4 +) + +pipe.transformer.enable_cache(mag_config) + +image = pipe("A cat playing chess", num_inference_steps=4).images[0] +``` + +> [!NOTE] +> `mag_ratios` represent the model's intrinsic magnitude decay curve. Ratios calibrated for a high number of steps (e.g., 50) can be reused for lower step counts (e.g., 20). The implementation uses interpolation to map the curve to the current number of inference steps. + +> [!TIP] +> For pipelines that run Classifier-Free Guidance sequentially (like Kandinsky 5.0), the calibration log might print two arrays: one for the Conditional pass and one for the Unconditional pass. In most cases, you should use the first array (Conditional). + +> [!TIP] +> For pipelines that run Classifier-Free Guidance in a **batched** manner (like SDXL or Flux), the `hidden_states` processed by the model contain both conditional and unconditional branches concatenated together. The calibration process automatically accounts for this, producing a single array of ratios that represents the joint behavior. You can use this resulting array directly without modification. diff --git a/docs/source/en/quantization/torchao.md b/docs/source/en/quantization/torchao.md index 18cc109e0785..de90c3006e8f 100644 --- a/docs/source/en/quantization/torchao.md +++ b/docs/source/en/quantization/torchao.md @@ -33,7 +33,7 @@ pipeline_quant_config = PipelineQuantizationConfig( ) pipeline = DiffusionPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", - quantzation_config=pipeline_quant_config, + quantization_config=pipeline_quant_config, torch_dtype=torch.bfloat16, device_map="cuda" ) @@ -50,7 +50,7 @@ pipeline_quant_config = PipelineQuantizationConfig( ) pipeline = DiffusionPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", - quantzation_config=pipeline_quant_config, + quantization_config=pipeline_quant_config, torch_dtype=torch.bfloat16, device_map="cuda" ) @@ -66,11 +66,11 @@ from diffusers import DiffusionPipeline, PipelineQuantizationConfig, TorchAoConf from torchao.quantization import Int4WeightOnlyConfig pipeline_quant_config = PipelineQuantizationConfig( - quant_mapping={"transformer": TorchAoConfig(Int4WeightOnlyConfig(group_size=128)))} + quant_mapping={"transformer": TorchAoConfig(Int4WeightOnlyConfig(group_size=128))} ) pipeline = DiffusionPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", - quantzation_config=pipeline_quant_config, + quantization_config=pipeline_quant_config, torch_dtype=torch.bfloat16, device_map="cuda" ) @@ -83,25 +83,6 @@ Refer to this [table](https://github.com/huggingface/diffusers/pull/10009#issue- > [!TIP] > The FP8 post-training quantization schemes in torchao are effective for GPUs with compute capability of at least 8.9 (RTX-4090, Hopper, etc.). FP8 often provides the best speed, memory, and quality trade-off when generating images and videos. We recommend combining FP8 and torch.compile if your GPU is compatible. -## autoquant - -torchao provides [autoquant](https://docs.pytorch.org/ao/stable/generated/torchao.quantization.autoquant.html#torchao.quantization.autoquant) an automatic quantization API. Autoquantization chooses the best quantization strategy by comparing the performance of each strategy on chosen input types and shapes. This is only supported in Diffusers for individual models at the moment. - -```py -import torch -from diffusers import DiffusionPipeline -from torchao.quantization import autoquant - -# Load the pipeline -pipeline = DiffusionPipeline.from_pretrained( - "black-forest-labs/FLUX.1-schnell", - torch_dtype=torch.bfloat16, - device_map="cuda" -) - -transformer = autoquant(pipeline.transformer) -``` - ## Supported quantization types torchao supports weight-only quantization and weight and dynamic-activation quantization for int8, float3-float8, and uint1-uint7. diff --git a/docs/source/en/training/distributed_inference.md b/docs/source/en/training/distributed_inference.md index 22e8a30427b9..c9a341df40c5 100644 --- a/docs/source/en/training/distributed_inference.md +++ b/docs/source/en/training/distributed_inference.md @@ -111,7 +111,7 @@ if __name__ == "__main__": Call `torchrun` to run the inference script and use the `--nproc_per_node` argument to set the number of GPUs to use. ```bash -torchrun run_distributed.py --nproc_per_node=2 +torchrun --nproc_per_node=2 run_distributed.py ``` ## device_map @@ -263,8 +263,8 @@ def main(): world_size = dist.get_world_size() pipeline = DiffusionPipeline.from_pretrained( - "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, device_map=device - ) + "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16 + ).to(device) pipeline.transformer.set_attention_backend("_native_cudnn") cp_config = ContextParallelConfig(ring_degree=world_size) @@ -314,6 +314,63 @@ Pass the [`ContextParallelConfig`] to [`~ModelMixin.enable_parallelism`]. pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2)) ``` +### Unified Attention + +[Unified Sequence Parallelism](https://huggingface.co/papers/2405.07719) combines Ring Attention and Ulysses Attention into a single approach for efficient long-sequence processing. It applies Ulysses's *all-to-all* communication first to redistribute heads and sequence tokens, then uses Ring Attention to process the redistributed data, and finally reverses the *all-to-all* to restore the original layout. + +This hybrid approach leverages the strengths of both methods: +- **Ulysses Attention** efficiently parallelizes across attention heads +- **Ring Attention** handles very long sequences with minimal memory overhead +- Together, they enable 2D parallelization across both heads and sequence dimensions + +[`ContextParallelConfig`] supports Unified Attention by specifying both `ulysses_degree` and `ring_degree`. The total number of devices used is `ulysses_degree * ring_degree`, arranged in a 2D grid where Ulysses and Ring groups are orthogonal (non-overlapping). +Pass the [`ContextParallelConfig`] with both `ulysses_degree` and `ring_degree` set to bigger than 1 to [`~ModelMixin.enable_parallelism`]. + +```py +pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ring_degree=2)) +``` + +> [!TIP] +> Unified Attention is to be used when there are enough devices to arrange in a 2D grid (at least 4 devices). + +We ran a benchmark with Ulysess, Ring, and Unified Attention with [this script](https://github.com/huggingface/diffusers/pull/12693#issuecomment-3694727532) on a node of 4 H100 GPUs. The results are summarized as follows: + +| CP Backend | Time / Iter (ms) | Steps / Sec | Peak Memory (GB) | +|--------------------|------------------|-------------|------------------| +| ulysses | 6670.789 | 7.50 | 33.85 | +| ring | 13076.492 | 3.82 | 56.02 | +| unified_balanced | 11068.705 | 4.52 | 33.85 | + +From the above table, it's clear that Ulysses provides better throughput, but the number of devices it can use remains limited to the number of attention heads, a limitation that is solved by unified attention. + + +### Ulysses Anything Attention + +The default Ulysses Attention mechanism requires that the sequence length of hidden states must be divisible by the number of devices. This imposes significant limitations on the practical application of Ulysses Attention. [Ulysses Anything Attention](https://github.com/huggingface/diffusers/pull/12996) is a variant of Ulysses Attention that supports arbitrary sequence lengths and arbitrary numbers of attention heads, thereby enhancing the versatility of Ulysses Attention in practical use. + +[`ContextParallelConfig`] supports Ulysses Anything Attention by specifying both `ulysses_degree` and `ulysses_anything`. Please note that Ulysses Anything Attention is not currently supported by Unified Attention. Pass the [`ContextParallelConfig`] with both `ulysses_degree` set to bigger than 1 and `ulysses_anything=True` to [`~ModelMixin.enable_parallelism`]. + +```py +pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ulysses_anything=True)) +``` + +> [!TIP] To avoid multiple forced CUDA sync caused by H2D and D2H transfers, please add the **gloo** backend in `init_process_group`. This will significantly reduce communication latency. + +We ran a benchmark for FLUX.1-dev with Ulysses, Ring, Unified Attention and Ulysses Anything Attention with [this script](https://github.com/huggingface/diffusers/pull/12996#issuecomment-3797695999) on a node of 4 L20 GPUs. The results are summarized as follows: + +| CP Backend | Time / Iter (ms) | Steps / Sec | Peak Memory (GB) | Shape (HxW)| +|--------------------|------------------|-------------|------------------|------------| +| ulysses | 281.07 | 3.56 | 37.11 | 1024x1024 | +| ring | 351.34 | 2.85 | 37.01 | 1024x1024 | +| unified_balanced | 324.37 | 3.08 | 37.16 | 1024x1024 | +| ulysses_anything | 280.94 | 3.56 | 37.11 | 1024x1024 | +| ulysses | failed | failed | failed | 1008x1008 | +| ring | failed | failed | failed | 1008x1008 | +| unified_balanced | failed | failed | failed | 1008x1008 | +| ulysses_anything | 278.40 | 3.59 | 36.99 | 1008x1008 | + +From the above table, it is clear that Ulysses Anything Attention offers better compatibility with arbitrary sequence lengths while maintaining the same performance as the standard Ulysses Attention. + ### parallel_config Pass `parallel_config` during model initialization to enable context parallelism. diff --git a/docs/source/en/using-diffusers/automodel.md b/docs/source/en/using-diffusers/automodel.md index 957cbd17e3f7..82d4d14a10a9 100644 --- a/docs/source/en/using-diffusers/automodel.md +++ b/docs/source/en/using-diffusers/automodel.md @@ -29,8 +29,31 @@ text_encoder = AutoModel.from_pretrained( ) ``` +## Custom models + [`AutoModel`] also loads models from the [Hub](https://huggingface.co/models) that aren't included in Diffusers. Set `trust_remote_code=True` in [`AutoModel.from_pretrained`] to load custom models. +A custom model repository needs a Python module with the model class, and a `config.json` with an `auto_map` entry that maps `"AutoModel"` to `"module_file.ClassName"`. + +``` +custom/custom-transformer-model/ +├── config.json +├── my_model.py +└── diffusion_pytorch_model.safetensors +``` + +The `config.json` includes the `auto_map` field pointing to the custom class. + +```json +{ + "auto_map": { + "AutoModel": "my_model.MyCustomModel" + } +} +``` + +Then load it with `trust_remote_code=True`. + ```py import torch from diffusers import AutoModel @@ -40,7 +63,66 @@ transformer = AutoModel.from_pretrained( ) ``` +For a real-world example, [Overworld/Waypoint-1-Small](https://huggingface.co/Overworld/Waypoint-1-Small/tree/main/transformer) hosts a custom `WorldModel` class across several modules in its `transformer` subfolder. + +``` +transformer/ +├── config.json # auto_map: "model.WorldModel" +├── model.py +├── attn.py +├── nn.py +├── cache.py +├── quantize.py +├── __init__.py +└── diffusion_pytorch_model.safetensors +``` + +```py +import torch +from diffusers import AutoModel + +transformer = AutoModel.from_pretrained( + "Overworld/Waypoint-1-Small", subfolder="transformer", trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="cuda" +) +``` + If the custom model inherits from the [`ModelMixin`] class, it gets access to the same features as Diffusers model classes, like [regional compilation](../optimization/fp16#regional-compilation) and [group offloading](../optimization/memory#group-offloading). +> [!WARNING] +> As a precaution with `trust_remote_code=True`, pass a commit hash to the `revision` argument in [`AutoModel.from_pretrained`] to make sure the code hasn't been updated with new malicious code (unless you fully trust the model owners). +> +> ```py +> transformer = AutoModel.from_pretrained( +> "Overworld/Waypoint-1-Small", subfolder="transformer", trust_remote_code=True, revision="a3d8cb2" +> ) +> ``` + +### Saving custom models + +Use [`~ConfigMixin.register_for_auto_class`] to add the `auto_map` entry to `config.json` automatically when saving. This avoids having to manually edit the config file. + +```py +# my_model.py +from diffusers import ModelMixin, ConfigMixin + +class MyCustomModel(ModelMixin, ConfigMixin): + ... + +MyCustomModel.register_for_auto_class("AutoModel") + +model = MyCustomModel(...) +model.save_pretrained("./my_model") +``` + +The saved `config.json` will include the `auto_map` field. + +```json +{ + "auto_map": { + "AutoModel": "my_model.MyCustomModel" + } +} +``` + > [!NOTE] > Learn more about implementing custom models in the [Community components](../using-diffusers/custom_pipeline_overview#community-components) guide. \ No newline at end of file diff --git a/docs/source/en/using-diffusers/consisid.md b/docs/source/en/using-diffusers/consisid.md index b6b04ddaf57e..96ece5b20c3a 100644 --- a/docs/source/en/using-diffusers/consisid.md +++ b/docs/source/en/using-diffusers/consisid.md @@ -60,7 +60,7 @@ export_to_video(video.frames[0], "output.mp4", fps=8) Face Image Video - DescriptionDescription diff --git a/docs/source/en/using-diffusers/helios.md b/docs/source/en/using-diffusers/helios.md new file mode 100644 index 000000000000..ced7c6298f23 --- /dev/null +++ b/docs/source/en/using-diffusers/helios.md @@ -0,0 +1,133 @@ + +# Helios + +[Helios](https://github.com/PKU-YuanGroup/Helios) is the first 14B video generation model that runs at 19.5 FPS on a single NVIDIA H100 GPU and supports minute-scale generation while matching the quality of a strong baseline, natively integrating T2V, I2V, and V2V tasks within a unified architecture. The main features of Helios are: + +- Without commonly used anti-drifting strategies (eg, self-forcing, error-banks, keyframe sampling, or inverted sampling), Helios generates minute-scale videos with high quality and strong coherence. +- Without standard acceleration techniques (eg, KV-cache, causal masking, sparse/linear attention, TinyVAE, progressive noise schedules, hidden-state caching, or quantization), Helios achieves 19.5 FPS in end-to-end inference for a 14B video generation model on a single H100 GPU. +- Introducing optimizations that improve both training and inference throughput while reducing memory consumption. These changes enable training a 14B video generation model without parallelism or sharding infrastructure, with batch sizes comparable to image models. + +This guide will walk you through using Helios for use cases. + +## Load Model Checkpoints + +Model weights may be stored in separate subfolders on the Hub or locally, in which case, you should use the [`~DiffusionPipeline.from_pretrained`] method. + +```python +import torch +from diffusers import HeliosPipeline, HeliosPyramidPipeline +from huggingface_hub import snapshot_download + +# For Best Quality +snapshot_download(repo_id="BestWishYsh/Helios-Base", local_dir="BestWishYsh/Helios-Base") +pipe = HeliosPipeline.from_pretrained("BestWishYsh/Helios-Base", torch_dtype=torch.bfloat16) +pipe.to("cuda") + +# Intermediate Weight +snapshot_download(repo_id="BestWishYsh/Helios-Mid", local_dir="BestWishYsh/Helios-Mid") +pipe = HeliosPyramidPipeline.from_pretrained("BestWishYsh/Helios-Mid", torch_dtype=torch.bfloat16) +pipe.to("cuda") + +# For Best Efficiency +snapshot_download(repo_id="BestWishYsh/Helios-Distilled", local_dir="BestWishYsh/Helios-Distilled") +pipe = HeliosPyramidPipeline.from_pretrained("BestWishYsh/Helios-Distilled", torch_dtype=torch.bfloat16) +pipe.to("cuda") +``` + +## Text-to-Video Showcases + + + + + + + + + + + + + + +
PromptGenerated Video
A Viking warrior driving a modern city bus filled with passengers. The Viking has long blonde hair tied back, a beard, and is adorned with a fur-lined helmet and armor. He wears a traditional tunic and trousers, but also sports a seatbelt as he focuses on navigating the busy streets. The interior of the bus is typical, with rows of seats occupied by diverse passengers going about their daily routines. The exterior shots show the bustling urban environment, including tall buildings and traffic. Medium shot focusing on the Viking at the wheel, with occasional close-ups of his determined expression. + + +
A documentary-style nature photography shot from a camera truck moving to the left, capturing a crab quickly scurrying into its burrow. The crab has a hard, greenish-brown shell and long claws, moving with determined speed across the sandy ground. Its body is slightly arched as it burrows into the sand, leaving a small trail behind. The background shows a shallow beach with scattered rocks and seashells, and the horizon features a gentle curve of the coastline. The photo has a natural and realistic texture, emphasizing the crab's natural movement and the texture of the sand. A close-up shot from a slightly elevated angle. + + +
+ +## Image-to-Video Showcases + + + + + + + + + + + + + + + + + +
ImagePromptGenerated Video
A sleek red Kia car speeds along a rural road under a cloudy sky, its modern design and dynamic movement emphasized by the blurred motion of the surrounding fields and trees stretching into the distance. The car's glossy exterior reflects the overcast sky, highlighting its aerodynamic shape and sporty stance. The license plate reads "KIA 626," and the vehicle's headlights are on, adding to the sense of motion and energy. The road curves gently, with the car positioned slightly off-center, creating a sense of forward momentum. A dynamic front three-quarter view captures the car's powerful presence against the serene backdrop of rolling hills and scattered trees. + + +
A close-up captures a fluffy orange cat with striking green eyes and white whiskers, gazing intently towards the camera. The cat's fur is soft and well-groomed, with a mix of warm orange and cream tones. Its large, expressive eyes are a vivid green, reflecting curiosity and alertness. The cat's nose is small and pink, and its mouth is slightly open, revealing a hint of its pink tongue. The background is softly blurred, suggesting a cozy indoor setting with neutral tones. The photo has a shallow depth of field, focusing sharply on the cat's face while the background remains out of focus. A close-up shot from a slightly elevated perspective. + + +
+ +## Interactive-Video Showcases + + + + + + + + + + + + + + +
PromptGenerated Video
The prompt can be found here + +
The prompt can be found here + +
+ +## Resources + +Learn more about Helios with the following resources. +- Watch [video1](https://www.youtube.com/watch?v=vd_AgHtOUFQ) and [video2](https://www.youtube.com/watch?v=1GeIU2Dn7UY) for a demonstration of Helios's key features. +- The research paper, [Helios: Real Real-Time Long Video Generation Model](https://huggingface.co/papers/2603.04379) for more details. diff --git a/docs/source/zh/_toctree.yml b/docs/source/zh/_toctree.yml index 337d010fc74d..ab9eaf6ec7fb 100644 --- a/docs/source/zh/_toctree.yml +++ b/docs/source/zh/_toctree.yml @@ -132,6 +132,8 @@ sections: - local: using-diffusers/consisid title: ConsisID + - local: using-diffusers/helios + title: Helios - title: Resources isExpanded: false diff --git a/docs/source/zh/community_projects.md b/docs/source/zh/community_projects.md index 0440142452f1..ffa45f1e9bb0 100644 --- a/docs/source/zh/community_projects.md +++ b/docs/source/zh/community_projects.md @@ -26,6 +26,14 @@ http://www.apache.org/licenses/LICENSE-2.0 项目名称 描述 + + helios + Helios:比1.3B更低开销、更快且更强的14B的实时长视频生成模型 + + + consisid + ConsisID:零样本身份保持的文本到视频生成模型 + dream-textures Stable Diffusion内置到Blender diff --git a/docs/source/zh/modular_diffusers/guiders.md b/docs/source/zh/modular_diffusers/guiders.md index 50436f90c4a5..2315625a197a 100644 --- a/docs/source/zh/modular_diffusers/guiders.md +++ b/docs/source/zh/modular_diffusers/guiders.md @@ -86,10 +86,7 @@ t2i_pipeline.guider ## 更改引导器参数 -引导器参数可以通过 [`~ComponentSpec.create`] 方法或 [`~ModularPipeline.update_components`] 方法进行调整。下面的示例更改了 `guidance_scale` 值。 - - - +引导器参数可以通过 [`~ComponentSpec.create`] 方法以及 [`~ModularPipeline.update_components`] 方法进行调整。下面的示例更改了 `guidance_scale` 值。 ```py guider_spec = t2i_pipeline.get_component_spec("guider") @@ -97,18 +94,6 @@ guider = guider_spec.create(guidance_scale=10) t2i_pipeline.update_components(guider=guider) ``` - - - -```py -guider_spec = t2i_pipeline.get_component_spec("guider") -guider_spec.config["guidance_scale"] = 10 -t2i_pipeline.update_components(guider=guider_spec) -``` - - - - ## 上传自定义引导器 在自定义引导器上调用 [`~utils.PushToHubMixin.push_to_hub`] 方法,将其分享到 Hub。 diff --git a/docs/source/zh/using-diffusers/helios.md b/docs/source/zh/using-diffusers/helios.md new file mode 100644 index 000000000000..5f7f067eb781 --- /dev/null +++ b/docs/source/zh/using-diffusers/helios.md @@ -0,0 +1,134 @@ + +# Helios + +[Helios](https://github.com/PKU-YuanGroup/Helios) 是首个能够在单张 NVIDIA H100 GPU 上以 19.5 FPS 运行的 14B 视频生成模型。它在支持分钟级视频生成的同时,拥有媲美强大基线模型的生成质量,并在统一架构下原生集成了文生视频(T2V)、图生视频(I2V)和视频生视频(V2V)任务。Helios 的主要特性包括: + +- 无需常用的防漂移策略(例如:自强制/self-forcing、误差库/error-banks、关键帧采样或逆采样),我们的模型即可生成高质量且高度连贯的分钟级视频。 +- 无需标准的加速技术(例如:KV 缓存、因果掩码、稀疏/线性注意力机制、TinyVAE、渐进式噪声调度、隐藏状态缓存或量化),作为一款 14B 规模的视频生成模型,我们在单张 H100 GPU 上的端到端推理速度便达到了 19.5 FPS。 +- 引入了多项优化方案,在降低显存消耗的同时,显著提升了训练与推理的吞吐量。这些改进使得我们无需借助并行或分片(sharding)等基础设施,即可使用与图像模型相当的批大小(batch sizes)来训练 14B 的视频生成模型。 + +本指南将引导您完成 Helios 在不同场景下的使用。 + +## Load Model Checkpoints + +模型权重可以存储在Hub上或本地的单独子文件夹中,在这种情况下,您应该使用 [`~DiffusionPipeline.from_pretrained`] 方法。 + +```python +import torch +from diffusers import HeliosPipeline, HeliosPyramidPipeline +from huggingface_hub import snapshot_download + +# For Best Quality +snapshot_download(repo_id="BestWishYsh/Helios-Base", local_dir="BestWishYsh/Helios-Base") +pipe = HeliosPipeline.from_pretrained("BestWishYsh/Helios-Base", torch_dtype=torch.bfloat16) +pipe.to("cuda") + +# Intermediate Weight +snapshot_download(repo_id="BestWishYsh/Helios-Mid", local_dir="BestWishYsh/Helios-Mid") +pipe = HeliosPyramidPipeline.from_pretrained("BestWishYsh/Helios-Mid", torch_dtype=torch.bfloat16) +pipe.to("cuda") + +# For Best Efficiency +snapshot_download(repo_id="BestWishYsh/Helios-Distilled", local_dir="BestWishYsh/Helios-Distilled") +pipe = HeliosPyramidPipeline.from_pretrained("BestWishYsh/Helios-Distilled", torch_dtype=torch.bfloat16) +pipe.to("cuda") +``` + +## Text-to-Video Showcases + + + + + + + + + + + + + + +
PromptGenerated Video
A Viking warrior driving a modern city bus filled with passengers. The Viking has long blonde hair tied back, a beard, and is adorned with a fur-lined helmet and armor. He wears a traditional tunic and trousers, but also sports a seatbelt as he focuses on navigating the busy streets. The interior of the bus is typical, with rows of seats occupied by diverse passengers going about their daily routines. The exterior shots show the bustling urban environment, including tall buildings and traffic. Medium shot focusing on the Viking at the wheel, with occasional close-ups of his determined expression. + + +
A documentary-style nature photography shot from a camera truck moving to the left, capturing a crab quickly scurrying into its burrow. The crab has a hard, greenish-brown shell and long claws, moving with determined speed across the sandy ground. Its body is slightly arched as it burrows into the sand, leaving a small trail behind. The background shows a shallow beach with scattered rocks and seashells, and the horizon features a gentle curve of the coastline. The photo has a natural and realistic texture, emphasizing the crab's natural movement and the texture of the sand. A close-up shot from a slightly elevated angle. + + +
+ +## Image-to-Video Showcases + + + + + + + + + + + + + + + + + +
ImagePromptGenerated Video
A sleek red Kia car speeds along a rural road under a cloudy sky, its modern design and dynamic movement emphasized by the blurred motion of the surrounding fields and trees stretching into the distance. The car's glossy exterior reflects the overcast sky, highlighting its aerodynamic shape and sporty stance. The license plate reads "KIA 626," and the vehicle's headlights are on, adding to the sense of motion and energy. The road curves gently, with the car positioned slightly off-center, creating a sense of forward momentum. A dynamic front three-quarter view captures the car's powerful presence against the serene backdrop of rolling hills and scattered trees. + + +
A close-up captures a fluffy orange cat with striking green eyes and white whiskers, gazing intently towards the camera. The cat's fur is soft and well-groomed, with a mix of warm orange and cream tones. Its large, expressive eyes are a vivid green, reflecting curiosity and alertness. The cat's nose is small and pink, and its mouth is slightly open, revealing a hint of its pink tongue. The background is softly blurred, suggesting a cozy indoor setting with neutral tones. The photo has a shallow depth of field, focusing sharply on the cat's face while the background remains out of focus. A close-up shot from a slightly elevated perspective. + + +
+ +## Interactive-Video Showcases + + + + + + + + + + + + + + +
PromptGenerated Video
The prompt can be found here + +
The prompt can be found here + +
+ +## Resources + +通过以下资源了解有关 Helios 的更多信息: + +- [视频1](https://www.youtube.com/watch?v=vd_AgHtOUFQ)和[视频2](https://www.youtube.com/watch?v=1GeIU2Dn7UY)演示了 Helios 的主要功能; +- 有关更多详细信息,请参阅研究论文 [Helios: Real Real-Time Long Video Generation Model](https://huggingface.co/papers/2603.04379)。 diff --git a/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py b/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py index 05f2b1ee17f3..608ab3ef3135 100644 --- a/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py +++ b/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py @@ -94,7 +94,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/advanced_diffusion_training/train_dreambooth_lora_sd15_advanced.py b/examples/advanced_diffusion_training/train_dreambooth_lora_sd15_advanced.py index 8fba00afc39e..a47e4dd96dcb 100644 --- a/examples/advanced_diffusion_training/train_dreambooth_lora_sd15_advanced.py +++ b/examples/advanced_diffusion_training/train_dreambooth_lora_sd15_advanced.py @@ -88,7 +88,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py b/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py index cf0a1588f39b..dcaa5a38fc37 100644 --- a/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py +++ b/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py @@ -95,7 +95,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) @@ -1929,6 +1929,8 @@ def compute_text_embeddings(prompt, text_encoders, tokenizers, clip_skip): if args.cache_latents: latents_cache = [] + # Store vae config before potential deletion + vae_scaling_factor = vae.config.scaling_factor for batch in tqdm(train_dataloader, desc="Caching latents"): with torch.no_grad(): batch["pixel_values"] = batch["pixel_values"].to( @@ -1940,6 +1942,8 @@ def compute_text_embeddings(prompt, text_encoders, tokenizers, clip_skip): del vae if torch.cuda.is_available(): torch.cuda.empty_cache() + else: + vae_scaling_factor = vae.config.scaling_factor # Scheduler and math around the number of training steps. # Check the PR https://github.com/huggingface/diffusers/pull/8312 for detailed explanation. @@ -2109,13 +2113,13 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): model_input = vae.encode(pixel_values).latent_dist.sample() if latents_mean is None and latents_std is None: - model_input = model_input * vae.config.scaling_factor + model_input = model_input * vae_scaling_factor if args.pretrained_vae_model_name_or_path is None: model_input = model_input.to(weight_dtype) else: latents_mean = latents_mean.to(device=model_input.device, dtype=model_input.dtype) latents_std = latents_std.to(device=model_input.device, dtype=model_input.dtype) - model_input = (model_input - latents_mean) * vae.config.scaling_factor / latents_std + model_input = (model_input - latents_mean) * vae_scaling_factor / latents_std model_input = model_input.to(dtype=weight_dtype) # Sample noise that we'll add to the latents diff --git a/examples/cogvideo/train_cogvideox_image_to_video_lora.py b/examples/cogvideo/train_cogvideox_image_to_video_lora.py index 113d9b57398e..17a9dd47d3ba 100644 --- a/examples/cogvideo/train_cogvideox_image_to_video_lora.py +++ b/examples/cogvideo/train_cogvideox_image_to_video_lora.py @@ -61,7 +61,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) @@ -149,13 +149,13 @@ def get_args(): "--validation_prompt", type=str, default=None, - help="One or more prompt(s) that is used during validation to verify that the model is learning. Multiple validation prompts should be separated by the '--validation_prompt_seperator' string.", + help="One or more prompt(s) that is used during validation to verify that the model is learning. Multiple validation prompts should be separated by the '--validation_prompt_separator' string.", ) parser.add_argument( "--validation_images", type=str, default=None, - help="One or more image path(s) that is used during validation to verify that the model is learning. Multiple validation paths should be separated by the '--validation_prompt_seperator' string. These should correspond to the order of the validation prompts.", + help="One or more image path(s) that is used during validation to verify that the model is learning. Multiple validation paths should be separated by the '--validation_prompt_separator' string. These should correspond to the order of the validation prompts.", ) parser.add_argument( "--validation_prompt_separator", @@ -432,9 +432,9 @@ def get_args(): class VideoDataset(Dataset): def __init__( self, - instance_data_root: Optional[str] = None, - dataset_name: Optional[str] = None, - dataset_config_name: Optional[str] = None, + instance_data_root: str | None = None, + dataset_name: str | None = None, + dataset_config_name: str | None = None, caption_column: str = "text", video_column: str = "video", height: int = 480, @@ -443,8 +443,8 @@ def __init__( max_num_frames: int = 49, skip_frames_start: int = 0, skip_frames_end: int = 0, - cache_dir: Optional[str] = None, - id_token: Optional[str] = None, + cache_dir: str | None = None, + id_token: str | None = None, ) -> None: super().__init__() diff --git a/examples/cogvideo/train_cogvideox_lora.py b/examples/cogvideo/train_cogvideox_lora.py index bcafe4ecf5d9..984ed697d7c7 100644 --- a/examples/cogvideo/train_cogvideox_lora.py +++ b/examples/cogvideo/train_cogvideox_lora.py @@ -52,7 +52,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) @@ -140,7 +140,7 @@ def get_args(): "--validation_prompt", type=str, default=None, - help="One or more prompt(s) that is used during validation to verify that the model is learning. Multiple validation prompts should be separated by the '--validation_prompt_seperator' string.", + help="One or more prompt(s) that is used during validation to verify that the model is learning. Multiple validation prompts should be separated by the '--validation_prompt_separator' string.", ) parser.add_argument( "--validation_prompt_separator", @@ -416,9 +416,9 @@ def get_args(): class VideoDataset(Dataset): def __init__( self, - instance_data_root: Optional[str] = None, - dataset_name: Optional[str] = None, - dataset_config_name: Optional[str] = None, + instance_data_root: str | None = None, + dataset_name: str | None = None, + dataset_config_name: str | None = None, caption_column: str = "text", video_column: str = "video", height: int = 480, @@ -428,8 +428,8 @@ def __init__( max_num_frames: int = 49, skip_frames_start: int = 0, skip_frames_end: int = 0, - cache_dir: Optional[str] = None, - id_token: Optional[str] = None, + cache_dir: str | None = None, + id_token: str | None = None, ) -> None: super().__init__() @@ -1232,22 +1232,49 @@ def load_model_hook(models, input_dir): id_token=args.id_token, ) - def encode_video(video, bar): - bar.update(1) + def encode_video(video): video = video.to(accelerator.device, dtype=vae.dtype).unsqueeze(0) video = video.permute(0, 2, 1, 3, 4) # [B, C, F, H, W] latent_dist = vae.encode(video).latent_dist return latent_dist + # Distribute video encoding across processes: each process only encodes its own shard + num_videos = len(train_dataset.instance_videos) + num_procs = accelerator.num_processes + local_rank = accelerator.process_index + local_count = len(range(local_rank, num_videos, num_procs)) + progress_encode_bar = tqdm( - range(0, len(train_dataset.instance_videos)), - desc="Loading Encode videos", + range(local_count), + desc="Encoding videos", + disable=not accelerator.is_local_main_process, ) - train_dataset.instance_videos = [ - encode_video(video, progress_encode_bar) for video in train_dataset.instance_videos - ] + + encoded_videos = [None] * num_videos + for i, video in enumerate(train_dataset.instance_videos): + if i % num_procs == local_rank: + encoded_videos[i] = encode_video(video) + progress_encode_bar.update(1) progress_encode_bar.close() + # Broadcast encoded latent distributions so every process has the full set + if num_procs > 1: + import torch.distributed as dist + + from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution + + ref_params = next(v for v in encoded_videos if v is not None).parameters + for i in range(num_videos): + src = i % num_procs + if encoded_videos[i] is not None: + params = encoded_videos[i].parameters.contiguous() + else: + params = torch.empty_like(ref_params) + dist.broadcast(params, src=src) + encoded_videos[i] = DiagonalGaussianDistribution(params) + + train_dataset.instance_videos = encoded_videos + def collate_fn(examples): videos = [example["instance_video"].sample() * vae.config.scaling_factor for example in examples] prompts = [example["instance_prompt"] for example in examples] diff --git a/examples/cogview4-control/train_control_cogview4.py b/examples/cogview4-control/train_control_cogview4.py index 6f06ed749635..d381a7902723 100644 --- a/examples/cogview4-control/train_control_cogview4.py +++ b/examples/cogview4-control/train_control_cogview4.py @@ -60,7 +60,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/community/README_community_scripts.md b/examples/community/README_community_scripts.md index 3c9ad0d89bb4..d790b67a44c5 100644 --- a/examples/community/README_community_scripts.md +++ b/examples/community/README_community_scripts.md @@ -260,7 +260,7 @@ class SDPromptSchedulingCallback(PipelineCallback): def callback_fn( self, pipeline, step_index, timestep, callback_kwargs - ) -> Dict[str, Any]: + ) -> dict[str, Any]: cutoff_step_ratio = self.config.cutoff_step_ratio cutoff_step_index = self.config.cutoff_step_index if isinstance(self.config.encoded_prompt, tuple): @@ -343,7 +343,7 @@ class SDXLPromptSchedulingCallback(PipelineCallback): def callback_fn( self, pipeline, step_index, timestep, callback_kwargs - ) -> Dict[str, Any]: + ) -> dict[str, Any]: cutoff_step_ratio = self.config.cutoff_step_ratio cutoff_step_index = self.config.cutoff_step_index if isinstance(self.config.encoded_prompt, tuple): diff --git a/examples/community/adaptive_mask_inpainting.py b/examples/community/adaptive_mask_inpainting.py index da67debe72ec..0378ffcdf985 100644 --- a/examples/community/adaptive_mask_inpainting.py +++ b/examples/community/adaptive_mask_inpainting.py @@ -871,7 +871,7 @@ def __call__( latents: Optional[torch.FloatTensor] = None, prompt_embeds: Optional[torch.FloatTensor] = None, negative_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/bit_diffusion.py b/examples/community/bit_diffusion.py index 67f4cd3fe199..c19bc95eefe0 100644 --- a/examples/community/bit_diffusion.py +++ b/examples/community/bit_diffusion.py @@ -231,9 +231,9 @@ def __call__( height: Optional[int] = 256, width: Optional[int] = 256, num_inference_steps: Optional[int] = 50, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, batch_size: Optional[int] = 1, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, **kwargs, ) -> Union[Tuple, ImagePipelineOutput]: diff --git a/examples/community/clip_guided_images_mixing_stable_diffusion.py b/examples/community/clip_guided_images_mixing_stable_diffusion.py index 2cd3daf68c24..bacf9f0ec8b4 100644 --- a/examples/community/clip_guided_images_mixing_stable_diffusion.py +++ b/examples/community/clip_guided_images_mixing_stable_diffusion.py @@ -235,8 +235,8 @@ def __call__( self, style_image: Union[torch.Tensor, PIL.Image.Image], content_image: Union[torch.Tensor, PIL.Image.Image], - style_prompt: Optional[str] = None, - content_prompt: Optional[str] = None, + style_prompt: str | None = None, + content_prompt: str | None = None, height: Optional[int] = 512, width: Optional[int] = 512, noise_strength: float = 0.6, @@ -245,8 +245,8 @@ def __call__( batch_size: Optional[int] = 1, eta: float = 0.0, clip_guidance_scale: Optional[float] = 100, - generator: Optional[torch.Generator] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | None = None, + output_type: str | None = "pil", return_dict: bool = True, slerp_latent_style_strength: float = 0.8, slerp_prompt_style_strength: float = 0.1, diff --git a/examples/community/clip_guided_stable_diffusion.py b/examples/community/clip_guided_stable_diffusion.py index bfd0858d245e..013df7acc7c7 100644 --- a/examples/community/clip_guided_stable_diffusion.py +++ b/examples/community/clip_guided_stable_diffusion.py @@ -179,9 +179,9 @@ def __call__( clip_prompt: Optional[Union[str, List[str]]] = None, num_cutouts: Optional[int] = 4, use_cutouts: Optional[bool] = True, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, ): if isinstance(prompt, str): diff --git a/examples/community/clip_guided_stable_diffusion_img2img.py b/examples/community/clip_guided_stable_diffusion_img2img.py index f3dd4903f851..739d6dafff73 100644 --- a/examples/community/clip_guided_stable_diffusion_img2img.py +++ b/examples/community/clip_guided_stable_diffusion_img2img.py @@ -316,9 +316,9 @@ def __call__( clip_prompt: Optional[Union[str, List[str]]] = None, num_cutouts: Optional[int] = 4, use_cutouts: Optional[bool] = True, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, ): if isinstance(prompt, str): diff --git a/examples/community/composable_stable_diffusion.py b/examples/community/composable_stable_diffusion.py index a7c610ad4355..466502e8394c 100644 --- a/examples/community/composable_stable_diffusion.py +++ b/examples/community/composable_stable_diffusion.py @@ -357,13 +357,13 @@ def __call__( negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, - weights: Optional[str] = "", + weights: str | None = "", ): r""" Function invoked when calling the pipeline for generation. diff --git a/examples/community/ddim_noise_comparative_analysis.py b/examples/community/ddim_noise_comparative_analysis.py index 829106c47f65..453f8525a552 100644 --- a/examples/community/ddim_noise_comparative_analysis.py +++ b/examples/community/ddim_noise_comparative_analysis.py @@ -110,7 +110,7 @@ def __call__( eta: float = 0.0, num_inference_steps: int = 50, use_clipped_model_output: Optional[bool] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, ) -> Union[ImagePipelineOutput, Tuple]: r""" diff --git a/examples/community/dps_pipeline.py b/examples/community/dps_pipeline.py index b29b06365bf3..6942d8ae608f 100755 --- a/examples/community/dps_pipeline.py +++ b/examples/community/dps_pipeline.py @@ -54,7 +54,7 @@ def __call__( batch_size: int = 1, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, num_inference_steps: int = 1000, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, zeta: float = 0.3, ) -> Union[ImagePipelineOutput, Tuple]: diff --git a/examples/community/edict_pipeline.py b/examples/community/edict_pipeline.py index a7bc892ddf93..bd96cdd6b2fc 100644 --- a/examples/community/edict_pipeline.py +++ b/examples/community/edict_pipeline.py @@ -1,5 +1,3 @@ -from typing import Optional - import torch from PIL import Image from tqdm.auto import tqdm @@ -39,7 +37,7 @@ def __init__( self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) def _encode_prompt( - self, prompt: str, negative_prompt: Optional[str] = None, do_classifier_free_guidance: bool = False + self, prompt: str, negative_prompt: str | None = None, do_classifier_free_guidance: bool = False ): text_inputs = self.tokenizer( prompt, @@ -141,7 +139,7 @@ def prepare_latents( text_embeds: torch.Tensor, timesteps: torch.Tensor, guidance_scale: float, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, ): do_classifier_free_guidance = guidance_scale > 1.0 @@ -194,9 +192,9 @@ def __call__( guidance_scale: float = 3.0, num_inference_steps: int = 50, strength: float = 0.8, - negative_prompt: Optional[str] = None, - generator: Optional[torch.Generator] = None, - output_type: Optional[str] = "pil", + negative_prompt: str | None = None, + generator: torch.Generator | None = None, + output_type: str | None = "pil", ): do_classifier_free_guidance = guidance_scale > 1.0 diff --git a/examples/community/fresco_v2v.py b/examples/community/fresco_v2v.py index b79834db5eed..974f71805baa 100644 --- a/examples/community/fresco_v2v.py +++ b/examples/community/fresco_v2v.py @@ -1208,7 +1208,7 @@ def apply_FRESCO_attn(pipe): def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -2064,7 +2064,7 @@ def __call__( negative_prompt_embeds: Optional[torch.FloatTensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, ip_adapter_image_embeds: Optional[List[torch.FloatTensor]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, controlnet_conditioning_scale: Union[float, List[float]] = 0.8, diff --git a/examples/community/gluegen.py b/examples/community/gluegen.py index 86813b63eca5..b98e0465561d 100644 --- a/examples/community/gluegen.py +++ b/examples/community/gluegen.py @@ -597,7 +597,7 @@ def __call__( latents: Optional[torch.Tensor] = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0, diff --git a/examples/community/hd_painter.py b/examples/community/hd_painter.py index 70e5656855ff..f412437b30ff 100644 --- a/examples/community/hd_painter.py +++ b/examples/community/hd_painter.py @@ -462,7 +462,7 @@ def __call__( num_inference_steps: int = 50, timesteps: List[int] = None, guidance_scale: float = 7.5, - positive_prompt: Optional[str] = "", + positive_prompt: str | None = "", negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.01, @@ -471,7 +471,7 @@ def __call__( prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, clip_skip: int = None, diff --git a/examples/community/iadb.py b/examples/community/iadb.py index 6262c3cb15fc..d61e2573b551 100644 --- a/examples/community/iadb.py +++ b/examples/community/iadb.py @@ -86,7 +86,7 @@ def __call__( batch_size: int = 1, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, num_inference_steps: int = 50, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, ) -> Union[ImagePipelineOutput, Tuple]: r""" diff --git a/examples/community/imagic_stable_diffusion.py b/examples/community/imagic_stable_diffusion.py index 091d0fbf8d3a..2efea594aff2 100644 --- a/examples/community/imagic_stable_diffusion.py +++ b/examples/community/imagic_stable_diffusion.py @@ -113,7 +113,7 @@ def train( image: Union[torch.Tensor, PIL.Image.Image], height: Optional[int] = 512, width: Optional[int] = 512, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, embedding_learning_rate: float = 0.001, diffusion_model_learning_rate: float = 2e-6, text_embedding_optimization_steps: int = 500, @@ -314,8 +314,8 @@ def __call__( height: Optional[int] = 512, width: Optional[int] = 512, num_inference_steps: Optional[int] = 50, - generator: Optional[torch.Generator] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | None = None, + output_type: str | None = "pil", return_dict: bool = True, guidance_scale: float = 7.5, eta: float = 0.0, diff --git a/examples/community/img2img_inpainting.py b/examples/community/img2img_inpainting.py index bef682425a2c..efc24a339d4f 100644 --- a/examples/community/img2img_inpainting.py +++ b/examples/community/img2img_inpainting.py @@ -143,9 +143,9 @@ def __call__( negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/instaflow_one_step.py b/examples/community/instaflow_one_step.py index 0f16707eadf5..d5c304b48d81 100644 --- a/examples/community/instaflow_one_step.py +++ b/examples/community/instaflow_one_step.py @@ -512,7 +512,7 @@ def __call__( latents: Optional[torch.Tensor] = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/interpolate_stable_diffusion.py b/examples/community/interpolate_stable_diffusion.py index 5b96c14d6367..448cf94216a8 100644 --- a/examples/community/interpolate_stable_diffusion.py +++ b/examples/community/interpolate_stable_diffusion.py @@ -131,9 +131,9 @@ def __call__( negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, @@ -401,8 +401,8 @@ def walk( prompts: List[str], seeds: List[int], num_interpolation_steps: Optional[int] = 6, - output_dir: Optional[str] = "./dreams", - name: Optional[str] = None, + output_dir: str | None = "./dreams", + name: str | None = None, batch_size: Optional[int] = 1, height: Optional[int] = 512, width: Optional[int] = 512, diff --git a/examples/community/ip_adapter_face_id.py b/examples/community/ip_adapter_face_id.py index d16aaf5a54c6..817d31e51af5 100644 --- a/examples/community/ip_adapter_face_id.py +++ b/examples/community/ip_adapter_face_id.py @@ -855,7 +855,7 @@ def __call__( prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, image_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0, diff --git a/examples/community/latent_consistency_img2img.py b/examples/community/latent_consistency_img2img.py index 0f5711f34b62..a3381dd9bd6f 100644 --- a/examples/community/latent_consistency_img2img.py +++ b/examples/community/latent_consistency_img2img.py @@ -286,7 +286,7 @@ def __call__( num_inference_steps: int = 4, lcm_origin_steps: int = 50, prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, ): diff --git a/examples/community/latent_consistency_interpolate.py b/examples/community/latent_consistency_interpolate.py index e8349ba317e9..ce87f61b13b0 100644 --- a/examples/community/latent_consistency_interpolate.py +++ b/examples/community/latent_consistency_interpolate.py @@ -669,7 +669,7 @@ def __call__( generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, latents: Optional[torch.Tensor] = None, prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, clip_skip: Optional[int] = None, diff --git a/examples/community/latent_consistency_txt2img.py b/examples/community/latent_consistency_txt2img.py index 0ce982065619..80dc7e2fc2c3 100755 --- a/examples/community/latent_consistency_txt2img.py +++ b/examples/community/latent_consistency_txt2img.py @@ -212,7 +212,7 @@ def __call__( num_inference_steps: int = 4, lcm_origin_steps: int = 50, prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, ): diff --git a/examples/community/llm_grounded_diffusion.py b/examples/community/llm_grounded_diffusion.py index 5bf6674a43e3..1ec565bac00e 100644 --- a/examples/community/llm_grounded_diffusion.py +++ b/examples/community/llm_grounded_diffusion.py @@ -769,7 +769,7 @@ def __call__( prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/lpw_stable_diffusion.py b/examples/community/lpw_stable_diffusion.py index 58e932bbcf74..3407c01e1edd 100644 --- a/examples/community/lpw_stable_diffusion.py +++ b/examples/community/lpw_stable_diffusion.py @@ -830,7 +830,7 @@ def __call__( prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, max_embeddings_multiples: Optional[int] = 3, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, is_cancelled_callback: Optional[Callable[[], bool]] = None, @@ -1091,7 +1091,7 @@ def text2img( prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, max_embeddings_multiples: Optional[int] = 3, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, is_cancelled_callback: Optional[Callable[[], bool]] = None, @@ -1209,7 +1209,7 @@ def img2img( prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, max_embeddings_multiples: Optional[int] = 3, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, is_cancelled_callback: Optional[Callable[[], bool]] = None, @@ -1323,7 +1323,7 @@ def inpaint( prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, max_embeddings_multiples: Optional[int] = 3, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, is_cancelled_callback: Optional[Callable[[], bool]] = None, diff --git a/examples/community/lpw_stable_diffusion_onnx.py b/examples/community/lpw_stable_diffusion_onnx.py index 92effc193329..eb7d86d09d3f 100644 --- a/examples/community/lpw_stable_diffusion_onnx.py +++ b/examples/community/lpw_stable_diffusion_onnx.py @@ -664,10 +664,10 @@ def __call__( strength: float = 0.8, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[np.ndarray] = None, max_embeddings_multiples: Optional[int] = 3, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, np.ndarray], None]] = None, is_cancelled_callback: Optional[Callable[[], bool]] = None, @@ -877,10 +877,10 @@ def text2img( guidance_scale: float = 7.5, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[np.ndarray] = None, max_embeddings_multiples: Optional[int] = 3, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, np.ndarray], None]] = None, callback_steps: int = 1, @@ -969,9 +969,9 @@ def img2img( guidance_scale: Optional[float] = 7.5, num_images_per_prompt: Optional[int] = 1, eta: Optional[float] = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, max_embeddings_multiples: Optional[int] = 3, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, np.ndarray], None]] = None, callback_steps: int = 1, @@ -1061,9 +1061,9 @@ def inpaint( guidance_scale: Optional[float] = 7.5, num_images_per_prompt: Optional[int] = 1, eta: Optional[float] = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, max_embeddings_multiples: Optional[int] = 3, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, np.ndarray], None]] = None, callback_steps: int = 1, diff --git a/examples/community/lpw_stable_diffusion_xl.py b/examples/community/lpw_stable_diffusion_xl.py index 1aba1f270198..b7a3623bf6ce 100644 --- a/examples/community/lpw_stable_diffusion_xl.py +++ b/examples/community/lpw_stable_diffusion_xl.py @@ -519,7 +519,7 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -724,12 +724,12 @@ def enable_model_cpu_offload(self, gpu_id=0): def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, + prompt_2: str | None = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, pooled_prompt_embeds: Optional[torch.Tensor] = None, @@ -1399,7 +1399,7 @@ def num_timesteps(self): def __call__( self, prompt: str = None, - prompt_2: Optional[str] = None, + prompt_2: str | None = None, image: Optional[PipelineImageInput] = None, mask_image: Optional[PipelineImageInput] = None, masked_image_latents: Optional[torch.Tensor] = None, @@ -1411,8 +1411,8 @@ def __call__( denoising_start: Optional[float] = None, denoising_end: Optional[float] = None, guidance_scale: float = 5.0, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, @@ -1422,7 +1422,7 @@ def __call__( negative_prompt_embeds: Optional[torch.Tensor] = None, pooled_prompt_embeds: Optional[torch.Tensor] = None, negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0, @@ -1955,7 +1955,7 @@ def denoising_value_valid(dnv): def text2img( self, prompt: str = None, - prompt_2: Optional[str] = None, + prompt_2: str | None = None, height: Optional[int] = None, width: Optional[int] = None, num_inference_steps: int = 50, @@ -1963,8 +1963,8 @@ def text2img( denoising_start: Optional[float] = None, denoising_end: Optional[float] = None, guidance_scale: float = 5.0, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, @@ -1974,7 +1974,7 @@ def text2img( negative_prompt_embeds: Optional[torch.Tensor] = None, pooled_prompt_embeds: Optional[torch.Tensor] = None, negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0, @@ -2028,7 +2028,7 @@ def text2img( def img2img( self, prompt: str = None, - prompt_2: Optional[str] = None, + prompt_2: str | None = None, image: Optional[PipelineImageInput] = None, height: Optional[int] = None, width: Optional[int] = None, @@ -2038,8 +2038,8 @@ def img2img( denoising_start: Optional[float] = None, denoising_end: Optional[float] = None, guidance_scale: float = 5.0, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, @@ -2049,7 +2049,7 @@ def img2img( negative_prompt_embeds: Optional[torch.Tensor] = None, pooled_prompt_embeds: Optional[torch.Tensor] = None, negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0, @@ -2105,7 +2105,7 @@ def img2img( def inpaint( self, prompt: str = None, - prompt_2: Optional[str] = None, + prompt_2: str | None = None, image: Optional[PipelineImageInput] = None, mask_image: Optional[PipelineImageInput] = None, masked_image_latents: Optional[torch.Tensor] = None, @@ -2117,8 +2117,8 @@ def inpaint( denoising_start: Optional[float] = None, denoising_end: Optional[float] = None, guidance_scale: float = 5.0, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, @@ -2128,7 +2128,7 @@ def inpaint( negative_prompt_embeds: Optional[torch.Tensor] = None, pooled_prompt_embeds: Optional[torch.Tensor] = None, negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0, diff --git a/examples/community/marigold_depth_estimation.py b/examples/community/marigold_depth_estimation.py index cd4473264e41..e1026cbafb06 100644 --- a/examples/community/marigold_depth_estimation.py +++ b/examples/community/marigold_depth_estimation.py @@ -43,7 +43,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") class MarigoldDepthOutput(BaseOutput): diff --git a/examples/community/masked_stable_diffusion_img2img.py b/examples/community/masked_stable_diffusion_img2img.py index 570bd0963e28..ac1612527d6a 100644 --- a/examples/community/masked_stable_diffusion_img2img.py +++ b/examples/community/masked_stable_diffusion_img2img.py @@ -32,7 +32,7 @@ def __call__( generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/masked_stable_diffusion_xl_img2img.py b/examples/community/masked_stable_diffusion_xl_img2img.py index 14d8c7c2da78..9e47b79e18be 100644 --- a/examples/community/masked_stable_diffusion_xl_img2img.py +++ b/examples/community/masked_stable_diffusion_xl_img2img.py @@ -59,7 +59,7 @@ def __call__( negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, ip_adapter_image_embeds: Optional[List[torch.FloatTensor]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0, diff --git a/examples/community/matryoshka.py b/examples/community/matryoshka.py index 09b1d1b24465..a4971f94829d 100644 --- a/examples/community/matryoshka.py +++ b/examples/community/matryoshka.py @@ -783,7 +783,7 @@ def __init__( norm_type: str = "layer_norm", num_attention_heads: int = 1, cross_attention_dim: int = 1280, - cross_attention_norm: Optional[str] = None, + cross_attention_norm: str | None = None, output_scale_factor: float = 1.0, downsample_padding: int = 1, add_downsample: bool = True, @@ -922,7 +922,7 @@ def __init__( num_attention_heads: int = 1, output_scale_factor: float = 1.0, cross_attention_dim: int = 1280, - cross_attention_norm: Optional[str] = None, + cross_attention_norm: str | None = None, dual_cross_attention: bool = False, use_linear_projection: bool = False, upcast_attention: bool = False, @@ -1055,7 +1055,7 @@ def __init__( norm_type: str = "layer_norm", num_attention_heads: int = 1, cross_attention_dim: int = 1280, - cross_attention_norm: Optional[str] = None, + cross_attention_norm: str | None = None, output_scale_factor: float = 1.0, add_upsample: bool = True, dual_cross_attention: bool = False, @@ -1617,10 +1617,10 @@ def get_down_block( attention_pre_only: bool = False, resnet_skip_time_act: bool = False, resnet_out_scale_factor: float = 1.0, - cross_attention_norm: Optional[str] = None, + cross_attention_norm: str | None = None, attention_head_dim: Optional[int] = None, use_attention_ffn: bool = True, - downsample_type: Optional[str] = None, + downsample_type: str | None = None, dropout: float = 0.0, ): # If attn head dim is not defined, we default it to the number of heads @@ -1695,7 +1695,7 @@ def get_mid_block( attention_type: str = "default", attention_pre_only: bool = False, resnet_skip_time_act: bool = False, - cross_attention_norm: Optional[str] = None, + cross_attention_norm: str | None = None, attention_head_dim: Optional[int] = 1, dropout: float = 0.0, ): @@ -1747,10 +1747,10 @@ def get_up_block( attention_pre_only: bool = False, resnet_skip_time_act: bool = False, resnet_out_scale_factor: float = 1.0, - cross_attention_norm: Optional[str] = None, + cross_attention_norm: str | None = None, attention_head_dim: Optional[int] = None, use_attention_ffn: bool = True, - upsample_type: Optional[str] = None, + upsample_type: str | None = None, dropout: float = 0.0, ) -> nn.Module: # If attn head dim is not defined, we default it to the number of heads @@ -1972,7 +1972,7 @@ def __init__( "CrossAttnDownBlock2D", "DownBlock2D", ), - mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn", + mid_block_type: str | None = "UNetMidBlock2DCrossAttn", up_block_types: Tuple[str, ...] = ( "UpBlock2D", "CrossAttnUpBlock2D", @@ -1993,14 +1993,14 @@ def __init__( transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1, reverse_transformer_layers_per_block: Optional[Tuple[Tuple[int]]] = None, encoder_hid_dim: Optional[int] = None, - encoder_hid_dim_type: Optional[str] = None, + encoder_hid_dim_type: str | None = None, attention_head_dim: Union[int, Tuple[int]] = 8, num_attention_heads: Optional[Union[int, Tuple[int]]] = None, dual_cross_attention: bool = False, use_attention_ffn: bool = True, use_linear_projection: bool = False, - class_embed_type: Optional[str] = None, - addition_embed_type: Optional[str] = None, + class_embed_type: str | None = None, + addition_embed_type: str | None = None, addition_time_embed_dim: Optional[int] = None, num_class_embeds: Optional[int] = None, upcast_attention: bool = False, @@ -2009,8 +2009,8 @@ def __init__( resnet_out_scale_factor: float = 1.0, time_embedding_type: str = "positional", time_embedding_dim: Optional[int] = None, - time_embedding_act_fn: Optional[str] = None, - timestep_post_act: Optional[str] = None, + time_embedding_act_fn: str | None = None, + timestep_post_act: str | None = None, time_cond_proj_dim: Optional[int] = None, conv_in_kernel: int = 3, conv_out_kernel: int = 3, @@ -2021,7 +2021,7 @@ def __init__( micro_conditioning_scale: int = None, class_embeddings_concat: bool = False, mid_block_only_cross_attention: Optional[bool] = None, - cross_attention_norm: Optional[str] = None, + cross_attention_norm: str | None = None, addition_embed_type_num_heads: int = 64, temporal_mode: bool = False, temporal_spatial_ds: bool = False, @@ -2384,7 +2384,7 @@ def _set_time_proj( def _set_encoder_hid_proj( self, - encoder_hid_dim_type: Optional[str], + encoder_hid_dim_type: str | None, cross_attention_dim: Union[int, Tuple[int]], encoder_hid_dim: Optional[int], ): @@ -2424,7 +2424,7 @@ def _set_encoder_hid_proj( def _set_class_embedding( self, - class_embed_type: Optional[str], + class_embed_type: str | None, act_fn: str, num_class_embeds: Optional[int], projection_class_embeddings_input_dim: Optional[int], @@ -2524,7 +2524,7 @@ def _set_pos_net_if_use_gligen(self, attention_type: str, cross_attention_dim: i ) @property - def attn_processors(self) -> Dict[str, AttentionProcessor]: + def attn_processors(self) -> dict[str, AttentionProcessor]: r""" Returns: `dict` of attention processors: A dictionary containing all attention processors used in the model with @@ -4264,7 +4264,7 @@ def __call__( negative_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0, diff --git a/examples/community/mixture_tiling_sdxl.py b/examples/community/mixture_tiling_sdxl.py index e09f5a25db73..dee1c3ef055e 100644 --- a/examples/community/mixture_tiling_sdxl.py +++ b/examples/community/mixture_tiling_sdxl.py @@ -388,12 +388,12 @@ class SeedTilesMode(Enum): def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, + prompt_2: str | None = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, pooled_prompt_embeds: Optional[torch.Tensor] = None, @@ -780,7 +780,7 @@ def __call__( num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, original_size: Optional[Tuple[int, int]] = None, diff --git a/examples/community/mod_controlnet_tile_sr_sdxl.py b/examples/community/mod_controlnet_tile_sr_sdxl.py index fe8bd73eabbe..959ef6453de0 100644 --- a/examples/community/mod_controlnet_tile_sr_sdxl.py +++ b/examples/community/mod_controlnet_tile_sr_sdxl.py @@ -243,7 +243,7 @@ def _tile2latent_indices( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -395,12 +395,12 @@ class TileWeightingMethod(Enum): def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, + prompt_2: str | None = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, pooled_prompt_embeds: Optional[torch.Tensor] = None, @@ -1265,7 +1265,7 @@ def __call__( eta: float = 0.0, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, controlnet_conditioning_scale: Union[float, List[float]] = 1.0, diff --git a/examples/community/multilingual_stable_diffusion.py b/examples/community/multilingual_stable_diffusion.py index 436803f201fe..444012acacec 100644 --- a/examples/community/multilingual_stable_diffusion.py +++ b/examples/community/multilingual_stable_diffusion.py @@ -146,9 +146,9 @@ def __call__( negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/pipeline_animatediff_controlnet.py b/examples/community/pipeline_animatediff_controlnet.py index 903bfd4fd57b..b195ba0f8a20 100644 --- a/examples/community/pipeline_animatediff_controlnet.py +++ b/examples/community/pipeline_animatediff_controlnet.py @@ -762,7 +762,7 @@ def __call__( ip_adapter_image: Optional[PipelineImageInput] = None, ip_adapter_image_embeds: Optional[PipelineImageInput] = None, conditioning_frames: Optional[List[PipelineImageInput]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, controlnet_conditioning_scale: Union[float, List[float]] = 1.0, diff --git a/examples/community/pipeline_animatediff_img2video.py b/examples/community/pipeline_animatediff_img2video.py index feba19f70fc6..5deadf6784c0 100644 --- a/examples/community/pipeline_animatediff_img2video.py +++ b/examples/community/pipeline_animatediff_img2video.py @@ -182,7 +182,7 @@ def tensor2vid(video: torch.Tensor, processor, output_type="np"): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -755,7 +755,7 @@ def __call__( negative_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, ip_adapter_image_embeds: Optional[PipelineImageInput] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: Optional[int] = 1, diff --git a/examples/community/pipeline_animatediff_ipex.py b/examples/community/pipeline_animatediff_ipex.py index 409ab9d6ad73..43d93bd3af2a 100644 --- a/examples/community/pipeline_animatediff_ipex.py +++ b/examples/community/pipeline_animatediff_ipex.py @@ -588,7 +588,7 @@ def __call__( negative_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, clip_skip: Optional[int] = None, @@ -844,7 +844,7 @@ def prepare_for_ipex( negative_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, clip_skip: Optional[int] = None, diff --git a/examples/community/pipeline_controlnet_xl_kolors.py b/examples/community/pipeline_controlnet_xl_kolors.py index 96c801bab379..d35176ba4f45 100644 --- a/examples/community/pipeline_controlnet_xl_kolors.py +++ b/examples/community/pipeline_controlnet_xl_kolors.py @@ -111,7 +111,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -800,7 +800,7 @@ def __call__( negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, controlnet_conditioning_scale: Union[float, List[float]] = 0.8, diff --git a/examples/community/pipeline_controlnet_xl_kolors_img2img.py b/examples/community/pipeline_controlnet_xl_kolors_img2img.py index 41f2550138b3..077241d9da22 100644 --- a/examples/community/pipeline_controlnet_xl_kolors_img2img.py +++ b/examples/community/pipeline_controlnet_xl_kolors_img2img.py @@ -131,7 +131,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -972,7 +972,7 @@ def __call__( negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, controlnet_conditioning_scale: Union[float, List[float]] = 0.8, diff --git a/examples/community/pipeline_controlnet_xl_kolors_inpaint.py b/examples/community/pipeline_controlnet_xl_kolors_inpaint.py index 07bd552a29e7..387a42ba22b3 100644 --- a/examples/community/pipeline_controlnet_xl_kolors_inpaint.py +++ b/examples/community/pipeline_controlnet_xl_kolors_inpaint.py @@ -118,7 +118,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -1139,7 +1139,7 @@ def __call__( negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, controlnet_conditioning_scale: Union[float, List[float]] = 0.8, diff --git a/examples/community/pipeline_demofusion_sdxl.py b/examples/community/pipeline_demofusion_sdxl.py index b5b68577982b..b0a3cab347a0 100644 --- a/examples/community/pipeline_demofusion_sdxl.py +++ b/examples/community/pipeline_demofusion_sdxl.py @@ -184,12 +184,12 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, + prompt_2: str | None = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, pooled_prompt_embeds: Optional[torch.Tensor] = None, @@ -637,7 +637,7 @@ def __call__( negative_prompt_embeds: Optional[torch.Tensor] = None, pooled_prompt_embeds: Optional[torch.Tensor] = None, negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = False, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/pipeline_fabric.py b/examples/community/pipeline_fabric.py index d29e98df5e50..fe7de1dd94a7 100644 --- a/examples/community/pipeline_fabric.py +++ b/examples/community/pipeline_fabric.py @@ -508,7 +508,7 @@ def __call__( num_images: int = 4, guidance_scale: float = 7.0, num_inference_steps: int = 20, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", feedback_start_ratio: float = 0.33, feedback_end_ratio: float = 0.66, min_weight: float = 0.05, diff --git a/examples/community/pipeline_faithdiff_stable_diffusion_xl.py b/examples/community/pipeline_faithdiff_stable_diffusion_xl.py index 4b564d5ee5c8..796a900b2f75 100644 --- a/examples/community/pipeline_faithdiff_stable_diffusion_xl.py +++ b/examples/community/pipeline_faithdiff_stable_diffusion_xl.py @@ -437,7 +437,7 @@ def __init__( "CrossAttnDownBlock2D", "DownBlock2D", ), - mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn", + mid_block_type: str | None = "UNetMidBlock2DCrossAttn", up_block_types: Tuple[str, ...] = ( "UpBlock2D", "CrossAttnUpBlock2D", @@ -457,13 +457,13 @@ def __init__( transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1, reverse_transformer_layers_per_block: Optional[Tuple[Tuple[int]]] = None, encoder_hid_dim: Optional[int] = None, - encoder_hid_dim_type: Optional[str] = None, + encoder_hid_dim_type: str | None = None, attention_head_dim: Union[int, Tuple[int]] = 8, num_attention_heads: Optional[Union[int, Tuple[int]]] = None, dual_cross_attention: bool = False, use_linear_projection: bool = False, - class_embed_type: Optional[str] = None, - addition_embed_type: Optional[str] = None, + class_embed_type: str | None = None, + addition_embed_type: str | None = None, addition_time_embed_dim: Optional[int] = None, num_class_embeds: Optional[int] = None, upcast_attention: bool = False, @@ -472,8 +472,8 @@ def __init__( resnet_out_scale_factor: float = 1.0, time_embedding_type: str = "positional", time_embedding_dim: Optional[int] = None, - time_embedding_act_fn: Optional[str] = None, - timestep_post_act: Optional[str] = None, + time_embedding_act_fn: str | None = None, + timestep_post_act: str | None = None, time_cond_proj_dim: Optional[int] = None, conv_in_kernel: int = 3, conv_out_kernel: int = 3, @@ -481,7 +481,7 @@ def __init__( attention_type: str = "default", class_embeddings_concat: bool = False, mid_block_only_cross_attention: Optional[bool] = None, - cross_attention_norm: Optional[str] = None, + cross_attention_norm: str | None = None, addition_embed_type_num_heads: int = 64, ): """Initialize the UnifiedUNet2DConditionModel.""" @@ -565,7 +565,7 @@ def init_extra_weights(self): self.agg_net = nn.ModuleList() def load_additional_layers( - self, dtype: Optional[torch.dtype] = torch.float16, channel: int = 512, weight_path: Optional[str] = None + self, dtype: Optional[torch.dtype] = torch.float16, channel: int = 512, weight_path: str | None = None ): """Load additional layers and weights from a file. @@ -1096,7 +1096,7 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): """Retrieve latents from an encoder output. @@ -1267,12 +1267,12 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, + prompt_2: str | None = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, prompt_embeds: Optional[torch.FloatTensor] = None, negative_prompt_embeds: Optional[torch.FloatTensor] = None, pooled_prompt_embeds: Optional[torch.FloatTensor] = None, @@ -1808,7 +1808,7 @@ def __call__( height: Optional[int] = None, width: Optional[int] = None, num_inference_steps: int = 50, - start_point: Optional[str] = "noise", + start_point: str | None = "noise", timesteps: List[int] = None, denoising_end: Optional[float] = None, overlap: float = 0.5, @@ -1823,7 +1823,7 @@ def __call__( negative_prompt_embeds: Optional[torch.FloatTensor] = None, pooled_prompt_embeds: Optional[torch.FloatTensor] = None, negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0, diff --git a/examples/community/pipeline_flux_differential_img2img.py b/examples/community/pipeline_flux_differential_img2img.py index 3677e73136f7..7cfb816d7310 100644 --- a/examples/community/pipeline_flux_differential_img2img.py +++ b/examples/community/pipeline_flux_differential_img2img.py @@ -97,7 +97,7 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -698,7 +698,7 @@ def __call__( latents: Optional[torch.FloatTensor] = None, prompt_embeds: Optional[torch.FloatTensor] = None, pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, joint_attention_kwargs: Optional[Dict[str, Any]] = None, callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, diff --git a/examples/community/pipeline_flux_kontext_multiple_images.py b/examples/community/pipeline_flux_kontext_multiple_images.py index 9e6ae427dbfa..fb344859e99e 100644 --- a/examples/community/pipeline_flux_kontext_multiple_images.py +++ b/examples/community/pipeline_flux_kontext_multiple_images.py @@ -188,7 +188,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -838,7 +838,7 @@ def __call__( negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, negative_prompt_embeds: Optional[torch.FloatTensor] = None, negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, joint_attention_kwargs: Optional[Dict[str, Any]] = None, callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, diff --git a/examples/community/pipeline_flux_rf_inversion.py b/examples/community/pipeline_flux_rf_inversion.py index 2cd6eb088cd8..16bb83bbcf34 100644 --- a/examples/community/pipeline_flux_rf_inversion.py +++ b/examples/community/pipeline_flux_rf_inversion.py @@ -685,7 +685,7 @@ def __call__( latents: Optional[torch.FloatTensor] = None, prompt_embeds: Optional[torch.FloatTensor] = None, pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, joint_attention_kwargs: Optional[Dict[str, Any]] = None, callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, diff --git a/examples/community/pipeline_flux_semantic_guidance.py b/examples/community/pipeline_flux_semantic_guidance.py index 74cd5c6981b0..a8d64b9caca6 100644 --- a/examples/community/pipeline_flux_semantic_guidance.py +++ b/examples/community/pipeline_flux_semantic_guidance.py @@ -802,7 +802,7 @@ def __call__( negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, negative_prompt_embeds: Optional[torch.FloatTensor] = None, negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, joint_attention_kwargs: Optional[Dict[str, Any]] = None, callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, diff --git a/examples/community/pipeline_flux_with_cfg.py b/examples/community/pipeline_flux_with_cfg.py index 5bc13f7e5e11..107901a76f57 100644 --- a/examples/community/pipeline_flux_with_cfg.py +++ b/examples/community/pipeline_flux_with_cfg.py @@ -622,7 +622,7 @@ def __call__( pooled_prompt_embeds: Optional[torch.FloatTensor] = None, negative_prompt_embeds: Optional[torch.FloatTensor] = None, negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, joint_attention_kwargs: Optional[Dict[str, Any]] = None, callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, diff --git a/examples/community/pipeline_hunyuandit_differential_img2img.py b/examples/community/pipeline_hunyuandit_differential_img2img.py index bc6841525b49..123e46e53bae 100644 --- a/examples/community/pipeline_hunyuandit_differential_img2img.py +++ b/examples/community/pipeline_hunyuandit_differential_img2img.py @@ -164,7 +164,7 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( encoder_output: torch.Tensor, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, sample_mode: str = "sample", ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": @@ -349,7 +349,7 @@ def encode_prompt( dtype: torch.dtype = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, + negative_prompt: str | None = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, prompt_attention_mask: Optional[torch.Tensor] = None, @@ -749,7 +749,7 @@ def __call__( prompt_attention_mask_2: Optional[torch.Tensor] = None, negative_prompt_attention_mask: Optional[torch.Tensor] = None, negative_prompt_attention_mask_2: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback_on_step_end: Optional[ Union[ diff --git a/examples/community/pipeline_kolors_differential_img2img.py b/examples/community/pipeline_kolors_differential_img2img.py index d456af8b3385..507f71ad61ef 100644 --- a/examples/community/pipeline_kolors_differential_img2img.py +++ b/examples/community/pipeline_kolors_differential_img2img.py @@ -67,7 +67,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -800,7 +800,7 @@ def __call__( negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, original_size: Optional[Tuple[int, int]] = None, diff --git a/examples/community/pipeline_kolors_inpainting.py b/examples/community/pipeline_kolors_inpainting.py index 85a3d698efa1..d5cb57580f58 100644 --- a/examples/community/pipeline_kolors_inpainting.py +++ b/examples/community/pipeline_kolors_inpainting.py @@ -239,7 +239,7 @@ def prepare_mask_and_masked_image(image, mask, height, width, return_image: bool # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -1100,7 +1100,7 @@ def __call__( negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0, diff --git a/examples/community/pipeline_prompt2prompt.py b/examples/community/pipeline_prompt2prompt.py index eb19667970b0..9ca7058892c4 100644 --- a/examples/community/pipeline_prompt2prompt.py +++ b/examples/community/pipeline_prompt2prompt.py @@ -571,7 +571,7 @@ def __call__( latents: Optional[torch.Tensor] = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: Optional[int] = 1, diff --git a/examples/community/pipeline_sdxl_style_aligned.py b/examples/community/pipeline_sdxl_style_aligned.py index ac7ad33b68c5..75f3cb9ead76 100644 --- a/examples/community/pipeline_sdxl_style_aligned.py +++ b/examples/community/pipeline_sdxl_style_aligned.py @@ -371,7 +371,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -508,12 +508,12 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, + prompt_2: str | None = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, pooled_prompt_embeds: Optional[torch.Tensor] = None, @@ -1394,7 +1394,7 @@ def __call__( pooled_prompt_embeds: Optional[torch.Tensor] = None, negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0, diff --git a/examples/community/pipeline_stable_diffusion_3_differential_img2img.py b/examples/community/pipeline_stable_diffusion_3_differential_img2img.py index 1803cf60cc4b..df5628e55eb6 100644 --- a/examples/community/pipeline_stable_diffusion_3_differential_img2img.py +++ b/examples/community/pipeline_stable_diffusion_3_differential_img2img.py @@ -68,7 +68,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -688,7 +688,7 @@ def __call__( negative_prompt_embeds: Optional[torch.FloatTensor] = None, pooled_prompt_embeds: Optional[torch.FloatTensor] = None, negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, clip_skip: Optional[int] = None, callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, diff --git a/examples/community/pipeline_stable_diffusion_3_instruct_pix2pix.py b/examples/community/pipeline_stable_diffusion_3_instruct_pix2pix.py index d9cee800e8ad..d4cb2924b9dc 100644 --- a/examples/community/pipeline_stable_diffusion_3_instruct_pix2pix.py +++ b/examples/community/pipeline_stable_diffusion_3_instruct_pix2pix.py @@ -102,7 +102,7 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -880,7 +880,7 @@ def __call__( negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, ip_adapter_image_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, joint_attention_kwargs: Optional[Dict[str, Any]] = None, clip_skip: Optional[int] = None, diff --git a/examples/community/pipeline_stable_diffusion_boxdiff.py b/examples/community/pipeline_stable_diffusion_boxdiff.py index c05a6143132c..69cff867ed6e 100644 --- a/examples/community/pipeline_stable_diffusion_boxdiff.py +++ b/examples/community/pipeline_stable_diffusion_boxdiff.py @@ -1341,7 +1341,7 @@ def __call__( prompt_embeds: Optional[torch.FloatTensor] = None, negative_prompt_embeds: Optional[torch.FloatTensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0, diff --git a/examples/community/pipeline_stable_diffusion_pag.py b/examples/community/pipeline_stable_diffusion_pag.py index 3f98dca0b986..f8061e11aa99 100644 --- a/examples/community/pipeline_stable_diffusion_pag.py +++ b/examples/community/pipeline_stable_diffusion_pag.py @@ -1114,7 +1114,7 @@ def __call__( negative_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0, diff --git a/examples/community/pipeline_stable_diffusion_upscale_ldm3d.py b/examples/community/pipeline_stable_diffusion_upscale_ldm3d.py index ba89556a5e4d..d537ef879711 100644 --- a/examples/community/pipeline_stable_diffusion_upscale_ldm3d.py +++ b/examples/community/pipeline_stable_diffusion_upscale_ldm3d.py @@ -523,7 +523,7 @@ def __call__( latents: Optional[torch.Tensor] = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/pipeline_stable_diffusion_xl_attentive_eraser.py b/examples/community/pipeline_stable_diffusion_xl_attentive_eraser.py index 822c804ec8a9..a001e458392c 100644 --- a/examples/community/pipeline_stable_diffusion_xl_attentive_eraser.py +++ b/examples/community/pipeline_stable_diffusion_xl_attentive_eraser.py @@ -435,7 +435,7 @@ def prepare_mask_and_masked_image(image, mask, height, width, return_image: bool # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -694,12 +694,12 @@ def prepare_ip_adapter_image_embeds( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, + prompt_2: str | None = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, prompt_embeds: Optional[torch.FloatTensor] = None, negative_prompt_embeds: Optional[torch.FloatTensor] = None, pooled_prompt_embeds: Optional[torch.FloatTensor] = None, @@ -1640,7 +1640,7 @@ def __call__( negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, ip_adapter_image_embeds: Optional[List[torch.FloatTensor]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0, diff --git a/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter.py b/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter.py index d435a6e146b2..e38801cd7647 100644 --- a/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter.py +++ b/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter.py @@ -241,12 +241,12 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, + prompt_2: str | None = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, pooled_prompt_embeds: Optional[torch.Tensor] = None, @@ -878,7 +878,7 @@ def __call__( negative_prompt_embeds: Optional[torch.Tensor] = None, pooled_prompt_embeds: Optional[torch.Tensor] = None, negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter_inpaint.py b/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter_inpaint.py index 8d70e8f3c2a1..2e05e3380316 100644 --- a/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter_inpaint.py +++ b/examples/community/pipeline_stable_diffusion_xl_controlnet_adapter_inpaint.py @@ -388,12 +388,12 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, + prompt_2: str | None = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, pooled_prompt_embeds: Optional[torch.Tensor] = None, @@ -1219,7 +1219,7 @@ def __call__( negative_prompt_embeds: Optional[torch.Tensor] = None, pooled_prompt_embeds: Optional[torch.Tensor] = None, negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/pipeline_stable_diffusion_xl_differential_img2img.py b/examples/community/pipeline_stable_diffusion_xl_differential_img2img.py index 99657facddf6..16144671d892 100644 --- a/examples/community/pipeline_stable_diffusion_xl_differential_img2img.py +++ b/examples/community/pipeline_stable_diffusion_xl_differential_img2img.py @@ -101,7 +101,7 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -268,12 +268,12 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, + prompt_2: str | None = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, pooled_prompt_embeds: Optional[torch.Tensor] = None, @@ -949,7 +949,7 @@ def __call__( negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0, diff --git a/examples/community/pipeline_stable_diffusion_xl_instandid_img2img.py b/examples/community/pipeline_stable_diffusion_xl_instandid_img2img.py index 33bf5ad346f4..1710f682d0ed 100644 --- a/examples/community/pipeline_stable_diffusion_xl_instandid_img2img.py +++ b/examples/community/pipeline_stable_diffusion_xl_instandid_img2img.py @@ -567,7 +567,7 @@ def __call__( pooled_prompt_embeds: Optional[torch.Tensor] = None, negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, image_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, controlnet_conditioning_scale: Union[float, List[float]] = 1.0, diff --git a/examples/community/pipeline_stable_diffusion_xl_instantid.py b/examples/community/pipeline_stable_diffusion_xl_instantid.py index 45876b91f7d8..4dfbcc194dd8 100644 --- a/examples/community/pipeline_stable_diffusion_xl_instantid.py +++ b/examples/community/pipeline_stable_diffusion_xl_instantid.py @@ -565,7 +565,7 @@ def __call__( pooled_prompt_embeds: Optional[torch.Tensor] = None, negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, image_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, controlnet_conditioning_scale: Union[float, List[float]] = 1.0, diff --git a/examples/community/pipeline_stable_diffusion_xl_ipex.py b/examples/community/pipeline_stable_diffusion_xl_ipex.py index 8e8357db3c7a..a0c07ccaeefb 100644 --- a/examples/community/pipeline_stable_diffusion_xl_ipex.py +++ b/examples/community/pipeline_stable_diffusion_xl_ipex.py @@ -268,12 +268,12 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, + prompt_2: str | None = None, device: Optional[torch.device] = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, pooled_prompt_embeds: Optional[torch.Tensor] = None, @@ -741,7 +741,7 @@ def __call__( pooled_prompt_embeds: Optional[torch.Tensor] = None, negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0, @@ -1181,7 +1181,7 @@ def prepare_for_ipex( pooled_prompt_embeds: Optional[torch.Tensor] = None, negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0, diff --git a/examples/community/pipeline_stg_hunyuan_video.py b/examples/community/pipeline_stg_hunyuan_video.py index 028d54d047e4..489125ff6a62 100644 --- a/examples/community/pipeline_stg_hunyuan_video.py +++ b/examples/community/pipeline_stg_hunyuan_video.py @@ -571,7 +571,7 @@ def __call__( prompt_embeds: Optional[torch.Tensor] = None, pooled_prompt_embeds: Optional[torch.Tensor] = None, prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, attention_kwargs: Optional[Dict[str, Any]] = None, callback_on_step_end: Optional[ diff --git a/examples/community/pipeline_stg_ltx.py b/examples/community/pipeline_stg_ltx.py index 70069a33f5d9..5bf98f97fa7d 100644 --- a/examples/community/pipeline_stg_ltx.py +++ b/examples/community/pipeline_stg_ltx.py @@ -506,7 +506,7 @@ def prepare_latents( num_frames: int = 161, dtype: Optional[torch.dtype] = None, device: Optional[torch.device] = None, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, ) -> torch.Tensor: if latents is not None: @@ -576,7 +576,7 @@ def __call__( negative_prompt_attention_mask: Optional[torch.Tensor] = None, decode_timestep: Union[float, List[float]] = 0.0, decode_noise_scale: Optional[Union[float, List[float]]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, attention_kwargs: Optional[Dict[str, Any]] = None, callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, diff --git a/examples/community/pipeline_stg_ltx_image2video.py b/examples/community/pipeline_stg_ltx_image2video.py index c32805e1419f..4eec632d7853 100644 --- a/examples/community/pipeline_stg_ltx_image2video.py +++ b/examples/community/pipeline_stg_ltx_image2video.py @@ -200,7 +200,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -535,7 +535,7 @@ def prepare_latents( num_frames: int = 161, dtype: Optional[torch.dtype] = None, device: Optional[torch.device] = None, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, ) -> torch.Tensor: height = height // self.vae_spatial_compression_ratio @@ -636,7 +636,7 @@ def __call__( negative_prompt_attention_mask: Optional[torch.Tensor] = None, decode_timestep: Union[float, List[float]] = 0.0, decode_noise_scale: Optional[Union[float, List[float]]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, attention_kwargs: Optional[Dict[str, Any]] = None, callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, diff --git a/examples/community/pipeline_stg_mochi.py b/examples/community/pipeline_stg_mochi.py index ad9317f6bc9d..350fba3bcf80 100644 --- a/examples/community/pipeline_stg_mochi.py +++ b/examples/community/pipeline_stg_mochi.py @@ -579,7 +579,7 @@ def __call__( prompt_attention_mask: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, attention_kwargs: Optional[Dict[str, Any]] = None, callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, diff --git a/examples/community/pipeline_stg_wan.py b/examples/community/pipeline_stg_wan.py index 39f208bad7c5..0833869bd35c 100644 --- a/examples/community/pipeline_stg_wan.py +++ b/examples/community/pipeline_stg_wan.py @@ -423,7 +423,7 @@ def __call__( latents: Optional[torch.Tensor] = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "np", + output_type: str | None = "np", return_dict: bool = True, attention_kwargs: Optional[Dict[str, Any]] = None, callback_on_step_end: Optional[ diff --git a/examples/community/pipeline_z_image_differential_img2img.py b/examples/community/pipeline_z_image_differential_img2img.py new file mode 100644 index 000000000000..6309e91abbdd --- /dev/null +++ b/examples/community/pipeline_z_image_differential_img2img.py @@ -0,0 +1,844 @@ +# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +from typing import Any, Callable, Dict, List, Optional, Union + +import torch +from transformers import AutoTokenizer, PreTrainedModel + +from diffusers.image_processor import PipelineImageInput, VaeImageProcessor +from diffusers.loaders import FromSingleFileMixin, ZImageLoraLoaderMixin +from diffusers.models.autoencoders import AutoencoderKL +from diffusers.models.transformers import ZImageTransformer2DModel +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from diffusers.pipelines.z_image.pipeline_output import ZImagePipelineOutput +from diffusers.schedulers import FlowMatchEulerDiscreteScheduler +from diffusers.utils import logging, replace_example_docstring +from diffusers.utils.torch_utils import randn_tensor + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from pipeline_z_image_differential_img2img import ZImageDifferentialImg2ImgPipeline + >>> from diffusers.utils import load_image + + >>> pipe = ZImageDifferentialImg2ImgPipeline.from_pretrained("Z-a-o/Z-Image-Turbo", torch_dtype=torch.bfloat16) + >>> pipe.to("cuda") + + >>> init_image = load_image( + >>> "https://github.com/exx8/differential-diffusion/blob/main/assets/input.jpg?raw=true", + >>> ) + + >>> mask = load_image( + >>> "https://github.com/exx8/differential-diffusion/blob/main/assets/map.jpg?raw=true", + >>> ) + + >>> prompt = "painting of a mountain landscape with a meadow and a forest, meadow background, anime countryside landscape, anime nature wallpap, anime landscape wallpaper, studio ghibli landscape, anime landscape, mountain behind meadow, anime background art, studio ghibli environment, background of flowery hill, anime beautiful peace scene, forrest background, anime scenery, landscape background, background art, anime scenery concept art" + + >>> image = pipe( + ... prompt, + ... image=init_image, + ... mask_image=mask, + ... strength=0.75, + ... num_inference_steps=9, + ... guidance_scale=0.0, + ... generator=torch.Generator("cuda").manual_seed(41), + ... ).images[0] + >>> image.save("image.png") + ``` +""" + + +# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents +def retrieve_latents( + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" +): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class ZImageDifferentialImg2ImgPipeline(DiffusionPipeline, ZImageLoraLoaderMixin, FromSingleFileMixin): + r""" + The ZImage pipeline for image-to-image generation. + + Args: + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + vae ([`AutoencoderKL`]): + Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. + text_encoder ([`PreTrainedModel`]): + A text encoder model to encode text prompts. + tokenizer ([`AutoTokenizer`]): + A tokenizer to tokenize text prompts. + transformer ([`ZImageTransformer2DModel`]): + A ZImage transformer model to denoise the encoded image latents. + """ + + model_cpu_offload_seq = "text_encoder->transformer->vae" + _optional_components = [] + _callback_tensor_inputs = ["latents", "prompt_embeds"] + + def __init__( + self, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKL, + text_encoder: PreTrainedModel, + tokenizer: AutoTokenizer, + transformer: ZImageTransformer2DModel, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + scheduler=scheduler, + transformer=transformer, + ) + self.vae_scale_factor = ( + 2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8 + ) + latent_channels = self.vae.config.latent_channels if getattr(self, "vae", None) else 16 + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) + + self.mask_processor = VaeImageProcessor( + vae_scale_factor=self.vae_scale_factor, + vae_latent_channels=latent_channels, + do_normalize=False, + do_binarize=False, + do_convert_grayscale=True, + ) + + # Copied from diffusers.pipelines.z_image.pipeline_z_image.ZImagePipeline.encode_prompt + def encode_prompt( + self, + prompt: Union[str, List[str]], + device: Optional[torch.device] = None, + do_classifier_free_guidance: bool = True, + negative_prompt: Optional[Union[str, List[str]]] = None, + prompt_embeds: Optional[List[torch.FloatTensor]] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + max_sequence_length: int = 512, + ): + prompt = [prompt] if isinstance(prompt, str) else prompt + prompt_embeds = self._encode_prompt( + prompt=prompt, + device=device, + prompt_embeds=prompt_embeds, + max_sequence_length=max_sequence_length, + ) + + if do_classifier_free_guidance: + if negative_prompt is None: + negative_prompt = ["" for _ in prompt] + else: + negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + assert len(prompt) == len(negative_prompt) + negative_prompt_embeds = self._encode_prompt( + prompt=negative_prompt, + device=device, + prompt_embeds=negative_prompt_embeds, + max_sequence_length=max_sequence_length, + ) + else: + negative_prompt_embeds = [] + return prompt_embeds, negative_prompt_embeds + + # Copied from diffusers.pipelines.z_image.pipeline_z_image.ZImagePipeline._encode_prompt + def _encode_prompt( + self, + prompt: Union[str, List[str]], + device: Optional[torch.device] = None, + prompt_embeds: Optional[List[torch.FloatTensor]] = None, + max_sequence_length: int = 512, + ) -> List[torch.FloatTensor]: + device = device or self._execution_device + + if prompt_embeds is not None: + return prompt_embeds + + if isinstance(prompt, str): + prompt = [prompt] + + for i, prompt_item in enumerate(prompt): + messages = [ + {"role": "user", "content": prompt_item}, + ] + prompt_item = self.tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=True, + ) + prompt[i] = prompt_item + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + return_tensors="pt", + ) + + text_input_ids = text_inputs.input_ids.to(device) + prompt_masks = text_inputs.attention_mask.to(device).bool() + + prompt_embeds = self.text_encoder( + input_ids=text_input_ids, + attention_mask=prompt_masks, + output_hidden_states=True, + ).hidden_states[-2] + + embeddings_list = [] + + for i in range(len(prompt_embeds)): + embeddings_list.append(prompt_embeds[i][prompt_masks[i]]) + + return embeddings_list + + # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3_img2img.StableDiffusion3Img2ImgPipeline.get_timesteps + def get_timesteps(self, num_inference_steps, strength, device): + # get the original timestep using init_timestep + init_timestep = min(num_inference_steps * strength, num_inference_steps) + + t_start = int(max(num_inference_steps - init_timestep, 0)) + timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :] + if hasattr(self.scheduler, "set_begin_index"): + self.scheduler.set_begin_index(t_start * self.scheduler.order) + + return timesteps, num_inference_steps - t_start + + @staticmethod + def _prepare_latent_image_ids(batch_size, height, width, device, dtype): + latent_image_ids = torch.zeros(height // 2, width // 2, 3) + latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height // 2)[:, None] + latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width // 2)[None, :] + + latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape + + latent_image_ids = latent_image_ids.reshape( + latent_image_id_height * latent_image_id_width, latent_image_id_channels + ) + + return latent_image_ids.to(device=device, dtype=dtype) + + def prepare_latents( + self, + image, + timestep, + batch_size, + num_channels_latents, + height, + width, + dtype, + device, + generator, + latents=None, + ): + height = 2 * (int(height) // (self.vae_scale_factor * 2)) + width = 2 * (int(width) // (self.vae_scale_factor * 2)) + + shape = (batch_size, num_channels_latents, height, width) + latent_image_ids = self._prepare_latent_image_ids(batch_size, height, width, device, dtype) + + if latents is not None: + return latents.to(device=device, dtype=dtype) + + # Encode the input image + image = image.to(device=device, dtype=dtype) + if image.shape[1] != num_channels_latents: + if isinstance(generator, list): + image_latents = [ + retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) + for i in range(image.shape[0]) + ] + image_latents = torch.cat(image_latents, dim=0) + else: + image_latents = retrieve_latents(self.vae.encode(image), generator=generator) + + # Apply scaling (inverse of decoding: decode does latents/scaling_factor + shift_factor) + image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor + else: + image_latents = image + + # Handle batch size expansion + if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0: + additional_image_per_prompt = batch_size // image_latents.shape[0] + image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0) + elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0: + raise ValueError( + f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts." + ) + + # Add noise using flow matching scale_noise + noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = self.scheduler.scale_noise(image_latents, timestep, noise) + + return latents, noise, image_latents, latent_image_ids + + def prepare_mask_latents( + self, + mask, + masked_image, + batch_size, + num_images_per_prompt, + height, + width, + dtype, + device, + generator, + ): + height = 2 * (int(height) // (self.vae_scale_factor * 2)) + width = 2 * (int(width) // (self.vae_scale_factor * 2)) + # resize the mask to latents shape as we concatenate the mask to the latents + # we do that before converting to dtype to avoid breaking in case we're using cpu_offload + # and half precision + mask = torch.nn.functional.interpolate(mask, size=(height, width)) + mask = mask.to(device=device, dtype=dtype) + + batch_size = batch_size * num_images_per_prompt + + masked_image = masked_image.to(device=device, dtype=dtype) + + if masked_image.shape[1] == 16: + masked_image_latents = masked_image + else: + masked_image_latents = retrieve_latents(self.vae.encode(masked_image), generator=generator) + + masked_image_latents = (masked_image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor + + # duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method + if mask.shape[0] < batch_size: + if not batch_size % mask.shape[0] == 0: + raise ValueError( + "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to" + f" a total batch size of {batch_size}, but {mask.shape[0]} masks were passed. Make sure the number" + " of masks that you pass is divisible by the total requested batch size." + ) + mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1) + if masked_image_latents.shape[0] < batch_size: + if not batch_size % masked_image_latents.shape[0] == 0: + raise ValueError( + "The passed images and the required batch size don't match. Images are supposed to be duplicated" + f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed." + " Make sure the number of images that you pass is divisible by the total requested batch size." + ) + masked_image_latents = masked_image_latents.repeat(batch_size // masked_image_latents.shape[0], 1, 1, 1) + + # aligning device to prevent device errors when concating it with the latent model input + masked_image_latents = masked_image_latents.to(device=device, dtype=dtype) + + return mask, masked_image_latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1 + + @property + def joint_attention_kwargs(self): + return self._joint_attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + image: PipelineImageInput = None, + mask_image: PipelineImageInput = None, + strength: float = 0.6, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 50, + sigmas: Optional[List[float]] = None, + guidance_scale: float = 5.0, + cfg_normalization: bool = False, + cfg_truncation: float = 1.0, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[List[torch.FloatTensor]] = None, + negative_prompt_embeds: Optional[List[torch.FloatTensor]] = None, + output_type: str | None = "pil", + return_dict: bool = True, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + max_sequence_length: int = 512, + ): + r""" + Function invoked when calling the pipeline for image-to-image generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both + numpy array and pytorch tensor, the expected value range is between `[0, 1]`. If it's a tensor or a + list of tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or + a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)`. + mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + `Image`, numpy array or tensor representing an image batch to mask `image`. Black pixels in the mask + are repainted while white pixels are preserved. If `mask_image` is a PIL image, it is converted to a + single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one + color channel (L) instead of 3, so the expected shape for pytorch tensor would be `(B, 1, H, W)`, `(B, + H, W)`, `(1, H, W)`, `(H, W)`. And for numpy array would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W, + 1)`, or `(H, W)`. + strength (`float`, *optional*, defaults to 0.6): + Indicates extent to transform the reference `image`. Must be between 0 and 1. `image` is used as a + starting point and more noise is added the higher the `strength`. The number of denoising steps depends + on the amount of noise initially added. When `strength` is 1, added noise is maximum and the denoising + process runs for the full number of iterations specified in `num_inference_steps`. A value of 1 + essentially ignores `image`. + height (`int`, *optional*, defaults to 1024): + The height in pixels of the generated image. If not provided, uses the input image height. + width (`int`, *optional*, defaults to 1024): + The width in pixels of the generated image. If not provided, uses the input image width. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + guidance_scale (`float`, *optional*, defaults to 5.0): + Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). + `guidance_scale` is defined as `w` of equation 2. of [Imagen + Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > + 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, + usually at the expense of lower image quality. + cfg_normalization (`bool`, *optional*, defaults to False): + Whether to apply configuration normalization. + cfg_truncation (`float`, *optional*, defaults to 1.0): + The truncation value for configuration. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + prompt_embeds (`List[torch.FloatTensor]`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`List[torch.FloatTensor]`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.stable_diffusion.ZImagePipelineOutput`] instead of a plain + tuple. + joint_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int`, *optional*, defaults to 512): + Maximum sequence length to use with the `prompt`. + + Examples: + + Returns: + [`~pipelines.z_image.ZImagePipelineOutput`] or `tuple`: [`~pipelines.z_image.ZImagePipelineOutput`] if + `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the + generated images. + """ + # 1. Check inputs and validate strength + if strength < 0 or strength > 1: + raise ValueError(f"The value of strength should be in [0.0, 1.0] but is {strength}") + + # 2. Preprocess image + init_image = self.image_processor.preprocess(image) + init_image = init_image.to(dtype=torch.float32) + + # Get dimensions from the preprocessed image if not specified + if height is None: + height = init_image.shape[-2] + if width is None: + width = init_image.shape[-1] + + vae_scale = self.vae_scale_factor * 2 + if height % vae_scale != 0: + raise ValueError( + f"Height must be divisible by {vae_scale} (got {height}). " + f"Please adjust the height to a multiple of {vae_scale}." + ) + if width % vae_scale != 0: + raise ValueError( + f"Width must be divisible by {vae_scale} (got {width}). " + f"Please adjust the width to a multiple of {vae_scale}." + ) + + device = self._execution_device + + self._guidance_scale = guidance_scale + self._joint_attention_kwargs = joint_attention_kwargs + self._interrupt = False + self._cfg_normalization = cfg_normalization + self._cfg_truncation = cfg_truncation + + # 3. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = len(prompt_embeds) + + # If prompt_embeds is provided and prompt is None, skip encoding + if prompt_embeds is not None and prompt is None: + if self.do_classifier_free_guidance and negative_prompt_embeds is None: + raise ValueError( + "When `prompt_embeds` is provided without `prompt`, " + "`negative_prompt_embeds` must also be provided for classifier-free guidance." + ) + else: + ( + prompt_embeds, + negative_prompt_embeds, + ) = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + device=device, + max_sequence_length=max_sequence_length, + ) + + # 4. Prepare latent variables + num_channels_latents = self.transformer.in_channels + + # Repeat prompt_embeds for num_images_per_prompt + if num_images_per_prompt > 1: + prompt_embeds = [pe for pe in prompt_embeds for _ in range(num_images_per_prompt)] + if self.do_classifier_free_guidance and negative_prompt_embeds: + negative_prompt_embeds = [npe for npe in negative_prompt_embeds for _ in range(num_images_per_prompt)] + + actual_batch_size = batch_size * num_images_per_prompt + + # Calculate latent dimensions for image_seq_len + latent_height = 2 * (int(height) // (self.vae_scale_factor * 2)) + latent_width = 2 * (int(width) // (self.vae_scale_factor * 2)) + image_seq_len = (latent_height // 2) * (latent_width // 2) + + # 5. Prepare timesteps + mu = calculate_shift( + image_seq_len, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.5), + self.scheduler.config.get("max_shift", 1.15), + ) + self.scheduler.sigma_min = 0.0 + scheduler_kwargs = {"mu": mu} + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + sigmas=sigmas, + **scheduler_kwargs, + ) + + # 6. Adjust timesteps based on strength + timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) + if num_inference_steps < 1: + raise ValueError( + f"After adjusting the num_inference_steps by strength parameter: {strength}, the number of pipeline " + f"steps is {num_inference_steps} which is < 1 and not appropriate for this pipeline." + ) + latent_timestep = timesteps[:1].repeat(actual_batch_size) + + # 7. Prepare latents from image + latents, noise, original_image_latents, latent_image_ids = self.prepare_latents( + init_image, + latent_timestep, + actual_batch_size, + num_channels_latents, + height, + width, + prompt_embeds[0].dtype, + device, + generator, + latents, + ) + resize_mode = "default" + crops_coords = None + + # start diff diff preparation + original_mask = self.mask_processor.preprocess( + mask_image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords + ) + + masked_image = init_image * original_mask + original_mask, _ = self.prepare_mask_latents( + original_mask, + masked_image, + batch_size, + num_images_per_prompt, + height, + width, + prompt_embeds[0].dtype, + device, + generator, + ) + mask_thresholds = torch.arange(num_inference_steps, dtype=original_mask.dtype) / num_inference_steps + mask_thresholds = mask_thresholds.reshape(-1, 1, 1, 1).to(device) + masks = original_mask > mask_thresholds + # end diff diff preparation + + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + # 8. Denoising loop + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latents.shape[0]) + timestep = (1000 - timestep) / 1000 + # Normalized time for time-aware config (0 at start, 1 at end) + t_norm = timestep[0].item() + + # Handle cfg truncation + current_guidance_scale = self.guidance_scale + if ( + self.do_classifier_free_guidance + and self._cfg_truncation is not None + and float(self._cfg_truncation) <= 1 + ): + if t_norm > self._cfg_truncation: + current_guidance_scale = 0.0 + + # Run CFG only if configured AND scale is non-zero + apply_cfg = self.do_classifier_free_guidance and current_guidance_scale > 0 + + if apply_cfg: + latents_typed = latents.to(self.transformer.dtype) + latent_model_input = latents_typed.repeat(2, 1, 1, 1) + prompt_embeds_model_input = prompt_embeds + negative_prompt_embeds + timestep_model_input = timestep.repeat(2) + else: + latent_model_input = latents.to(self.transformer.dtype) + prompt_embeds_model_input = prompt_embeds + timestep_model_input = timestep + + latent_model_input = latent_model_input.unsqueeze(2) + latent_model_input_list = list(latent_model_input.unbind(dim=0)) + + model_out_list = self.transformer( + latent_model_input_list, + timestep_model_input, + prompt_embeds_model_input, + )[0] + + if apply_cfg: + # Perform CFG + pos_out = model_out_list[:actual_batch_size] + neg_out = model_out_list[actual_batch_size:] + + noise_pred = [] + for j in range(actual_batch_size): + pos = pos_out[j].float() + neg = neg_out[j].float() + + pred = pos + current_guidance_scale * (pos - neg) + + # Renormalization + if self._cfg_normalization and float(self._cfg_normalization) > 0.0: + ori_pos_norm = torch.linalg.vector_norm(pos) + new_pos_norm = torch.linalg.vector_norm(pred) + max_new_norm = ori_pos_norm * float(self._cfg_normalization) + if new_pos_norm > max_new_norm: + pred = pred * (max_new_norm / new_pos_norm) + + noise_pred.append(pred) + + noise_pred = torch.stack(noise_pred, dim=0) + else: + noise_pred = torch.stack([t.float() for t in model_out_list], dim=0) + + noise_pred = noise_pred.squeeze(2) + noise_pred = -noise_pred + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred.to(torch.float32), t, latents, return_dict=False)[0] + assert latents.dtype == torch.float32 + + # start diff diff + image_latent = original_image_latents + latents_dtype = latents.dtype + if i < len(timesteps) - 1: + noise_timestep = timesteps[i + 1] + image_latent = self.scheduler.scale_noise( + original_image_latents, torch.tensor([noise_timestep]), noise + ) + + mask = masks[i].to(latents_dtype) + latents = image_latent * mask + latents * (1 - mask) + # end diff diff + + if latents.dtype != latents_dtype: + if torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 + latents = latents.to(latents_dtype) + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if output_type == "latent": + image = latents + + else: + latents = latents.to(self.vae.dtype) + latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor + + image = self.vae.decode(latents, return_dict=False)[0] + image = self.image_processor.postprocess(image, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return ZImagePipelineOutput(images=image) diff --git a/examples/community/pipeline_zero1to3.py b/examples/community/pipeline_zero1to3.py index 1be59fd8323e..125566d7cf2d 100644 --- a/examples/community/pipeline_zero1to3.py +++ b/examples/community/pipeline_zero1to3.py @@ -614,7 +614,7 @@ def __call__( latents: Optional[torch.Tensor] = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/pipline_flux_fill_controlnet_Inpaint.py b/examples/community/pipline_flux_fill_controlnet_Inpaint.py index cc642a767f87..88dcbecf8de0 100644 --- a/examples/community/pipline_flux_fill_controlnet_Inpaint.py +++ b/examples/community/pipline_flux_fill_controlnet_Inpaint.py @@ -99,7 +99,7 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -112,7 +112,7 @@ def retrieve_latents( def retrieve_latents_fill( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -855,7 +855,7 @@ def __call__( latents: Optional[torch.FloatTensor] = None, prompt_embeds: Optional[torch.FloatTensor] = None, pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, joint_attention_kwargs: Optional[Dict[str, Any]] = None, callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, diff --git a/examples/community/regional_prompting_stable_diffusion.py b/examples/community/regional_prompting_stable_diffusion.py index 3bc780cfcf7a..daea47b3c3dd 100644 --- a/examples/community/regional_prompting_stable_diffusion.py +++ b/examples/community/regional_prompting_stable_diffusion.py @@ -148,9 +148,9 @@ def __call__( negative_prompt: str = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, rp_args: Dict[str, str] = None, ): @@ -800,7 +800,7 @@ def stable_diffusion_call( negative_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0, diff --git a/examples/community/rerender_a_video.py b/examples/community/rerender_a_video.py index 840f9e206d4d..b2680ebb2d4f 100644 --- a/examples/community/rerender_a_video.py +++ b/examples/community/rerender_a_video.py @@ -607,7 +607,7 @@ def __call__( latents: Optional[torch.Tensor] = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/run_onnx_controlnet.py b/examples/community/run_onnx_controlnet.py index 2b56e8a1e51b..2c631ae37ed7 100644 --- a/examples/community/run_onnx_controlnet.py +++ b/examples/community/run_onnx_controlnet.py @@ -136,7 +136,7 @@ def _encode_prompt( prompt: Union[str, List[str]], num_images_per_prompt: Optional[int], do_classifier_free_guidance: bool, - negative_prompt: Optional[str], + negative_prompt: str | None, prompt_embeds: Optional[np.ndarray] = None, negative_prompt_embeds: Optional[np.ndarray] = None, ): @@ -534,7 +534,7 @@ def __call__( latents: Optional[torch.Tensor] = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/run_tensorrt_controlnet.py b/examples/community/run_tensorrt_controlnet.py index b62eb4f58e8f..23c13e6a501c 100644 --- a/examples/community/run_tensorrt_controlnet.py +++ b/examples/community/run_tensorrt_controlnet.py @@ -240,7 +240,7 @@ def _encode_prompt( prompt: Union[str, List[str]], num_images_per_prompt: Optional[int], do_classifier_free_guidance: bool, - negative_prompt: Optional[str], + negative_prompt: str | None, prompt_embeds: Optional[np.ndarray] = None, negative_prompt_embeds: Optional[np.ndarray] = None, ): @@ -638,7 +638,7 @@ def __call__( latents: Optional[torch.Tensor] = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/scheduling_ufogen.py b/examples/community/scheduling_ufogen.py index fada2636e98d..2d109ef36909 100644 --- a/examples/community/scheduling_ufogen.py +++ b/examples/community/scheduling_ufogen.py @@ -377,7 +377,7 @@ def step( model_output: torch.Tensor, timestep: int, sample: torch.Tensor, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, return_dict: bool = True, ) -> Union[UFOGenSchedulerOutput, Tuple]: """ diff --git a/examples/community/sd_text2img_k_diffusion.py b/examples/community/sd_text2img_k_diffusion.py index e351420f786a..690182714cde 100755 --- a/examples/community/sd_text2img_k_diffusion.py +++ b/examples/community/sd_text2img_k_diffusion.py @@ -285,9 +285,9 @@ def __call__( negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/sde_drag.py b/examples/community/sde_drag.py index 63899ce738bb..0b654730fcc4 100644 --- a/examples/community/sde_drag.py +++ b/examples/community/sde_drag.py @@ -73,7 +73,7 @@ def __call__( image_scale: Optional[float] = 0.3, adapt_radius: Optional[int] = 5, min_lora_scale: Optional[float] = 0.5, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, ): r""" Function invoked when calling the pipeline for image editing. diff --git a/examples/community/seed_resize_stable_diffusion.py b/examples/community/seed_resize_stable_diffusion.py index eafe7572aab5..373095fa2ca2 100644 --- a/examples/community/seed_resize_stable_diffusion.py +++ b/examples/community/seed_resize_stable_diffusion.py @@ -80,9 +80,9 @@ def __call__( negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/speech_to_image_diffusion.py b/examples/community/speech_to_image_diffusion.py index a8ec1620a2eb..695dca2465b4 100644 --- a/examples/community/speech_to_image_diffusion.py +++ b/examples/community/speech_to_image_diffusion.py @@ -75,9 +75,9 @@ def __call__( negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/stable_diffusion_comparison.py b/examples/community/stable_diffusion_comparison.py index ce6e77c87ff0..8ce9c39227e7 100644 --- a/examples/community/stable_diffusion_comparison.py +++ b/examples/community/stable_diffusion_comparison.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable, List, Optional, Union import torch from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer @@ -81,7 +81,7 @@ def __init__( self.register_modules(pipeline1=self.pipe1, pipeline2=self.pipe2, pipeline3=self.pipe3, pipeline4=self.pipe4) @property - def layers(self) -> Dict[str, Any]: + def layers(self) -> dict[str, Any]: return {k: getattr(self, k) for k in self.config.keys() if not k.startswith("_")} @torch.no_grad() @@ -95,9 +95,9 @@ def text2img_sd1_1( negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, @@ -132,9 +132,9 @@ def text2img_sd1_2( negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, @@ -169,9 +169,9 @@ def text2img_sd1_3( negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, @@ -206,9 +206,9 @@ def text2img_sd1_4( negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, @@ -243,9 +243,9 @@ def _call_( negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/stable_diffusion_controlnet_img2img.py b/examples/community/stable_diffusion_controlnet_img2img.py index aa116112be1c..03c6fe7f6466 100644 --- a/examples/community/stable_diffusion_controlnet_img2img.py +++ b/examples/community/stable_diffusion_controlnet_img2img.py @@ -614,7 +614,7 @@ def __call__( latents: Optional[torch.Tensor] = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/stable_diffusion_controlnet_inpaint.py b/examples/community/stable_diffusion_controlnet_inpaint.py index 6d710e0d73c7..9b76faf56a8a 100644 --- a/examples/community/stable_diffusion_controlnet_inpaint.py +++ b/examples/community/stable_diffusion_controlnet_inpaint.py @@ -757,7 +757,7 @@ def __call__( latents: Optional[torch.Tensor] = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/stable_diffusion_controlnet_inpaint_img2img.py b/examples/community/stable_diffusion_controlnet_inpaint_img2img.py index fcb5ed059bb4..299dad58707d 100644 --- a/examples/community/stable_diffusion_controlnet_inpaint_img2img.py +++ b/examples/community/stable_diffusion_controlnet_inpaint_img2img.py @@ -745,7 +745,7 @@ def __call__( latents: Optional[torch.Tensor] = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/stable_diffusion_controlnet_reference.py b/examples/community/stable_diffusion_controlnet_reference.py index 74c81b63626d..18c79a0853f9 100644 --- a/examples/community/stable_diffusion_controlnet_reference.py +++ b/examples/community/stable_diffusion_controlnet_reference.py @@ -119,7 +119,7 @@ def __call__( latents: Optional[torch.Tensor] = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/stable_diffusion_ipex.py b/examples/community/stable_diffusion_ipex.py index 4f545aa09ded..ddc39b5b7b1a 100644 --- a/examples/community/stable_diffusion_ipex.py +++ b/examples/community/stable_diffusion_ipex.py @@ -574,7 +574,7 @@ def __call__( latents: Optional[torch.Tensor] = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/stable_diffusion_mega.py b/examples/community/stable_diffusion_mega.py index c67ebc80b006..a22ecb209ada 100644 --- a/examples/community/stable_diffusion_mega.py +++ b/examples/community/stable_diffusion_mega.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable, List, Optional, Union import PIL.Image import torch @@ -92,7 +92,7 @@ def __init__( self.register_to_config(requires_safety_checker=requires_safety_checker) @property - def components(self) -> Dict[str, Any]: + def components(self) -> dict[str, Any]: return {k: getattr(self, k) for k in self.config.keys() if not k.startswith("_")} @torch.no_grad() @@ -107,8 +107,8 @@ def inpaint( negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: Optional[float] = 0.0, - generator: Optional[torch.Generator] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | None = None, + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, @@ -141,8 +141,8 @@ def img2img( negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: Optional[float] = 0.0, - generator: Optional[torch.Generator] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | None = None, + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, @@ -176,9 +176,9 @@ def text2img( negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/stable_diffusion_reference.py b/examples/community/stable_diffusion_reference.py index d0372bbeba65..7a32827a9350 100644 --- a/examples/community/stable_diffusion_reference.py +++ b/examples/community/stable_diffusion_reference.py @@ -276,7 +276,7 @@ def check_inputs( height: int, width: int, callback_steps: Optional[int], - negative_prompt: Optional[str] = None, + negative_prompt: str | None = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[torch.Tensor] = None, @@ -291,7 +291,7 @@ def check_inputs( height (int): The height of the input image. width (int): The width of the input image. callback_steps (Optional[int]): The number of steps to perform the callback on. - negative_prompt (Optional[str]): The negative prompt text. + negative_prompt (str | None): The negative prompt text. prompt_embeds (Optional[torch.Tensor]): The prompt embeddings. negative_prompt_embeds (Optional[torch.Tensor]): The negative prompt embeddings. ip_adapter_image (Optional[torch.Tensor]): The input adapter image. @@ -411,11 +411,11 @@ def _encode_prompt( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_prompt def encode_prompt( self, - prompt: Optional[str], + prompt: str | None, device: torch.device, num_images_per_prompt: int, do_classifier_free_guidance: bool, - negative_prompt: Optional[str] = None, + negative_prompt: str | None = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, lora_scale: Optional[float] = None, @@ -642,7 +642,7 @@ def prepare_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs def prepare_extra_step_kwargs( self, generator: Union[torch.Generator, List[torch.Generator]], eta: float - ) -> Dict[str, Any]: + ) -> dict[str, Any]: r""" Prepare extra keyword arguments for the scheduler step. @@ -834,7 +834,7 @@ def __call__( latents: Optional[torch.Tensor] = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/stable_diffusion_repaint.py b/examples/community/stable_diffusion_repaint.py index b974e3c7ae74..7248e4101206 100644 --- a/examples/community/stable_diffusion_repaint.py +++ b/examples/community/stable_diffusion_repaint.py @@ -625,7 +625,7 @@ def __call__( latents: Optional[torch.Tensor] = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/stable_diffusion_xl_controlnet_reference.py b/examples/community/stable_diffusion_xl_controlnet_reference.py index 421e67f5bba6..a458ee7c6506 100644 --- a/examples/community/stable_diffusion_xl_controlnet_reference.py +++ b/examples/community/stable_diffusion_xl_controlnet_reference.py @@ -345,7 +345,7 @@ def __call__( negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, controlnet_conditioning_scale: Union[float, List[float]] = 1.0, diff --git a/examples/community/stable_diffusion_xl_reference.py b/examples/community/stable_diffusion_xl_reference.py index eb055574966d..c6be397144c5 100644 --- a/examples/community/stable_diffusion_xl_reference.py +++ b/examples/community/stable_diffusion_xl_reference.py @@ -290,7 +290,7 @@ def __call__( negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0, diff --git a/examples/community/stable_unclip.py b/examples/community/stable_unclip.py index f13c4e0a490b..22c3b8c6fe9c 100644 --- a/examples/community/stable_unclip.py +++ b/examples/community/stable_unclip.py @@ -190,7 +190,7 @@ def __call__( width: Optional[int] = None, num_images_per_prompt: int = 1, prior_num_inference_steps: int = 25, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, prior_latents: Optional[torch.Tensor] = None, text_model_output: Optional[Union[CLIPTextModelOutput, Tuple]] = None, text_attention_mask: Optional[torch.Tensor] = None, @@ -199,7 +199,7 @@ def __call__( decoder_num_inference_steps: int = 50, decoder_num_images_per_prompt: Optional[int] = 1, decoder_eta: float = 0.0, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, ): if prompt is not None: diff --git a/examples/community/text_inpainting.py b/examples/community/text_inpainting.py index bdf9eca498e7..6df957089883 100644 --- a/examples/community/text_inpainting.py +++ b/examples/community/text_inpainting.py @@ -134,9 +134,9 @@ def __call__( negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/tiled_upscaling.py b/examples/community/tiled_upscaling.py index 7a5e77155cd0..5fe67dce43d8 100644 --- a/examples/community/tiled_upscaling.py +++ b/examples/community/tiled_upscaling.py @@ -192,7 +192,7 @@ def __call__( negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/community/unclip_image_interpolation.py b/examples/community/unclip_image_interpolation.py index 65b52578601e..54e3a9813908 100644 --- a/examples/community/unclip_image_interpolation.py +++ b/examples/community/unclip_image_interpolation.py @@ -216,7 +216,7 @@ def __call__( decoder_latents: Optional[torch.Tensor] = None, super_res_latents: Optional[torch.Tensor] = None, decoder_guidance_scale: float = 8.0, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, ): """ diff --git a/examples/community/unclip_text_interpolation.py b/examples/community/unclip_text_interpolation.py index 6fd4f348f48d..e2478811a059 100644 --- a/examples/community/unclip_text_interpolation.py +++ b/examples/community/unclip_text_interpolation.py @@ -225,7 +225,7 @@ def __call__( decoder_guidance_scale: float = 8.0, enable_sequential_cpu_offload=True, gpu_id=0, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, ): """ diff --git a/examples/community/wildcard_stable_diffusion.py b/examples/community/wildcard_stable_diffusion.py index d40221e5b1cf..342913b1508e 100644 --- a/examples/community/wildcard_stable_diffusion.py +++ b/examples/community/wildcard_stable_diffusion.py @@ -165,9 +165,9 @@ def __call__( negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/consistency_distillation/train_lcm_distill_lora_sd_wds.py b/examples/consistency_distillation/train_lcm_distill_lora_sd_wds.py index 26a3ecc87935..38885d4bdf11 100644 --- a/examples/consistency_distillation/train_lcm_distill_lora_sd_wds.py +++ b/examples/consistency_distillation/train_lcm_distill_lora_sd_wds.py @@ -74,7 +74,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py b/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py index ef50e8eb2da4..4dd7cbb60ce1 100644 --- a/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py +++ b/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py @@ -67,7 +67,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/consistency_distillation/train_lcm_distill_lora_sdxl_wds.py b/examples/consistency_distillation/train_lcm_distill_lora_sdxl_wds.py index a3302d7147b9..f4eb70e61e0f 100644 --- a/examples/consistency_distillation/train_lcm_distill_lora_sdxl_wds.py +++ b/examples/consistency_distillation/train_lcm_distill_lora_sdxl_wds.py @@ -80,7 +80,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/consistency_distillation/train_lcm_distill_sd_wds.py b/examples/consistency_distillation/train_lcm_distill_sd_wds.py index 79bc706bcca3..ef1c57bb9e18 100644 --- a/examples/consistency_distillation/train_lcm_distill_sd_wds.py +++ b/examples/consistency_distillation/train_lcm_distill_sd_wds.py @@ -73,7 +73,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/consistency_distillation/train_lcm_distill_sdxl_wds.py b/examples/consistency_distillation/train_lcm_distill_sdxl_wds.py index d6b2dd895766..6f6fcdfa286a 100644 --- a/examples/consistency_distillation/train_lcm_distill_sdxl_wds.py +++ b/examples/consistency_distillation/train_lcm_distill_sdxl_wds.py @@ -79,7 +79,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/controlnet/train_controlnet.py b/examples/controlnet/train_controlnet.py index 198501da725e..690325e24eb8 100644 --- a/examples/controlnet/train_controlnet.py +++ b/examples/controlnet/train_controlnet.py @@ -61,7 +61,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/controlnet/train_controlnet_flax.py b/examples/controlnet/train_controlnet_flax.py index 588f6b1f4ca0..4d60598104ba 100644 --- a/examples/controlnet/train_controlnet_flax.py +++ b/examples/controlnet/train_controlnet_flax.py @@ -61,7 +61,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = logging.getLogger(__name__) diff --git a/examples/controlnet/train_controlnet_flux.py b/examples/controlnet/train_controlnet_flux.py index 5d54e34eaa06..70355870e9e8 100644 --- a/examples/controlnet/train_controlnet_flux.py +++ b/examples/controlnet/train_controlnet_flux.py @@ -66,7 +66,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) if is_torch_npu_available(): diff --git a/examples/controlnet/train_controlnet_sd3.py b/examples/controlnet/train_controlnet_sd3.py index 1d130a38c97e..66f2bc2eadce 100644 --- a/examples/controlnet/train_controlnet_sd3.py +++ b/examples/controlnet/train_controlnet_sd3.py @@ -63,7 +63,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/controlnet/train_controlnet_sdxl.py b/examples/controlnet/train_controlnet_sdxl.py index b853a32c4483..62757c7f6eb2 100644 --- a/examples/controlnet/train_controlnet_sdxl.py +++ b/examples/controlnet/train_controlnet_sdxl.py @@ -62,7 +62,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) if is_torch_npu_available(): diff --git a/examples/custom_diffusion/test_custom_diffusion.py b/examples/custom_diffusion/test_custom_diffusion.py index 9af84ec7598f..ad18eb246777 100644 --- a/examples/custom_diffusion/test_custom_diffusion.py +++ b/examples/custom_diffusion/test_custom_diffusion.py @@ -17,6 +17,9 @@ import os import sys import tempfile +import unittest + +from diffusers.utils import is_transformers_version sys.path.append("..") @@ -30,6 +33,7 @@ logger.addHandler(stream_handler) +@unittest.skipIf(is_transformers_version(">=", "4.57.5"), "Size mismatch") class CustomDiffusion(ExamplesTestsAccelerate): def test_custom_diffusion(self): with tempfile.TemporaryDirectory() as tmpdir: diff --git a/examples/custom_diffusion/train_custom_diffusion.py b/examples/custom_diffusion/train_custom_diffusion.py index 5922b7443c10..2ce451917709 100644 --- a/examples/custom_diffusion/train_custom_diffusion.py +++ b/examples/custom_diffusion/train_custom_diffusion.py @@ -64,7 +64,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/README_flux2.md b/examples/dreambooth/README_flux2.md index 1d1777811387..ad5d61f1f9e2 100644 --- a/examples/dreambooth/README_flux2.md +++ b/examples/dreambooth/README_flux2.md @@ -1,14 +1,22 @@ -# DreamBooth training example for FLUX.2 [dev] +# DreamBooth training example for FLUX.2 [dev] and FLUX 2 [klein] [DreamBooth](https://huggingface.co/papers/2208.12242) is a method to personalize image generation models given just a few (3~5) images of a subject/concept. +[LoRA](https://huggingface.co/docs/peft/conceptual_guides/adapter#low-rank-adaptation-lora) is a popular parameter-efficient fine-tuning technique that allows you to achieve full-finetuning like performance but with a fraction of learnable parameters. + +The `train_dreambooth_lora_flux2.py`, `train_dreambooth_lora_flux2_klein.py` scripts shows how to implement the training procedure for [LoRAs](https://huggingface.co/blog/lora) and adapt it for [FLUX.2 [dev]](https://huggingface.co/black-forest-labs/FLUX.2-dev) and [FLUX 2 [klein]](https://huggingface.co/black-forest-labs/FLUX.2-klein). -The `train_dreambooth_lora_flux2.py` script shows how to implement the training procedure for [LoRAs](https://huggingface.co/blog/lora) and adapt it for [FLUX.2 [dev]](https://github.com/black-forest-labs/flux2). +> [!NOTE] +> **Model Variants** +> +> We support two FLUX model families: +> - **FLUX.2 [dev]**: The full-size model using Mistral Small 3.1 as the text encoder. Very capable but memory intensive. +> - **FLUX 2 [klein]**: Available in 4B and 9B parameter variants, using Qwen VL as the text encoder. Much more memory efficient and suitable for consumer hardware. > [!NOTE] > **Memory consumption** > -> Flux can be quite expensive to run on consumer hardware devices and as a result finetuning it comes with high memory requirements - -> a LoRA with a rank of 16 can exceed XXGB of VRAM for training. below we provide some tips and tricks to reduce memory consumption during training. +> FLUX.2 [dev] can be quite expensive to run on consumer hardware devices and as a result finetuning it comes with high memory requirements - +> a LoRA with a rank of 16 can exceed XXGB of VRAM for training. FLUX 2 [klein] models (4B and 9B) are significantly more memory efficient alternatives. Below we provide some tips and tricks to reduce memory consumption during training. > For more tips & guidance on training on a resource-constrained device and general good practices please check out these great guides and trainers for FLUX: > 1) [`@bghira`'s guide](https://github.com/bghira/SimpleTuner/blob/main/documentation/quickstart/FLUX2.md) @@ -17,7 +25,7 @@ The `train_dreambooth_lora_flux2.py` script shows how to implement the training > [!NOTE] > **Gated model** > -> As the model is gated, before using it with diffusers you first need to go to the [FLUX.2 [dev] Hugging Face page](https://huggingface.co/black-forest-labs/FLUX.2-dev), fill in the form and accept the gate. Once you are in, you need to log in so that your system knows you’ve accepted the gate. Use the command below to log in: +> As the model is gated, before using it with diffusers you first need to go to the [FLUX.2 [dev] Hugging Face page](https://huggingface.co/black-forest-labs/FLUX.2-dev), fill in the form and accept the gate. Once you are in, you need to log in so that your system knows you've accepted the gate. Use the command below to log in: ```bash hf auth login @@ -88,20 +96,32 @@ snapshot_download( This will also allow us to push the trained LoRA parameters to the Hugging Face Hub platform. -As mentioned, Flux2 LoRA training is *very* memory intensive. Here are memory optimizations we can use (some still experimental) for a more memory efficient training: +As mentioned, Flux2 LoRA training is *very* memory intensive (especially for FLUX.2 [dev]). Here are memory optimizations we can use (some still experimental) for a more memory efficient training: ## Memory Optimizations > [!NOTE] many of these techniques complement each other and can be used together to further reduce memory consumption. > However some techniques may be mutually exclusive so be sure to check before launching a training run. + ### Remote Text Encoder -Flux.2 uses Mistral Small 3.1 as text encoder which is quite large and can take up a lot of memory. To mitigate this, we can use the `--remote_text_encoder` flag to enable remote computation of the prompt embeddings using the HuggingFace Inference API. +FLUX.2 [dev] uses Mistral Small 3.1 as text encoder which is quite large and can take up a lot of memory. To mitigate this, we can use the `--remote_text_encoder` flag to enable remote computation of the prompt embeddings using the HuggingFace Inference API. This way, the text encoder model is not loaded into memory during training. + +> [!IMPORTANT] +> **Remote text encoder is only supported for FLUX.2 [dev]**. FLUX 2 [klein] models use the Qwen VL text encoder and do not support remote text encoding. + > [!NOTE] > to enable remote text encoding you must either be logged in to your HuggingFace account (`hf auth login`) OR pass a token with `--hub_token`. + +### FSDP Text Encoder +FLUX.2 [dev] uses Mistral Small 3.1 as text encoder which is quite large and can take up a lot of memory. To mitigate this, we can use the `--fsdp_text_encoder` flag to enable distributed computation of the prompt embeddings. +This way, it distributes the memory cost across multiple nodes. + ### CPU Offloading To offload parts of the model to CPU memory, you can use `--offload` flag. This will offload the vae and text encoder to CPU memory and only move them to GPU when needed. + ### Latent Caching Pre-encode the training images with the vae, and then delete it to free up some memory. To enable `latent_caching` simply pass `--cache_latents`. + ### QLoRA: Low Precision Training with Quantization Perform low precision training using 8-bit or 4-bit quantization to reduce memory usage. You can use the following flags: - **FP8 training** with `torchao`: @@ -111,22 +131,29 @@ enable FP8 training by passing `--do_fp8_training`. - **NF4 training** with `bitsandbytes`: Alternatively, you can use 8-bit or 4-bit quantization with `bitsandbytes` by passing: `--bnb_quantization_config_path` to enable 4-bit NF4 quantization. + ### Gradient Checkpointing and Accumulation * `--gradient accumulation` refers to the number of updates steps to accumulate before performing a backward/update pass. by passing a value > 1 you can reduce the amount of backward/update passes and hence also memory reqs. * with `--gradient checkpointing` we can save memory by not storing all intermediate activations during the forward pass. Instead, only a subset of these activations (the checkpoints) are stored and the rest is recomputed as needed during the backward pass. Note that this comes at the expanse of a slower backward pass. + ### 8-bit-Adam Optimizer When training with `AdamW`(doesn't apply to `prodigy`) You can pass `--use_8bit_adam` to reduce the memory requirements of training. Make sure to install `bitsandbytes` if you want to do so. + ### Image Resolution An easy way to mitigate some of the memory requirements is through `--resolution`. `--resolution` refers to the resolution for input images, all the images in the train/validation dataset are resized to this. Note that by default, images are resized to resolution of 512, but it's good to keep in mind in case you're accustomed to training on higher resolutions. + ### Precision of saved LoRA layers By default, trained transformer layers are saved in the precision dtype in which training was performed. E.g. when training in mixed precision is enabled with `--mixed_precision="bf16"`, final finetuned layers will be saved in `torch.bfloat16` as well. This reduces memory requirements significantly w/o a significant quality loss. Note that if you do wish to save the final layers in float32 at the expanse of more memory usage, you can do so by passing `--upcast_before_saving`. +## Training Examples +### FLUX.2 [dev] Training +To perform DreamBooth with LoRA on FLUX.2 [dev], run: ```bash export MODEL_NAME="black-forest-labs/FLUX.2-dev" export INSTANCE_DIR="dog" @@ -158,19 +185,104 @@ accelerate launch train_dreambooth_lora_flux2.py \ --push_to_hub ``` -To better track our training experiments, we're using the following flags in the command above: +### FLUX 2 [klein] Training -* `report_to="wandb` will ensure the training runs are tracked on [Weights and Biases](https://wandb.ai/site). To use it, be sure to install `wandb` with `pip install wandb`. Don't forget to call `wandb login ` before training if you haven't done it before. -* `validation_prompt` and `validation_epochs` to allow the script to do a few validation inference runs. This allows us to qualitatively check if the training is progressing as expected. +FLUX 2 [klein] models are more memory efficient alternatives available in 4B and 9B parameter variants. They use the Qwen VL text encoder instead of Mistral Small 3.1. > [!NOTE] -> If you want to train using long prompts with the T5 text encoder, you can use `--max_sequence_length` to set the token limit. The default is 77, but it can be increased to as high as 512. Note that this will use more resources and may slow down the training in some cases. +> The `--remote_text_encoder` flag is **not supported** for FLUX 2 [klein] models. The Qwen VL text encoder must be loaded locally, but offloading is still supported. -## LoRA + DreamBooth +**FLUX 2 [klein] 4B:** -[LoRA](https://huggingface.co/docs/peft/conceptual_guides/adapter#low-rank-adaptation-lora) is a popular parameter-efficient fine-tuning technique that allows you to achieve full-finetuning like performance but with a fraction of learnable parameters. +```bash +export MODEL_NAME="black-forest-labs/FLUX.2-klein-4B" +export INSTANCE_DIR="dog" +export OUTPUT_DIR="trained-flux2-klein-4b" -Note also that we use PEFT library as backend for LoRA training, make sure to have `peft>=0.6.0` installed in your environment. +accelerate launch train_dreambooth_lora_flux2_klein.py \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --instance_data_dir=$INSTANCE_DIR \ + --output_dir=$OUTPUT_DIR \ + --do_fp8_training \ + --gradient_checkpointing \ + --cache_latents \ + --instance_prompt="a photo of sks dog" \ + --resolution=1024 \ + --train_batch_size=1 \ + --guidance_scale=1 \ + --use_8bit_adam \ + --gradient_accumulation_steps=4 \ + --optimizer="adamW" \ + --learning_rate=1e-4 \ + --report_to="wandb" \ + --lr_scheduler="constant" \ + --lr_warmup_steps=100 \ + --max_train_steps=500 \ + --validation_prompt="A photo of sks dog in a bucket" \ + --validation_epochs=25 \ + --seed="0" \ + --push_to_hub +``` + +**FLUX 2 [klein] 9B:** + +```bash +export MODEL_NAME="black-forest-labs/FLUX.2-klein-9B" +export INSTANCE_DIR="dog" +export OUTPUT_DIR="trained-flux2-klein-9b" + +accelerate launch train_dreambooth_lora_flux2_klein.py \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --instance_data_dir=$INSTANCE_DIR \ + --output_dir=$OUTPUT_DIR \ + --do_fp8_training \ + --gradient_checkpointing \ + --cache_latents \ + --instance_prompt="a photo of sks dog" \ + --resolution=1024 \ + --train_batch_size=1 \ + --guidance_scale=1 \ + --use_8bit_adam \ + --gradient_accumulation_steps=4 \ + --optimizer="adamW" \ + --learning_rate=1e-4 \ + --report_to="wandb" \ + --lr_scheduler="constant" \ + --lr_warmup_steps=100 \ + --max_train_steps=500 \ + --validation_prompt="A photo of sks dog in a bucket" \ + --validation_epochs=25 \ + --seed="0" \ + --push_to_hub +``` + +To better track our training experiments, we're using the following flags in the command above: + +* `report_to="wandb` will ensure the training runs are tracked on [Weights and Biases](https://wandb.ai/site). To use it, be sure to install `wandb` with `pip install wandb`. Don't forget to call `wandb login ` before training if you haven't done it before. +* `validation_prompt` and `validation_epochs` to allow the script to do a few validation inference runs. This allows us to qualitatively check if the training is progressing as expected. + +> [!NOTE] +> If you want to train using long prompts, you can use `--max_sequence_length` to set the token limit. Note that this will use more resources and may slow down the training in some cases. + +### FSDP on the transformer +By setting the accelerate configuration with FSDP, the transformer block will be wrapped automatically. E.g. set the configuration to: + +```shell +distributed_type: FSDP +fsdp_config: + fsdp_version: 2 + fsdp_offload_params: false + fsdp_sharding_strategy: HYBRID_SHARD + fsdp_auto_wrap_policy: TRANSFOMER_BASED_WRAP + fsdp_transformer_layer_cls_to_wrap: Flux2TransformerBlock, Flux2SingleTransformerBlock + fsdp_forward_prefetch: true + fsdp_sync_module_states: false + fsdp_state_dict_type: FULL_STATE_DICT + fsdp_use_orig_params: false + fsdp_activation_checkpointing: true + fsdp_reshard_after_forward: true + fsdp_cpu_ram_efficient_loading: false +``` ### Prodigy Optimizer Prodigy is an adaptive optimizer that dynamically adjusts the learning rate learned parameters based on past gradients, allowing for more efficient convergence. @@ -183,8 +295,6 @@ to use prodigy, first make sure to install the prodigyopt library: `pip install > [!TIP] > When using prodigy it's generally good practice to set- `--learning_rate=1.0` -To perform DreamBooth with LoRA, run: - ```bash export MODEL_NAME="black-forest-labs/FLUX.2-dev" export INSTANCE_DIR="dog" @@ -248,13 +358,10 @@ the exact modules for LoRA training. Here are some examples of target modules yo > keep in mind that while training more layers can improve quality and expressiveness, it also increases the size of the output LoRA weights. - ## Training Image-to-Image Flux.2 lets us perform image editing as well as image generation. We provide a simple script for image-to-image(I2I) LoRA fine-tuning in [train_dreambooth_lora_flux2_img2img.py](./train_dreambooth_lora_flux2_img2img.py) for both T2I and I2I. The optimizations discussed above apply this script, too. -**important** - **Important** To make sure you can successfully run the latest version of the image-to-image example script, we highly recommend installing from source, specifically from the commit mentioned below. To do this, execute the following steps in a new virtual environment: @@ -311,5 +418,6 @@ we've added aspect ratio bucketing support which allows training on images with To enable aspect ratio bucketing, pass `--aspect_ratio_buckets` argument with a semicolon-separated list of height,width pairs, such as: `--aspect_ratio_buckets="672,1568;688,1504;720,1456;752,1392;800,1328;832,1248;880,1184;944,1104;1024,1024;1104,944;1184,880;1248,832;1328,800;1392,752;1456,720;1504,688;1568,672" -` -Since Flux.2 finetuning is still an experimental phase, we encourage you to explore different settings and share your insights! 🤗 + + +Since Flux.2 finetuning is still an experimental phase, we encourage you to explore different settings and share your insights! 🤗 \ No newline at end of file diff --git a/examples/dreambooth/README_sana.md b/examples/dreambooth/README_sana.md index 7972434b5e6f..8bddacf975d8 100644 --- a/examples/dreambooth/README_sana.md +++ b/examples/dreambooth/README_sana.md @@ -111,6 +111,25 @@ To better track our training experiments, we're using the following flags in the ## Notes +### LoRA Rank and Alpha +Two key LoRA hyperparameters are LoRA rank and LoRA alpha. +- `--rank`: Defines the dimension of the trainable LoRA matrices. A higher rank means more expressiveness and capacity to learn (and more parameters). +- `--lora_alpha`: A scaling factor for the LoRA's output. The LoRA update is scaled by lora_alpha / lora_rank. +- lora_alpha vs. rank: +This ratio dictates the LoRA's effective strength: +lora_alpha == rank: Scaling factor is 1. The LoRA is applied with its learned strength. (e.g., alpha=16, rank=16) +lora_alpha < rank: Scaling factor < 1. Reduces the LoRA's impact. Useful for subtle changes or to prevent overpowering the base model. (e.g., alpha=8, rank=16) +lora_alpha > rank: Scaling factor > 1. Amplifies the LoRA's impact. Allows a lower rank LoRA to have a stronger effect. (e.g., alpha=32, rank=16) + +> [!TIP] +> A common starting point is to set `lora_alpha` equal to `rank`. +> Some also set `lora_alpha` to be twice the `rank` (e.g., lora_alpha=32 for lora_rank=16) +> to give the LoRA updates more influence without increasing parameter count. +> If you find your LoRA is "overcooking" or learning too aggressively, consider setting `lora_alpha` to half of `rank` +> (e.g., lora_alpha=8 for rank=16). Experimentation is often key to finding the optimal balance for your use case. + +### Additional CLI arguments + Additionally, we welcome you to explore the following CLI arguments: * `--lora_layers`: The transformer modules to apply LoRA training on. Please specify the layers in a comma separated. E.g. - "to_k,to_q,to_v" will result in lora training of attention layers only. diff --git a/examples/dreambooth/README_z_image.md b/examples/dreambooth/README_z_image.md new file mode 100644 index 000000000000..cded38f3f11f --- /dev/null +++ b/examples/dreambooth/README_z_image.md @@ -0,0 +1,347 @@ +# DreamBooth training example for Z-Image + +[DreamBooth](https://huggingface.co/papers/2208.12242) is a method to personalize image generation models given just a few (3~5) images of a subject/concept. +[LoRA](https://huggingface.co/docs/peft/conceptual_guides/adapter#low-rank-adaptation-lora) is a popular parameter-efficient fine-tuning technique that allows you to achieve full-finetuning like performance but with a fraction of learnable parameters. + +The `train_dreambooth_lora_z_image.py` script shows how to implement the training procedure for [LoRAs](https://huggingface.co/blog/lora) and adapt it for [Z-Image](https://huggingface.co/Tongyi-MAI/Z-Image). + +> [!NOTE] +> **About Z-Image** +> +> Z-Image is a high-quality text-to-image generation model from Alibaba's Tongyi Lab. It uses a DiT (Diffusion Transformer) architecture with Qwen3 as the text encoder. The model excels at generating images with accurate text rendering, especially for Chinese characters. + +> [!NOTE] +> **Memory consumption** +> +> Z-Image is relatively memory efficient compared to other large-scale diffusion models. Below we provide some tips and tricks to further reduce memory consumption during training. + +## Running locally with PyTorch + +### Installing the dependencies + +Before running the scripts, make sure to install the library's training dependencies: + +**Important** + +To make sure you can successfully run the latest versions of the example scripts, we highly recommend **installing from source** and keeping the install up to date as we update the example scripts frequently and install some example-specific requirements. To do this, execute the following steps in a new virtual environment: + +```bash +git clone https://github.com/huggingface/diffusers +cd diffusers +pip install -e . +``` + +Then cd in the `examples/dreambooth` folder and run +```bash +pip install -r requirements_z_image.txt +``` + +And initialize an [🤗Accelerate](https://github.com/huggingface/accelerate/) environment with: + +```bash +accelerate config +``` + +Or for a default accelerate configuration without answering questions about your environment + +```bash +accelerate config default +``` + +Or if your environment doesn't support an interactive shell (e.g., a notebook) + +```python +from accelerate.utils import write_basic_config +write_basic_config() +``` + +When running `accelerate config`, if we specify torch compile mode to True there can be dramatic speedups. +Note also that we use PEFT library as backend for LoRA training, make sure to have `peft>=0.6.0` installed in your environment. + + +### Dog toy example + +Now let's get our dataset. For this example we will use some dog images: https://huggingface.co/datasets/diffusers/dog-example. + +Let's first download it locally: + +```python +from huggingface_hub import snapshot_download + +local_dir = "./dog" +snapshot_download( + "diffusers/dog-example", + local_dir=local_dir, repo_type="dataset", + ignore_patterns=".gitattributes", +) +``` + +This will also allow us to push the trained LoRA parameters to the Hugging Face Hub platform. + +## Memory Optimizations + +> [!NOTE] +> Many of these techniques complement each other and can be used together to further reduce memory consumption. However some techniques may be mutually exclusive so be sure to check before launching a training run. + +### CPU Offloading +To offload parts of the model to CPU memory, you can use `--offload` flag. This will offload the VAE and text encoder to CPU memory and only move them to GPU when needed. + +### Latent Caching +Pre-encode the training images with the VAE, and then delete it to free up some memory. To enable `latent_caching` simply pass `--cache_latents`. + +### QLoRA: Low Precision Training with Quantization +Perform low precision training using 8-bit or 4-bit quantization to reduce memory usage. You can use the following flags: + +- **FP8 training** with `torchao`: +Enable FP8 training by passing `--do_fp8_training`. +> [!IMPORTANT] +> Since we are utilizing FP8 tensor cores we need CUDA GPUs with compute capability at least 8.9 or greater. If you're looking for memory-efficient training on relatively older cards, we encourage you to check out other trainers. + +- **NF4 training** with `bitsandbytes`: +Alternatively, you can use 8-bit or 4-bit quantization with `bitsandbytes` by passing `--bnb_quantization_config_path` to enable 4-bit NF4 quantization. + +### Gradient Checkpointing and Accumulation +* `--gradient_accumulation` refers to the number of updates steps to accumulate before performing a backward/update pass. By passing a value > 1 you can reduce the amount of backward/update passes and hence also memory requirements. +* With `--gradient_checkpointing` we can save memory by not storing all intermediate activations during the forward pass. Instead, only a subset of these activations (the checkpoints) are stored and the rest is recomputed as needed during the backward pass. Note that this comes at the expense of a slower backward pass. + +### 8-bit-Adam Optimizer +When training with `AdamW` (doesn't apply to `prodigy`) you can pass `--use_8bit_adam` to reduce the memory requirements of training. Make sure to install `bitsandbytes` if you want to do so. + +### Image Resolution +An easy way to mitigate some of the memory requirements is through `--resolution`. `--resolution` refers to the resolution for input images, all the images in the train/validation dataset are resized to this. +Note that by default, images are resized to resolution of 1024, but it's good to keep in mind in case you're training on higher resolutions. + +### Precision of saved LoRA layers +By default, trained transformer layers are saved in the precision dtype in which training was performed. E.g. when training in mixed precision is enabled with `--mixed_precision="bf16"`, final finetuned layers will be saved in `torch.bfloat16` as well. +This reduces memory requirements significantly without a significant quality loss. Note that if you do wish to save the final layers in float32 at the expense of more memory usage, you can do so by passing `--upcast_before_saving`. + +## Training Examples + +### Z-Image Training + +To perform DreamBooth with LoRA on Z-Image, run: + +```bash +export MODEL_NAME="Tongyi-MAI/Z-Image" +export INSTANCE_DIR="dog" +export OUTPUT_DIR="trained-z-image-lora" + +accelerate launch train_dreambooth_lora_z_image.py \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --instance_data_dir=$INSTANCE_DIR \ + --output_dir=$OUTPUT_DIR \ + --mixed_precision="bf16" \ + --gradient_checkpointing \ + --cache_latents \ + --instance_prompt="a photo of sks dog" \ + --resolution=1024 \ + --train_batch_size=1 \ + --guidance_scale=5.0 \ + --use_8bit_adam \ + --gradient_accumulation_steps=4 \ + --optimizer="adamW" \ + --learning_rate=1e-4 \ + --report_to="wandb" \ + --lr_scheduler="constant" \ + --lr_warmup_steps=100 \ + --max_train_steps=500 \ + --validation_prompt="A photo of sks dog in a bucket" \ + --validation_epochs=25 \ + --seed="0" \ + --push_to_hub +``` + +To better track our training experiments, we're using the following flags in the command above: + +* `report_to="wandb"` will ensure the training runs are tracked on [Weights and Biases](https://wandb.ai/site). To use it, be sure to install `wandb` with `pip install wandb`. Don't forget to call `wandb login ` before training if you haven't done it before. +* `validation_prompt` and `validation_epochs` to allow the script to do a few validation inference runs. This allows us to qualitatively check if the training is progressing as expected. + +> [!NOTE] +> If you want to train using long prompts, you can use `--max_sequence_length` to set the token limit. The default is 512. Note that this will use more resources and may slow down the training in some cases. + +### Training with FP8 Quantization + +For reduced memory usage with FP8 training: + +```bash +export MODEL_NAME="Tongyi-MAI/Z-Image" +export INSTANCE_DIR="dog" +export OUTPUT_DIR="trained-z-image-lora-fp8" + +accelerate launch train_dreambooth_lora_z_image.py \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --instance_data_dir=$INSTANCE_DIR \ + --output_dir=$OUTPUT_DIR \ + --do_fp8_training \ + --gradient_checkpointing \ + --cache_latents \ + --instance_prompt="a photo of sks dog" \ + --resolution=1024 \ + --train_batch_size=1 \ + --guidance_scale=5.0 \ + --use_8bit_adam \ + --gradient_accumulation_steps=4 \ + --optimizer="adamW" \ + --learning_rate=1e-4 \ + --report_to="wandb" \ + --lr_scheduler="constant" \ + --lr_warmup_steps=100 \ + --max_train_steps=500 \ + --validation_prompt="A photo of sks dog in a bucket" \ + --validation_epochs=25 \ + --seed="0" \ + --push_to_hub +``` + +### FSDP on the transformer + +By setting the accelerate configuration with FSDP, the transformer block will be wrapped automatically. E.g. set the configuration to: + +```yaml +distributed_type: FSDP +fsdp_config: + fsdp_version: 2 + fsdp_offload_params: false + fsdp_sharding_strategy: HYBRID_SHARD + fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP + fsdp_transformer_layer_cls_to_wrap: ZImageTransformerBlock + fsdp_forward_prefetch: true + fsdp_sync_module_states: false + fsdp_state_dict_type: FULL_STATE_DICT + fsdp_use_orig_params: false + fsdp_activation_checkpointing: true + fsdp_reshard_after_forward: true + fsdp_cpu_ram_efficient_loading: false +``` + +### Prodigy Optimizer + +Prodigy is an adaptive optimizer that dynamically adjusts the learning rate learned parameters based on past gradients, allowing for more efficient convergence. +By using prodigy we can "eliminate" the need for manual learning rate tuning. Read more [here](https://huggingface.co/blog/sdxl_lora_advanced_script#adaptive-optimizers). + +To use prodigy, first make sure to install the prodigyopt library: `pip install prodigyopt`, and then specify: +```bash +--optimizer="prodigy" +``` + +> [!TIP] +> When using prodigy it's generally good practice to set `--learning_rate=1.0` + +```bash +export MODEL_NAME="Tongyi-MAI/Z-Image" +export INSTANCE_DIR="dog" +export OUTPUT_DIR="trained-z-image-lora-prodigy" + +accelerate launch train_dreambooth_lora_z_image.py \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --instance_data_dir=$INSTANCE_DIR \ + --output_dir=$OUTPUT_DIR \ + --mixed_precision="bf16" \ + --gradient_checkpointing \ + --cache_latents \ + --instance_prompt="a photo of sks dog" \ + --resolution=1024 \ + --train_batch_size=1 \ + --guidance_scale=5.0 \ + --gradient_accumulation_steps=4 \ + --optimizer="prodigy" \ + --learning_rate=1.0 \ + --report_to="wandb" \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --max_train_steps=500 \ + --validation_prompt="A photo of sks dog in a bucket" \ + --validation_epochs=25 \ + --seed="0" \ + --push_to_hub +``` + +### LoRA Rank and Alpha + +Two key LoRA hyperparameters are LoRA rank and LoRA alpha: + +- `--rank`: Defines the dimension of the trainable LoRA matrices. A higher rank means more expressiveness and capacity to learn (and more parameters). +- `--lora_alpha`: A scaling factor for the LoRA's output. The LoRA update is scaled by `lora_alpha / lora_rank`. + +**lora_alpha vs. rank:** + +This ratio dictates the LoRA's effective strength: +- `lora_alpha == rank`: Scaling factor is 1. The LoRA is applied with its learned strength. (e.g., alpha=16, rank=16) +- `lora_alpha < rank`: Scaling factor < 1. Reduces the LoRA's impact. Useful for subtle changes or to prevent overpowering the base model. (e.g., alpha=8, rank=16) +- `lora_alpha > rank`: Scaling factor > 1. Amplifies the LoRA's impact. Allows a lower rank LoRA to have a stronger effect. (e.g., alpha=32, rank=16) + +> [!TIP] +> A common starting point is to set `lora_alpha` equal to `rank`. +> Some also set `lora_alpha` to be twice the `rank` (e.g., lora_alpha=32 for lora_rank=16) +> to give the LoRA updates more influence without increasing parameter count. +> If you find your LoRA is "overcooking" or learning too aggressively, consider setting `lora_alpha` to half of `rank` +> (e.g., lora_alpha=8 for rank=16). Experimentation is often key to finding the optimal balance for your use case. + +### Target Modules + +When LoRA was first adapted from language models to diffusion models, it was applied to the cross-attention layers in the UNet that relate the image representations with the prompts that describe them. +More recently, SOTA text-to-image diffusion models replaced the UNet with a diffusion Transformer (DiT). With this change, we may also want to explore applying LoRA training onto different types of layers and blocks. + +To allow more flexibility and control over the targeted modules we added `--lora_layers`, in which you can specify in a comma separated string the exact modules for LoRA training. Here are some examples of target modules you can provide: + +- For attention only layers: `--lora_layers="to_k,to_q,to_v,to_out.0"` +- For attention and feed-forward layers: `--lora_layers="to_k,to_q,to_v,to_out.0,ff.net.0.proj,ff.net.2"` + +> [!NOTE] +> `--lora_layers` can also be used to specify which **blocks** to apply LoRA training to. To do so, simply add a block prefix to each layer in the comma separated string. + +> [!NOTE] +> Keep in mind that while training more layers can improve quality and expressiveness, it also increases the size of the output LoRA weights. + +### Aspect Ratio Bucketing + +We've added aspect ratio bucketing support which allows training on images with different aspect ratios without cropping them to a single square resolution. This technique helps preserve the original composition of training images and can improve training efficiency. + +To enable aspect ratio bucketing, pass `--aspect_ratio_buckets` argument with a semicolon-separated list of height,width pairs, such as: + +```bash +--aspect_ratio_buckets="672,1568;688,1504;720,1456;752,1392;800,1328;832,1248;880,1184;944,1104;1024,1024;1104,944;1184,880;1248,832;1328,800;1392,752;1456,720;1504,688;1568,672" +``` + +### Bilingual Prompts + +Z-Image has strong support for both Chinese and English prompts. When training with Chinese prompts, ensure your dataset captions are properly encoded in UTF-8: + +```bash +--instance_prompt="一只sks狗的照片" +--validation_prompt="一只sks狗在桶里的照片" +``` + +> [!TIP] +> Z-Image excels at text rendering in generated images, especially for Chinese characters. If your use case involves generating images with text, consider including text-related examples in your training data. + +## Inference + +Once you have trained a LoRA, you can load it for inference: + +```python +import torch +from diffusers import ZImagePipeline + +pipe = ZImagePipeline.from_pretrained("Tongyi-MAI/Z-Image", torch_dtype=torch.bfloat16) +pipe.to("cuda") + +# Load your trained LoRA +pipe.load_lora_weights("path/to/your/trained-z-image-lora") + +# Generate an image +image = pipe( + prompt="A photo of sks dog in a bucket", + height=1024, + width=1024, + num_inference_steps=50, + guidance_scale=5.0, + generator=torch.Generator("cuda").manual_seed(42), +).images[0] + +image.save("output.png") +``` + +--- + +Since Z-Image finetuning is still in an experimental phase, we encourage you to explore different settings and share your insights! 🤗 \ No newline at end of file diff --git a/examples/dreambooth/test_dreambooth_lora_flux2_klein.py b/examples/dreambooth/test_dreambooth_lora_flux2_klein.py new file mode 100644 index 000000000000..0e5506e1a3eb --- /dev/null +++ b/examples/dreambooth/test_dreambooth_lora_flux2_klein.py @@ -0,0 +1,262 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import logging +import os +import sys +import tempfile + +import safetensors + +from diffusers.loaders.lora_base import LORA_ADAPTER_METADATA_KEY + + +sys.path.append("..") +from test_examples_utils import ExamplesTestsAccelerate, run_command # noqa: E402 + + +logging.basicConfig(level=logging.DEBUG) + +logger = logging.getLogger() +stream_handler = logging.StreamHandler(sys.stdout) +logger.addHandler(stream_handler) + + +class DreamBoothLoRAFlux2Klein(ExamplesTestsAccelerate): + instance_data_dir = "docs/source/en/imgs" + instance_prompt = "dog" + pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-klein" + script_path = "examples/dreambooth/train_dreambooth_lora_flux2_klein.py" + transformer_layer_type = "single_transformer_blocks.0.attn.to_qkv_mlp_proj" + + def test_dreambooth_lora_flux2(self): + with tempfile.TemporaryDirectory() as tmpdir: + test_args = f""" + {self.script_path} + --pretrained_model_name_or_path {self.pretrained_model_name_or_path} + --instance_data_dir {self.instance_data_dir} + --instance_prompt {self.instance_prompt} + --resolution 64 + --train_batch_size 1 + --gradient_accumulation_steps 1 + --max_train_steps 2 + --learning_rate 5.0e-04 + --scale_lr + --lr_scheduler constant + --lr_warmup_steps 0 + --max_sequence_length 8 + --text_encoder_out_layers 1 + --output_dir {tmpdir} + """.split() + + run_command(self._launch_args + test_args) + # save_pretrained smoke test + self.assertTrue(os.path.isfile(os.path.join(tmpdir, "pytorch_lora_weights.safetensors"))) + + # make sure the state_dict has the correct naming in the parameters. + lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors")) + is_lora = all("lora" in k for k in lora_state_dict.keys()) + self.assertTrue(is_lora) + + # when not training the text encoder, all the parameters in the state dict should start + # with `"transformer"` in their names. + starts_with_transformer = all(key.startswith("transformer") for key in lora_state_dict.keys()) + self.assertTrue(starts_with_transformer) + + def test_dreambooth_lora_latent_caching(self): + with tempfile.TemporaryDirectory() as tmpdir: + test_args = f""" + {self.script_path} + --pretrained_model_name_or_path {self.pretrained_model_name_or_path} + --instance_data_dir {self.instance_data_dir} + --instance_prompt {self.instance_prompt} + --resolution 64 + --train_batch_size 1 + --gradient_accumulation_steps 1 + --max_train_steps 2 + --cache_latents + --learning_rate 5.0e-04 + --scale_lr + --lr_scheduler constant + --lr_warmup_steps 0 + --max_sequence_length 8 + --text_encoder_out_layers 1 + --output_dir {tmpdir} + """.split() + + run_command(self._launch_args + test_args) + # save_pretrained smoke test + self.assertTrue(os.path.isfile(os.path.join(tmpdir, "pytorch_lora_weights.safetensors"))) + + # make sure the state_dict has the correct naming in the parameters. + lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors")) + is_lora = all("lora" in k for k in lora_state_dict.keys()) + self.assertTrue(is_lora) + + # when not training the text encoder, all the parameters in the state dict should start + # with `"transformer"` in their names. + starts_with_transformer = all(key.startswith("transformer") for key in lora_state_dict.keys()) + self.assertTrue(starts_with_transformer) + + def test_dreambooth_lora_layers(self): + with tempfile.TemporaryDirectory() as tmpdir: + test_args = f""" + {self.script_path} + --pretrained_model_name_or_path {self.pretrained_model_name_or_path} + --instance_data_dir {self.instance_data_dir} + --instance_prompt {self.instance_prompt} + --resolution 64 + --train_batch_size 1 + --gradient_accumulation_steps 1 + --max_train_steps 2 + --cache_latents + --learning_rate 5.0e-04 + --scale_lr + --lora_layers {self.transformer_layer_type} + --lr_scheduler constant + --lr_warmup_steps 0 + --max_sequence_length 8 + --text_encoder_out_layers 1 + --output_dir {tmpdir} + """.split() + + run_command(self._launch_args + test_args) + # save_pretrained smoke test + self.assertTrue(os.path.isfile(os.path.join(tmpdir, "pytorch_lora_weights.safetensors"))) + + # make sure the state_dict has the correct naming in the parameters. + lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors")) + is_lora = all("lora" in k for k in lora_state_dict.keys()) + self.assertTrue(is_lora) + + # when not training the text encoder, all the parameters in the state dict should start + # with `"transformer"` in their names. In this test, we only params of + # transformer.single_transformer_blocks.0.attn.to_k should be in the state dict + starts_with_transformer = all( + key.startswith(f"transformer.{self.transformer_layer_type}") for key in lora_state_dict.keys() + ) + self.assertTrue(starts_with_transformer) + + def test_dreambooth_lora_flux2_checkpointing_checkpoints_total_limit(self): + with tempfile.TemporaryDirectory() as tmpdir: + test_args = f""" + {self.script_path} + --pretrained_model_name_or_path={self.pretrained_model_name_or_path} + --instance_data_dir={self.instance_data_dir} + --output_dir={tmpdir} + --instance_prompt={self.instance_prompt} + --resolution=64 + --train_batch_size=1 + --gradient_accumulation_steps=1 + --max_train_steps=6 + --checkpoints_total_limit=2 + --max_sequence_length 8 + --checkpointing_steps=2 + --text_encoder_out_layers 1 + """.split() + + run_command(self._launch_args + test_args) + + self.assertEqual( + {x for x in os.listdir(tmpdir) if "checkpoint" in x}, + {"checkpoint-4", "checkpoint-6"}, + ) + + def test_dreambooth_lora_flux2_checkpointing_checkpoints_total_limit_removes_multiple_checkpoints(self): + with tempfile.TemporaryDirectory() as tmpdir: + test_args = f""" + {self.script_path} + --pretrained_model_name_or_path={self.pretrained_model_name_or_path} + --instance_data_dir={self.instance_data_dir} + --output_dir={tmpdir} + --instance_prompt={self.instance_prompt} + --resolution=64 + --train_batch_size=1 + --gradient_accumulation_steps=1 + --max_train_steps=4 + --checkpointing_steps=2 + --max_sequence_length 8 + --text_encoder_out_layers 1 + """.split() + + run_command(self._launch_args + test_args) + + self.assertEqual({x for x in os.listdir(tmpdir) if "checkpoint" in x}, {"checkpoint-2", "checkpoint-4"}) + + resume_run_args = f""" + {self.script_path} + --pretrained_model_name_or_path={self.pretrained_model_name_or_path} + --instance_data_dir={self.instance_data_dir} + --output_dir={tmpdir} + --instance_prompt={self.instance_prompt} + --resolution=64 + --train_batch_size=1 + --gradient_accumulation_steps=1 + --max_train_steps=8 + --checkpointing_steps=2 + --resume_from_checkpoint=checkpoint-4 + --checkpoints_total_limit=2 + --max_sequence_length 8 + --text_encoder_out_layers 1 + """.split() + + run_command(self._launch_args + resume_run_args) + + self.assertEqual({x for x in os.listdir(tmpdir) if "checkpoint" in x}, {"checkpoint-6", "checkpoint-8"}) + + def test_dreambooth_lora_with_metadata(self): + # Use a `lora_alpha` that is different from `rank`. + lora_alpha = 8 + rank = 4 + with tempfile.TemporaryDirectory() as tmpdir: + test_args = f""" + {self.script_path} + --pretrained_model_name_or_path {self.pretrained_model_name_or_path} + --instance_data_dir {self.instance_data_dir} + --instance_prompt {self.instance_prompt} + --resolution 64 + --train_batch_size 1 + --gradient_accumulation_steps 1 + --max_train_steps 2 + --lora_alpha={lora_alpha} + --rank={rank} + --learning_rate 5.0e-04 + --scale_lr + --lr_scheduler constant + --lr_warmup_steps 0 + --max_sequence_length 8 + --text_encoder_out_layers 1 + --output_dir {tmpdir} + """.split() + + run_command(self._launch_args + test_args) + # save_pretrained smoke test + state_dict_file = os.path.join(tmpdir, "pytorch_lora_weights.safetensors") + self.assertTrue(os.path.isfile(state_dict_file)) + + # Check if the metadata was properly serialized. + with safetensors.torch.safe_open(state_dict_file, framework="pt", device="cpu") as f: + metadata = f.metadata() or {} + + metadata.pop("format", None) + raw = metadata.get(LORA_ADAPTER_METADATA_KEY) + if raw: + raw = json.loads(raw) + + loaded_lora_alpha = raw["transformer.lora_alpha"] + self.assertTrue(loaded_lora_alpha == lora_alpha) + loaded_lora_rank = raw["transformer.r"] + self.assertTrue(loaded_lora_rank == rank) diff --git a/examples/dreambooth/train_dreambooth.py b/examples/dreambooth/train_dreambooth.py index 2e66e1f724e7..d3a2b32aaef5 100644 --- a/examples/dreambooth/train_dreambooth.py +++ b/examples/dreambooth/train_dreambooth.py @@ -64,7 +64,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_flax.py b/examples/dreambooth/train_dreambooth_flax.py index e68d9df5e424..0580fb4b96b0 100644 --- a/examples/dreambooth/train_dreambooth_flax.py +++ b/examples/dreambooth/train_dreambooth_flax.py @@ -35,7 +35,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") # Cache compiled models across invocations of this script. cc.initialize_cache(os.path.expanduser("~/.cache/jax/compilation_cache")) diff --git a/examples/dreambooth/train_dreambooth_flux.py b/examples/dreambooth/train_dreambooth_flux.py index 468f6fce3ecb..c7e0c290fa8e 100644 --- a/examples/dreambooth/train_dreambooth_flux.py +++ b/examples/dreambooth/train_dreambooth_flux.py @@ -80,7 +80,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora.py b/examples/dreambooth/train_dreambooth_lora.py index 2d15684f9107..b6baccc4bc99 100644 --- a/examples/dreambooth/train_dreambooth_lora.py +++ b/examples/dreambooth/train_dreambooth_lora.py @@ -75,7 +75,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_flux.py b/examples/dreambooth/train_dreambooth_lora_flux.py index 8ae2ddd9796b..e0e7d2e40e56 100644 --- a/examples/dreambooth/train_dreambooth_lora_flux.py +++ b/examples/dreambooth/train_dreambooth_lora_flux.py @@ -92,7 +92,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_flux2.py b/examples/dreambooth/train_dreambooth_lora_flux2.py index 81306940af8f..24d098add017 100644 --- a/examples/dreambooth/train_dreambooth_lora_flux2.py +++ b/examples/dreambooth/train_dreambooth_lora_flux2.py @@ -44,6 +44,7 @@ import warnings from contextlib import nullcontext from pathlib import Path +from typing import Any import numpy as np import torch @@ -75,13 +76,16 @@ from diffusers.optimization import get_scheduler from diffusers.training_utils import ( _collate_lora_metadata, + _to_cpu_contiguous, cast_training_params, compute_density_for_timestep_sampling, compute_loss_weighting_for_sd3, find_nearest_bucket, free_memory, + get_fsdp_kwargs_from_accelerator, offload_models, parse_buckets_string, + wrap_with_fsdp, ) from diffusers.utils import ( check_min_version, @@ -93,11 +97,14 @@ from diffusers.utils.torch_utils import is_compiled_module +if getattr(torch, "distributed", None) is not None: + import torch.distributed as dist + if is_wandb_available(): import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) @@ -722,6 +729,7 @@ def parse_args(input_args=None): ) parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") parser.add_argument("--enable_npu_flash_attention", action="store_true", help="Enabla Flash Attention for NPU") + parser.add_argument("--fsdp_text_encoder", action="store_true", help="Use FSDP for text encoder") if input_args is not None: args = parser.parse_args(input_args) @@ -1219,7 +1227,11 @@ def main(args): if args.bnb_quantization_config_path is not None else {"device": accelerator.device, "dtype": weight_dtype} ) - transformer.to(**transformer_to_kwargs) + + is_fsdp = getattr(accelerator.state, "fsdp_plugin", None) is not None + if not is_fsdp: + transformer.to(**transformer_to_kwargs) + if args.do_fp8_training: convert_to_float8_training( transformer, module_filter_fn=module_filter_fn, config=Float8LinearConfig(pad_inner_dim=True) @@ -1263,17 +1275,42 @@ def unwrap_model(model): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format def save_model_hook(models, weights, output_dir): + transformer_cls = type(unwrap_model(transformer)) + + # 1) Validate and pick the transformer model + modules_to_save: dict[str, Any] = {} + transformer_model = None + + for model in models: + if isinstance(unwrap_model(model), transformer_cls): + transformer_model = model + modules_to_save["transformer"] = model + else: + raise ValueError(f"unexpected save model: {model.__class__}") + + if transformer_model is None: + raise ValueError("No transformer model found in 'models'") + + # 2) Optionally gather FSDP state dict once + state_dict = accelerator.get_state_dict(model) if is_fsdp else None + + # 3) Only main process materializes the LoRA state dict + transformer_lora_layers_to_save = None if accelerator.is_main_process: - transformer_lora_layers_to_save = None - modules_to_save = {} - for model in models: - if isinstance(model, type(unwrap_model(transformer))): - transformer_lora_layers_to_save = get_peft_model_state_dict(model) - modules_to_save["transformer"] = model - else: - raise ValueError(f"unexpected save model: {model.__class__}") + peft_kwargs = {} + if is_fsdp: + peft_kwargs["state_dict"] = state_dict + + transformer_lora_layers_to_save = get_peft_model_state_dict( + unwrap_model(transformer_model) if is_fsdp else transformer_model, + **peft_kwargs, + ) - # make sure to pop weight so that corresponding model is not saved again + if is_fsdp: + transformer_lora_layers_to_save = _to_cpu_contiguous(transformer_lora_layers_to_save) + + # make sure to pop weight so that corresponding model is not saved again + if weights: weights.pop() Flux2Pipeline.save_lora_weights( @@ -1285,13 +1322,20 @@ def save_model_hook(models, weights, output_dir): def load_model_hook(models, input_dir): transformer_ = None - while len(models) > 0: - model = models.pop() + if not is_fsdp: + while len(models) > 0: + model = models.pop() - if isinstance(model, type(unwrap_model(transformer))): - transformer_ = model - else: - raise ValueError(f"unexpected save model: {model.__class__}") + if isinstance(unwrap_model(model), type(unwrap_model(transformer))): + transformer_ = unwrap_model(model) + else: + raise ValueError(f"unexpected save model: {model.__class__}") + else: + transformer_ = Flux2Transformer2DModel.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="transformer", + ) + transformer_.add_adapter(transformer_lora_config) lora_state_dict = Flux2Pipeline.lora_state_dict(input_dir) @@ -1507,6 +1551,21 @@ def _encode_single(prompt: str): args.validation_prompt, text_encoding_pipeline ) + # Init FSDP for text encoder + if args.fsdp_text_encoder: + fsdp_kwargs = get_fsdp_kwargs_from_accelerator(accelerator) + text_encoder_fsdp = wrap_with_fsdp( + model=text_encoding_pipeline.text_encoder, + device=accelerator.device, + offload=args.offload, + limit_all_gathers=True, + use_orig_params=True, + fsdp_kwargs=fsdp_kwargs, + ) + + text_encoding_pipeline.text_encoder = text_encoder_fsdp + dist.barrier() + # If custom instance prompts are NOT provided (i.e. the instance prompt is used for all images), # pack the statically computed variables appropriately here. This is so that we don't # have to pass them to the dataloader. @@ -1536,6 +1595,8 @@ def _encode_single(prompt: str): if train_dataset.custom_instance_prompts: if args.remote_text_encoder: prompt_embeds, text_ids = compute_remote_text_embeddings(batch["prompts"]) + elif args.fsdp_text_encoder: + prompt_embeds, text_ids = compute_text_embeddings(batch["prompts"], text_encoding_pipeline) else: with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload): prompt_embeds, text_ids = compute_text_embeddings(batch["prompts"], text_encoding_pipeline) @@ -1777,7 +1838,7 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): progress_bar.update(1) global_step += 1 - if accelerator.is_main_process: + if accelerator.is_main_process or is_fsdp: if global_step % args.checkpointing_steps == 0: # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` if args.checkpoints_total_limit is not None: @@ -1836,15 +1897,41 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): # Save the lora layers accelerator.wait_for_everyone() + + if is_fsdp: + transformer = unwrap_model(transformer) + state_dict = accelerator.get_state_dict(transformer) if accelerator.is_main_process: modules_to_save = {} - transformer = unwrap_model(transformer) - if args.bnb_quantization_config_path is None: - if args.upcast_before_saving: - transformer.to(torch.float32) - else: - transformer = transformer.to(weight_dtype) - transformer_lora_layers = get_peft_model_state_dict(transformer) + if is_fsdp: + if args.bnb_quantization_config_path is None: + if args.upcast_before_saving: + state_dict = { + k: v.to(torch.float32) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items() + } + else: + state_dict = { + k: v.to(weight_dtype) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items() + } + + transformer_lora_layers = get_peft_model_state_dict( + transformer, + state_dict=state_dict, + ) + transformer_lora_layers = { + k: v.detach().cpu().contiguous() if isinstance(v, torch.Tensor) else v + for k, v in transformer_lora_layers.items() + } + + else: + transformer = unwrap_model(transformer) + if args.bnb_quantization_config_path is None: + if args.upcast_before_saving: + transformer.to(torch.float32) + else: + transformer = transformer.to(weight_dtype) + transformer_lora_layers = get_peft_model_state_dict(transformer) + modules_to_save["transformer"] = transformer Flux2Pipeline.save_lora_weights( diff --git a/examples/dreambooth/train_dreambooth_lora_flux2_img2img.py b/examples/dreambooth/train_dreambooth_lora_flux2_img2img.py index 0b9b9f993094..e18909e6dfd7 100644 --- a/examples/dreambooth/train_dreambooth_lora_flux2_img2img.py +++ b/examples/dreambooth/train_dreambooth_lora_flux2_img2img.py @@ -43,6 +43,7 @@ import shutil from contextlib import nullcontext from pathlib import Path +from typing import Any import numpy as np import torch @@ -74,13 +75,16 @@ from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor from diffusers.training_utils import ( _collate_lora_metadata, + _to_cpu_contiguous, cast_training_params, compute_density_for_timestep_sampling, compute_loss_weighting_for_sd3, find_nearest_bucket, free_memory, + get_fsdp_kwargs_from_accelerator, offload_models, parse_buckets_string, + wrap_with_fsdp, ) from diffusers.utils import ( check_min_version, @@ -93,11 +97,14 @@ from diffusers.utils.torch_utils import is_compiled_module +if getattr(torch, "distributed", None) is not None: + import torch.distributed as dist + if is_wandb_available(): import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) @@ -120,7 +127,7 @@ def save_model_card( ) model_description = f""" -# Flux DreamBooth LoRA - {repo_id} +# Flux.2 DreamBooth LoRA - {repo_id} @@ -339,7 +346,7 @@ def parse_args(input_args=None): "--instance_prompt", type=str, default=None, - required=True, + required=False, help="The prompt with identifier specifying the instance, e.g. 'photo of a TOK dog', 'in the style of TOK'", ) parser.add_argument( @@ -691,6 +698,7 @@ def parse_args(input_args=None): parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") parser.add_argument("--enable_npu_flash_attention", action="store_true", help="Enabla Flash Attention for NPU") + parser.add_argument("--fsdp_text_encoder", action="store_true", help="Use FSDP for text encoder") if input_args is not None: args = parser.parse_args(input_args) @@ -827,15 +835,28 @@ def __init__( dest_image = self.cond_images[i] image_width, image_height = dest_image.size if image_width * image_height > 1024 * 1024: - dest_image = Flux2ImageProcessor.image_processor._resize_to_target_area(dest_image, 1024 * 1024) + dest_image = Flux2ImageProcessor._resize_to_target_area(dest_image, 1024 * 1024) image_width, image_height = dest_image.size multiple_of = 2 ** (4 - 1) # 2 ** (len(vae.config.block_out_channels) - 1), temp! image_width = (image_width // multiple_of) * multiple_of image_height = (image_height // multiple_of) * multiple_of - dest_image = Flux2ImageProcessor.image_processor.preprocess( + image_processor = Flux2ImageProcessor() + dest_image = image_processor.preprocess( dest_image, height=image_height, width=image_width, resize_mode="crop" ) + # Convert back to PIL + dest_image = dest_image.squeeze(0) + if dest_image.min() < 0: + dest_image = (dest_image + 1) / 2 + dest_image = (torch.clamp(dest_image, 0, 1) * 255).byte().cpu() + + if dest_image.shape[0] == 1: + # Gray scale image + dest_image = Image.fromarray(dest_image.squeeze().numpy(), mode="L") + else: + # RGB scale image: (C, H, W) -> (H, W, C) + dest_image = TF.to_pil_image(dest_image) dest_image = exif_transpose(dest_image) if not dest_image.mode == "RGB": @@ -1156,7 +1177,11 @@ def main(args): if args.bnb_quantization_config_path is not None else {"device": accelerator.device, "dtype": weight_dtype} ) - transformer.to(**transformer_to_kwargs) + + is_fsdp = getattr(accelerator.state, "fsdp_plugin", None) is not None + if not is_fsdp: + transformer.to(**transformer_to_kwargs) + if args.do_fp8_training: convert_to_float8_training( transformer, module_filter_fn=module_filter_fn, config=Float8LinearConfig(pad_inner_dim=True) @@ -1200,17 +1225,42 @@ def unwrap_model(model): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format def save_model_hook(models, weights, output_dir): + transformer_cls = type(unwrap_model(transformer)) + + # 1) Validate and pick the transformer model + modules_to_save: dict[str, Any] = {} + transformer_model = None + + for model in models: + if isinstance(unwrap_model(model), transformer_cls): + transformer_model = model + modules_to_save["transformer"] = model + else: + raise ValueError(f"unexpected save model: {model.__class__}") + + if transformer_model is None: + raise ValueError("No transformer model found in 'models'") + + # 2) Optionally gather FSDP state dict once + state_dict = accelerator.get_state_dict(model) if is_fsdp else None + + # 3) Only main process materializes the LoRA state dict + transformer_lora_layers_to_save = None if accelerator.is_main_process: - transformer_lora_layers_to_save = None - modules_to_save = {} - for model in models: - if isinstance(model, type(unwrap_model(transformer))): - transformer_lora_layers_to_save = get_peft_model_state_dict(model) - modules_to_save["transformer"] = model - else: - raise ValueError(f"unexpected save model: {model.__class__}") + peft_kwargs = {} + if is_fsdp: + peft_kwargs["state_dict"] = state_dict + + transformer_lora_layers_to_save = get_peft_model_state_dict( + unwrap_model(transformer_model) if is_fsdp else transformer_model, + **peft_kwargs, + ) + + if is_fsdp: + transformer_lora_layers_to_save = _to_cpu_contiguous(transformer_lora_layers_to_save) - # make sure to pop weight so that corresponding model is not saved again + # make sure to pop weight so that corresponding model is not saved again + if weights: weights.pop() Flux2Pipeline.save_lora_weights( @@ -1222,13 +1272,20 @@ def save_model_hook(models, weights, output_dir): def load_model_hook(models, input_dir): transformer_ = None - while len(models) > 0: - model = models.pop() + if not is_fsdp: + while len(models) > 0: + model = models.pop() - if isinstance(model, type(unwrap_model(transformer))): - transformer_ = model - else: - raise ValueError(f"unexpected save model: {model.__class__}") + if isinstance(unwrap_model(model), type(unwrap_model(transformer))): + transformer_ = unwrap_model(model) + else: + raise ValueError(f"unexpected save model: {model.__class__}") + else: + transformer_ = Flux2Transformer2DModel.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="transformer", + ) + transformer_.add_adapter(transformer_lora_config) lora_state_dict = Flux2Pipeline.lora_state_dict(input_dir) @@ -1419,9 +1476,9 @@ def _encode_single(prompt: str): args.instance_prompt, text_encoding_pipeline ) - validation_image = load_image(args.validation_image_path).convert("RGB") - validation_kwargs = {"image": validation_image} if args.validation_prompt is not None: + validation_image = load_image(args.validation_image_path).convert("RGB") + validation_kwargs = {"image": validation_image} if args.remote_text_encoder: validation_kwargs["prompt_embeds"] = compute_remote_text_embeddings(args.validation_prompt) else: @@ -1430,6 +1487,21 @@ def _encode_single(prompt: str): args.validation_prompt, text_encoding_pipeline ) + # Init FSDP for text encoder + if args.fsdp_text_encoder: + fsdp_kwargs = get_fsdp_kwargs_from_accelerator(accelerator) + text_encoder_fsdp = wrap_with_fsdp( + model=text_encoding_pipeline.text_encoder, + device=accelerator.device, + offload=args.offload, + limit_all_gathers=True, + use_orig_params=True, + fsdp_kwargs=fsdp_kwargs, + ) + + text_encoding_pipeline.text_encoder = text_encoder_fsdp + dist.barrier() + # If custom instance prompts are NOT provided (i.e. the instance prompt is used for all images), # pack the statically computed variables appropriately here. This is so that we don't # have to pass them to the dataloader. @@ -1461,6 +1533,8 @@ def _encode_single(prompt: str): if train_dataset.custom_instance_prompts: if args.remote_text_encoder: prompt_embeds, text_ids = compute_remote_text_embeddings(batch["prompts"]) + elif args.fsdp_text_encoder: + prompt_embeds, text_ids = compute_text_embeddings(batch["prompts"], text_encoding_pipeline) else: with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload): prompt_embeds, text_ids = compute_text_embeddings(batch["prompts"], text_encoding_pipeline) @@ -1621,9 +1695,13 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): cond_model_input = (cond_model_input - latents_bn_mean) / latents_bn_std model_input_ids = Flux2Pipeline._prepare_latent_ids(model_input).to(device=model_input.device) - cond_model_input_ids = Flux2Pipeline._prepare_image_ids(cond_model_input).to( + cond_model_input_list = [cond_model_input[i].unsqueeze(0) for i in range(cond_model_input.shape[0])] + cond_model_input_ids = Flux2Pipeline._prepare_image_ids(cond_model_input_list).to( device=cond_model_input.device ) + cond_model_input_ids = cond_model_input_ids.view( + cond_model_input.shape[0], -1, model_input_ids.shape[-1] + ) # Sample noise that we'll add to the latents noise = torch.randn_like(model_input) @@ -1650,6 +1728,9 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): packed_noisy_model_input = Flux2Pipeline._pack_latents(noisy_model_input) packed_cond_model_input = Flux2Pipeline._pack_latents(cond_model_input) + orig_input_shape = packed_noisy_model_input.shape + orig_input_ids_shape = model_input_ids.shape + # concatenate the model inputs with the cond inputs packed_noisy_model_input = torch.cat([packed_noisy_model_input, packed_cond_model_input], dim=1) model_input_ids = torch.cat([model_input_ids, cond_model_input_ids], dim=1) @@ -1668,7 +1749,8 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): img_ids=model_input_ids, # B, image_seq_len, 4 return_dict=False, )[0] - model_pred = model_pred[:, : packed_noisy_model_input.size(1) :] + model_pred = model_pred[:, : orig_input_shape[1], :] + model_input_ids = model_input_ids[:, : orig_input_ids_shape[1], :] model_pred = Flux2Pipeline._unpack_latents_with_ids(model_pred, model_input_ids) @@ -1700,7 +1782,7 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): progress_bar.update(1) global_step += 1 - if accelerator.is_main_process: + if accelerator.is_main_process or is_fsdp: if global_step % args.checkpointing_steps == 0: # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` if args.checkpoints_total_limit is not None: @@ -1759,15 +1841,41 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): # Save the lora layers accelerator.wait_for_everyone() + + if is_fsdp: + transformer = unwrap_model(transformer) + state_dict = accelerator.get_state_dict(transformer) if accelerator.is_main_process: modules_to_save = {} - transformer = unwrap_model(transformer) - if args.bnb_quantization_config_path is None: - if args.upcast_before_saving: - transformer.to(torch.float32) - else: - transformer = transformer.to(weight_dtype) - transformer_lora_layers = get_peft_model_state_dict(transformer) + if is_fsdp: + if args.bnb_quantization_config_path is None: + if args.upcast_before_saving: + state_dict = { + k: v.to(torch.float32) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items() + } + else: + state_dict = { + k: v.to(weight_dtype) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items() + } + + transformer_lora_layers = get_peft_model_state_dict( + transformer, + state_dict=state_dict, + ) + transformer_lora_layers = { + k: v.detach().cpu().contiguous() if isinstance(v, torch.Tensor) else v + for k, v in transformer_lora_layers.items() + } + + else: + transformer = unwrap_model(transformer) + if args.bnb_quantization_config_path is None: + if args.upcast_before_saving: + transformer.to(torch.float32) + else: + transformer = transformer.to(weight_dtype) + transformer_lora_layers = get_peft_model_state_dict(transformer) + modules_to_save["transformer"] = transformer Flux2Pipeline.save_lora_weights( diff --git a/examples/dreambooth/train_dreambooth_lora_flux2_klein.py b/examples/dreambooth/train_dreambooth_lora_flux2_klein.py new file mode 100644 index 000000000000..268d0148e446 --- /dev/null +++ b/examples/dreambooth/train_dreambooth_lora_flux2_klein.py @@ -0,0 +1,1942 @@ +#!/usr/bin/env python +# coding=utf-8 +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# /// script +# dependencies = [ +# "diffusers @ git+https://github.com/huggingface/diffusers.git", +# "torch>=2.0.0", +# "accelerate>=0.31.0", +# "transformers>=4.41.2", +# "ftfy", +# "tensorboard", +# "Jinja2", +# "peft>=0.11.1", +# "sentencepiece", +# "torchvision", +# "datasets", +# "bitsandbytes", +# "prodigyopt", +# ] +# /// + +import argparse +import copy +import itertools +import json +import logging +import math +import os +import random +import shutil +import warnings +from contextlib import nullcontext +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import transformers +from accelerate import Accelerator +from accelerate.logging import get_logger +from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration, set_seed +from huggingface_hub import create_repo, upload_folder +from huggingface_hub.utils import insecure_hashlib +from peft import LoraConfig, prepare_model_for_kbit_training, set_peft_model_state_dict +from peft.utils import get_peft_model_state_dict +from PIL import Image +from PIL.ImageOps import exif_transpose +from torch.utils.data import Dataset +from torch.utils.data.sampler import BatchSampler +from torchvision import transforms +from torchvision.transforms import functional as TF +from tqdm.auto import tqdm +from transformers import Qwen2TokenizerFast, Qwen3ForCausalLM + +import diffusers +from diffusers import ( + AutoencoderKLFlux2, + BitsAndBytesConfig, + FlowMatchEulerDiscreteScheduler, + Flux2KleinPipeline, + Flux2Transformer2DModel, +) +from diffusers.optimization import get_scheduler +from diffusers.training_utils import ( + _collate_lora_metadata, + _to_cpu_contiguous, + cast_training_params, + compute_density_for_timestep_sampling, + compute_loss_weighting_for_sd3, + find_nearest_bucket, + free_memory, + get_fsdp_kwargs_from_accelerator, + offload_models, + parse_buckets_string, + wrap_with_fsdp, +) +from diffusers.utils import ( + check_min_version, + convert_unet_state_dict_to_peft, + is_wandb_available, +) +from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card +from diffusers.utils.import_utils import is_torch_npu_available +from diffusers.utils.torch_utils import is_compiled_module + + +if getattr(torch, "distributed", None) is not None: + import torch.distributed as dist + +if is_wandb_available(): + import wandb + +# Will error if the minimal version of diffusers is not installed. Remove at your own risks. +check_min_version("0.38.0.dev0") + +logger = get_logger(__name__) + + +def save_model_card( + repo_id: str, + images=None, + base_model: str = None, + instance_prompt=None, + validation_prompt=None, + repo_folder=None, + quant_training=None, +): + widget_dict = [] + if images is not None: + for i, image in enumerate(images): + image.save(os.path.join(repo_folder, f"image_{i}.png")) + widget_dict.append( + {"text": validation_prompt if validation_prompt else " ", "output": {"url": f"image_{i}.png"}} + ) + + model_description = f""" +# Flux.2 [Klein] DreamBooth LoRA - {repo_id} + + + +## Model description + +These are {repo_id} DreamBooth LoRA weights for {base_model}. + +The weights were trained using [DreamBooth](https://dreambooth.github.io/) with the [Flux2 diffusers trainer](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/README_flux2.md). + +Quant training? {quant_training} + +## Trigger words + +You should use `{instance_prompt}` to trigger the image generation. + +## Download model + +[Download the *.safetensors LoRA]({repo_id}/tree/main) in the Files & versions tab. + +## Use it with the [🧨 diffusers library](https://github.com/huggingface/diffusers) + +```py +from diffusers import AutoPipelineForText2Image +import torch +pipeline = AutoPipelineForText2Image.from_pretrained("black-forest-labs/FLUX.2", torch_dtype=torch.bfloat16).to('cuda') +pipeline.load_lora_weights('{repo_id}', weight_name='pytorch_lora_weights.safetensors') +image = pipeline('{validation_prompt if validation_prompt else instance_prompt}').images[0] +``` + +For more details, including weighting, merging and fusing LoRAs, check the [documentation on loading LoRAs in diffusers](https://huggingface.co/docs/diffusers/main/en/using-diffusers/loading_adapters) + +## License + +Please adhere to the licensing terms as described [here](https://huggingface.co/black-forest-labs/FLUX.2/blob/main/LICENSE.md). +""" + model_card = load_or_create_model_card( + repo_id_or_path=repo_id, + from_training=True, + license="other", + base_model=base_model, + prompt=instance_prompt, + model_description=model_description, + widget=widget_dict, + ) + tags = [ + "text-to-image", + "diffusers-training", + "diffusers", + "lora", + "flux2-klein", + "flux2-klein-diffusers", + "template:sd-lora", + ] + + model_card = populate_model_card(model_card, tags=tags) + model_card.save(os.path.join(repo_folder, "README.md")) + + +def log_validation( + pipeline, + args, + accelerator, + pipeline_args, + epoch, + torch_dtype, + is_final_validation=False, +): + args.num_validation_images = args.num_validation_images if args.num_validation_images else 1 + logger.info( + f"Running validation... \n Generating {args.num_validation_images} images with prompt:" + f" {args.validation_prompt}." + ) + pipeline = pipeline.to(dtype=torch_dtype) + pipeline.enable_model_cpu_offload() + pipeline.set_progress_bar_config(disable=True) + + # run inference + generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed is not None else None + autocast_ctx = torch.autocast(accelerator.device.type) if not is_final_validation else nullcontext() + + images = [] + for _ in range(args.num_validation_images): + with autocast_ctx: + image = pipeline( + prompt_embeds=pipeline_args["prompt_embeds"], + generator=generator, + ).images[0] + images.append(image) + + for tracker in accelerator.trackers: + phase_name = "test" if is_final_validation else "validation" + if tracker.name == "tensorboard": + np_images = np.stack([np.asarray(img) for img in images]) + tracker.writer.add_images(phase_name, np_images, epoch, dataformats="NHWC") + if tracker.name == "wandb": + tracker.log( + { + phase_name: [ + wandb.Image(image, caption=f"{i}: {args.validation_prompt}") for i, image in enumerate(images) + ] + } + ) + + del pipeline + free_memory() + + return images + + +def module_filter_fn(mod: torch.nn.Module, fqn: str): + # don't convert the output module + if fqn == "proj_out": + return False + # don't convert linear modules with weight dimensions not divisible by 16 + if isinstance(mod, torch.nn.Linear): + if mod.in_features % 16 != 0 or mod.out_features % 16 != 0: + return False + return True + + +def parse_args(input_args=None): + parser = argparse.ArgumentParser(description="Simple example of a training script.") + parser.add_argument( + "--pretrained_model_name_or_path", + type=str, + default=None, + required=True, + help="Path to pretrained model or model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--revision", + type=str, + default=None, + required=False, + help="Revision of pretrained model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--bnb_quantization_config_path", + type=str, + default=None, + help="Quantization config in a JSON file that will be used to define the bitsandbytes quant config of the DiT.", + ) + parser.add_argument( + "--do_fp8_training", + action="store_true", + help="if we are doing FP8 training.", + ) + parser.add_argument( + "--variant", + type=str, + default=None, + help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16", + ) + parser.add_argument( + "--dataset_name", + type=str, + default=None, + help=( + "The name of the Dataset (from the HuggingFace hub) containing the training data of instance images (could be your own, possibly private," + " dataset). It can also be a path pointing to a local copy of a dataset in your filesystem," + " or to a folder containing files that 🤗 Datasets can understand." + ), + ) + parser.add_argument( + "--dataset_config_name", + type=str, + default=None, + help="The config of the Dataset, leave as None if there's only one config.", + ) + parser.add_argument( + "--instance_data_dir", + type=str, + default=None, + help=("A folder containing the training data. "), + ) + + parser.add_argument( + "--cache_dir", + type=str, + default=None, + help="The directory where the downloaded models and datasets will be stored.", + ) + + parser.add_argument( + "--image_column", + type=str, + default="image", + help="The column of the dataset containing the target image. By " + "default, the standard Image Dataset maps out 'file_name' " + "to 'image'.", + ) + parser.add_argument( + "--caption_column", + type=str, + default=None, + help="The column of the dataset containing the instance prompt for each image", + ) + + parser.add_argument("--repeats", type=int, default=1, help="How many times to repeat the training data.") + + parser.add_argument( + "--class_data_dir", + type=str, + default=None, + required=False, + help="A folder containing the training data of class images.", + ) + parser.add_argument( + "--instance_prompt", + type=str, + default=None, + required=True, + help="The prompt with identifier specifying the instance, e.g. 'photo of a TOK dog', 'in the style of TOK'", + ) + parser.add_argument( + "--class_prompt", + type=str, + default=None, + help="The prompt to specify images in the same class as provided instance images.", + ) + parser.add_argument( + "--max_sequence_length", + type=int, + default=512, + help="Maximum sequence length to use with with the T5 text encoder", + ) + parser.add_argument( + "--text_encoder_out_layers", + type=int, + nargs="+", + default=[10, 20, 30], + help="Text encoder hidden layers to compute the final text embeddings.", + ) + parser.add_argument( + "--validation_prompt", + type=str, + default=None, + help="A prompt that is used during validation to verify that the model is learning.", + ) + parser.add_argument( + "--skip_final_inference", + default=False, + action="store_true", + help="Whether to skip the final inference step with loaded lora weights upon training completion. This will run intermediate validation inference if `validation_prompt` is provided. Specify to reduce memory.", + ) + parser.add_argument( + "--final_validation_prompt", + type=str, + default=None, + help="A prompt that is used during a final validation to verify that the model is learning. Ignored if `--validation_prompt` is provided.", + ) + parser.add_argument( + "--num_validation_images", + type=int, + default=4, + help="Number of images that should be generated during validation with `validation_prompt`.", + ) + parser.add_argument( + "--validation_epochs", + type=int, + default=50, + help=( + "Run dreambooth validation every X epochs. Dreambooth validation consists of running the prompt" + " `args.validation_prompt` multiple times: `args.num_validation_images`." + ), + ) + parser.add_argument( + "--rank", + type=int, + default=4, + help=("The dimension of the LoRA update matrices."), + ) + parser.add_argument( + "--lora_alpha", + type=int, + default=4, + help="LoRA alpha to be used for additional scaling.", + ) + parser.add_argument("--lora_dropout", type=float, default=0.0, help="Dropout probability for LoRA layers") + + parser.add_argument( + "--with_prior_preservation", + default=False, + action="store_true", + help="Flag to add prior preservation loss.", + ) + parser.add_argument("--prior_loss_weight", type=float, default=1.0, help="The weight of prior preservation loss.") + parser.add_argument( + "--num_class_images", + type=int, + default=100, + help=( + "Minimal class images for prior preservation loss. If there are not enough images already present in" + " class_data_dir, additional images will be sampled with class_prompt." + ), + ) + parser.add_argument( + "--output_dir", + type=str, + default="flux-dreambooth-lora", + help="The output directory where the model predictions and checkpoints will be written.", + ) + parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") + parser.add_argument( + "--resolution", + type=int, + default=512, + help=( + "The resolution for input images, all the images in the train/validation dataset will be resized to this" + " resolution" + ), + ) + parser.add_argument( + "--aspect_ratio_buckets", + type=str, + default=None, + help=( + "Aspect ratio buckets to use for training. Define as a string of 'h1,w1;h2,w2;...'. " + "e.g. '1024,1024;768,1360;1360,768;880,1168;1168,880;1248,832;832,1248'" + "Images will be resized and cropped to fit the nearest bucket. If provided, --resolution is ignored." + ), + ) + parser.add_argument( + "--center_crop", + default=False, + action="store_true", + help=( + "Whether to center crop the input images to the resolution. If not set, the images will be randomly" + " cropped. The images will be resized to the resolution first before cropping." + ), + ) + parser.add_argument( + "--random_flip", + action="store_true", + help="whether to randomly flip images horizontally", + ) + parser.add_argument( + "--train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader." + ) + parser.add_argument( + "--sample_batch_size", type=int, default=4, help="Batch size (per device) for sampling images." + ) + parser.add_argument("--num_train_epochs", type=int, default=1) + parser.add_argument( + "--max_train_steps", + type=int, + default=None, + help="Total number of training steps to perform. If provided, overrides num_train_epochs.", + ) + parser.add_argument( + "--checkpointing_steps", + type=int, + default=500, + help=( + "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final" + " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming" + " training using `--resume_from_checkpoint`." + ), + ) + parser.add_argument( + "--checkpoints_total_limit", + type=int, + default=None, + help=("Max number of checkpoints to store."), + ) + parser.add_argument( + "--resume_from_checkpoint", + type=str, + default=None, + help=( + "Whether training should be resumed from a previous checkpoint. Use a path saved by" + ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' + ), + ) + parser.add_argument( + "--gradient_accumulation_steps", + type=int, + default=1, + help="Number of updates steps to accumulate before performing a backward/update pass.", + ) + parser.add_argument( + "--gradient_checkpointing", + action="store_true", + help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.", + ) + parser.add_argument( + "--learning_rate", + type=float, + default=1e-4, + help="Initial learning rate (after the potential warmup period) to use.", + ) + + parser.add_argument( + "--guidance_scale", + type=float, + default=3.5, + help="the FLUX.1 dev variant is a guidance distilled model", + ) + + parser.add_argument( + "--text_encoder_lr", + type=float, + default=5e-6, + help="Text encoder learning rate to use.", + ) + parser.add_argument( + "--scale_lr", + action="store_true", + default=False, + help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.", + ) + parser.add_argument( + "--lr_scheduler", + type=str, + default="constant", + help=( + 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' + ' "constant", "constant_with_warmup"]' + ), + ) + parser.add_argument( + "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler." + ) + parser.add_argument( + "--lr_num_cycles", + type=int, + default=1, + help="Number of hard resets of the lr in cosine_with_restarts scheduler.", + ) + parser.add_argument("--lr_power", type=float, default=1.0, help="Power factor of the polynomial scheduler.") + parser.add_argument( + "--dataloader_num_workers", + type=int, + default=0, + help=( + "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process." + ), + ) + parser.add_argument( + "--weighting_scheme", + type=str, + default="none", + choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "none"], + help=('We default to the "none" weighting scheme for uniform sampling and uniform loss'), + ) + parser.add_argument( + "--logit_mean", type=float, default=0.0, help="mean to use when using the `'logit_normal'` weighting scheme." + ) + parser.add_argument( + "--logit_std", type=float, default=1.0, help="std to use when using the `'logit_normal'` weighting scheme." + ) + parser.add_argument( + "--mode_scale", + type=float, + default=1.29, + help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.", + ) + parser.add_argument( + "--optimizer", + type=str, + default="AdamW", + help=('The optimizer type to use. Choose between ["AdamW", "prodigy"]'), + ) + + parser.add_argument( + "--use_8bit_adam", + action="store_true", + help="Whether or not to use 8-bit Adam from bitsandbytes. Ignored if optimizer is not set to AdamW", + ) + + parser.add_argument( + "--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam and Prodigy optimizers." + ) + parser.add_argument( + "--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam and Prodigy optimizers." + ) + parser.add_argument( + "--prodigy_beta3", + type=float, + default=None, + help="coefficients for computing the Prodigy stepsize using running averages. If set to None, " + "uses the value of square root of beta2. Ignored if optimizer is adamW", + ) + parser.add_argument("--prodigy_decouple", type=bool, default=True, help="Use AdamW style decoupled weight decay") + parser.add_argument("--adam_weight_decay", type=float, default=1e-04, help="Weight decay to use for unet params") + parser.add_argument( + "--adam_weight_decay_text_encoder", type=float, default=1e-03, help="Weight decay to use for text_encoder" + ) + + parser.add_argument( + "--lora_layers", + type=str, + default=None, + help=( + 'The transformer modules to apply LoRA training on. Please specify the layers in a comma separated. E.g. - "to_k,to_q,to_v,to_out.0" will result in lora training of attention layers only' + ), + ) + + parser.add_argument( + "--adam_epsilon", + type=float, + default=1e-08, + help="Epsilon value for the Adam optimizer and Prodigy optimizers.", + ) + + parser.add_argument( + "--prodigy_use_bias_correction", + type=bool, + default=True, + help="Turn on Adam's bias correction. True by default. Ignored if optimizer is adamW", + ) + parser.add_argument( + "--prodigy_safeguard_warmup", + type=bool, + default=True, + help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage. True by default. " + "Ignored if optimizer is adamW", + ) + parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") + parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.") + parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.") + parser.add_argument( + "--hub_model_id", + type=str, + default=None, + help="The name of the repository to keep in sync with the local `output_dir`.", + ) + parser.add_argument( + "--logging_dir", + type=str, + default="logs", + help=( + "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to" + " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***." + ), + ) + parser.add_argument( + "--allow_tf32", + action="store_true", + help=( + "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" + " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" + ), + ) + parser.add_argument( + "--cache_latents", + action="store_true", + default=False, + help="Cache the VAE latents", + ) + parser.add_argument( + "--report_to", + type=str, + default="tensorboard", + help=( + 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' + ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' + ), + ) + parser.add_argument( + "--mixed_precision", + type=str, + default=None, + choices=["no", "fp16", "bf16"], + help=( + "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" + " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" + " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." + ), + ) + parser.add_argument( + "--upcast_before_saving", + action="store_true", + default=False, + help=( + "Whether to upcast the trained transformer layers to float32 before saving (at the end of training). " + "Defaults to precision dtype used for training to save memory" + ), + ) + parser.add_argument( + "--offload", + action="store_true", + help="Whether to offload the VAE and the text encoder to CPU when they are not used.", + ) + parser.add_argument( + "--prior_generation_precision", + type=str, + default=None, + choices=["no", "fp32", "fp16", "bf16"], + help=( + "Choose prior generation precision between fp32, fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" + " 1.10.and an Nvidia Ampere GPU. Default to fp16 if a GPU is available else fp32." + ), + ) + parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") + parser.add_argument("--enable_npu_flash_attention", action="store_true", help="Enabla Flash Attention for NPU") + parser.add_argument("--fsdp_text_encoder", action="store_true", help="Use FSDP for text encoder") + + if input_args is not None: + args = parser.parse_args(input_args) + else: + args = parser.parse_args() + + if args.dataset_name is None and args.instance_data_dir is None: + raise ValueError("Specify either `--dataset_name` or `--instance_data_dir`") + + if args.dataset_name is not None and args.instance_data_dir is not None: + raise ValueError("Specify only one of `--dataset_name` or `--instance_data_dir`") + if args.do_fp8_training and args.bnb_quantization_config_path: + raise ValueError("Both `do_fp8_training` and `bnb_quantization_config_path` cannot be passed.") + + env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) + if env_local_rank != -1 and env_local_rank != args.local_rank: + args.local_rank = env_local_rank + + if args.with_prior_preservation: + if args.class_data_dir is None: + raise ValueError("You must specify a data directory for class images.") + if args.class_prompt is None: + raise ValueError("You must specify prompt for class images.") + else: + # logger is not available yet + if args.class_data_dir is not None: + warnings.warn("You need not use --class_data_dir without --with_prior_preservation.") + if args.class_prompt is not None: + warnings.warn("You need not use --class_prompt without --with_prior_preservation.") + + return args + + +class DreamBoothDataset(Dataset): + """ + A dataset to prepare the instance and class images with the prompts for fine-tuning the model. + It pre-processes the images. + """ + + def __init__( + self, + instance_data_root, + instance_prompt, + class_prompt, + class_data_root=None, + class_num=None, + size=1024, + repeats=1, + center_crop=False, + buckets=None, + ): + self.size = size + self.center_crop = center_crop + + self.instance_prompt = instance_prompt + self.custom_instance_prompts = None + self.class_prompt = class_prompt + + self.buckets = buckets + + # if --dataset_name is provided or a metadata jsonl file is provided in the local --instance_data directory, + # we load the training data using load_dataset + if args.dataset_name is not None: + try: + from datasets import load_dataset + except ImportError: + raise ImportError( + "You are trying to load your data using the datasets library. If you wish to train using custom " + "captions please install the datasets library: `pip install datasets`. If you wish to load a " + "local folder containing images only, specify --instance_data_dir instead." + ) + # Downloading and loading a dataset from the hub. + # See more about loading custom images at + # https://huggingface.co/docs/datasets/v2.0.0/en/dataset_script + dataset = load_dataset( + args.dataset_name, + args.dataset_config_name, + cache_dir=args.cache_dir, + ) + # Preprocessing the datasets. + column_names = dataset["train"].column_names + + # 6. Get the column names for input/target. + if args.image_column is None: + image_column = column_names[0] + logger.info(f"image column defaulting to {image_column}") + else: + image_column = args.image_column + if image_column not in column_names: + raise ValueError( + f"`--image_column` value '{args.image_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}" + ) + instance_images = dataset["train"][image_column] + + if args.caption_column is None: + logger.info( + "No caption column provided, defaulting to instance_prompt for all images. If your dataset " + "contains captions/prompts for the images, make sure to specify the " + "column as --caption_column" + ) + self.custom_instance_prompts = None + else: + if args.caption_column not in column_names: + raise ValueError( + f"`--caption_column` value '{args.caption_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}" + ) + custom_instance_prompts = dataset["train"][args.caption_column] + # create final list of captions according to --repeats + self.custom_instance_prompts = [] + for caption in custom_instance_prompts: + self.custom_instance_prompts.extend(itertools.repeat(caption, repeats)) + else: + self.instance_data_root = Path(instance_data_root) + if not self.instance_data_root.exists(): + raise ValueError("Instance images root doesn't exists.") + + instance_images = [Image.open(path) for path in list(Path(instance_data_root).iterdir())] + self.custom_instance_prompts = None + + self.instance_images = [] + for img in instance_images: + self.instance_images.extend(itertools.repeat(img, repeats)) + + self.pixel_values = [] + for i, image in enumerate(self.instance_images): + image = exif_transpose(image) + if not image.mode == "RGB": + image = image.convert("RGB") + + width, height = image.size + + # Find the closest bucket + bucket_idx = find_nearest_bucket(height, width, self.buckets) + target_height, target_width = self.buckets[bucket_idx] + self.size = (target_height, target_width) + + # based on the bucket assignment, define the transformations + image = self.train_transform( + image, + size=self.size, + center_crop=args.center_crop, + random_flip=args.random_flip, + ) + self.pixel_values.append((image, bucket_idx)) + + self.num_instance_images = len(self.instance_images) + self._length = self.num_instance_images + + if class_data_root is not None: + self.class_data_root = Path(class_data_root) + self.class_data_root.mkdir(parents=True, exist_ok=True) + self.class_images_path = list(self.class_data_root.iterdir()) + if class_num is not None: + self.num_class_images = min(len(self.class_images_path), class_num) + else: + self.num_class_images = len(self.class_images_path) + self._length = max(self.num_class_images, self.num_instance_images) + else: + self.class_data_root = None + + self.image_transforms = transforms.Compose( + [ + transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR), + transforms.CenterCrop(size) if center_crop else transforms.RandomCrop(size), + transforms.ToTensor(), + transforms.Normalize([0.5], [0.5]), + ] + ) + + def __len__(self): + return self._length + + def __getitem__(self, index): + example = {} + instance_image, bucket_idx = self.pixel_values[index % self.num_instance_images] + example["instance_images"] = instance_image + example["bucket_idx"] = bucket_idx + if self.custom_instance_prompts: + caption = self.custom_instance_prompts[index % self.num_instance_images] + if caption: + example["instance_prompt"] = caption + else: + example["instance_prompt"] = self.instance_prompt + + else: # custom prompts were provided, but length does not match size of image dataset + example["instance_prompt"] = self.instance_prompt + + if self.class_data_root: + class_image = Image.open(self.class_images_path[index % self.num_class_images]) + class_image = exif_transpose(class_image) + + if not class_image.mode == "RGB": + class_image = class_image.convert("RGB") + example["class_images"] = self.image_transforms(class_image) + example["class_prompt"] = self.class_prompt + + return example + + def train_transform(self, image, size=(224, 224), center_crop=False, random_flip=False): + # 1. Resize (deterministic) + resize = transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR) + image = resize(image) + + # 2. Crop: either center or SAME random crop + if center_crop: + crop = transforms.CenterCrop(size) + image = crop(image) + else: + # get_params returns (i, j, h, w) + i, j, h, w = transforms.RandomCrop.get_params(image, output_size=size) + image = TF.crop(image, i, j, h, w) + + # 3. Random horizontal flip with the SAME coin flip + if random_flip: + do_flip = random.random() < 0.5 + if do_flip: + image = TF.hflip(image) + + # 4. ToTensor + Normalize (deterministic) + to_tensor = transforms.ToTensor() + normalize = transforms.Normalize([0.5], [0.5]) + image = normalize(to_tensor(image)) + + return image + + +def collate_fn(examples, with_prior_preservation=False): + pixel_values = [example["instance_images"] for example in examples] + prompts = [example["instance_prompt"] for example in examples] + + # Concat class and instance examples for prior preservation. + # We do this to avoid doing two forward passes. + if with_prior_preservation: + pixel_values += [example["class_images"] for example in examples] + prompts += [example["class_prompt"] for example in examples] + + pixel_values = torch.stack(pixel_values) + pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float() + + batch = {"pixel_values": pixel_values, "prompts": prompts} + return batch + + +class BucketBatchSampler(BatchSampler): + def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False): + if not isinstance(batch_size, int) or batch_size <= 0: + raise ValueError("batch_size should be a positive integer value, but got batch_size={}".format(batch_size)) + if not isinstance(drop_last, bool): + raise ValueError("drop_last should be a boolean value, but got drop_last={}".format(drop_last)) + + self.dataset = dataset + self.batch_size = batch_size + self.drop_last = drop_last + + # Group indices by bucket + self.bucket_indices = [[] for _ in range(len(self.dataset.buckets))] + for idx, (_, bucket_idx) in enumerate(self.dataset.pixel_values): + self.bucket_indices[bucket_idx].append(idx) + + self.sampler_len = 0 + self.batches = [] + + # Pre-generate batches for each bucket + for indices_in_bucket in self.bucket_indices: + # Shuffle indices within the bucket + random.shuffle(indices_in_bucket) + # Create batches + for i in range(0, len(indices_in_bucket), self.batch_size): + batch = indices_in_bucket[i : i + self.batch_size] + if len(batch) < self.batch_size and self.drop_last: + continue # Skip partial batch if drop_last is True + self.batches.append(batch) + self.sampler_len += 1 # Count the number of batches + + def __iter__(self): + # Shuffle the order of the batches each epoch + random.shuffle(self.batches) + for batch in self.batches: + yield batch + + def __len__(self): + return self.sampler_len + + +class PromptDataset(Dataset): + "A simple dataset to prepare the prompts to generate class images on multiple GPUs." + + def __init__(self, prompt, num_samples): + self.prompt = prompt + self.num_samples = num_samples + + def __len__(self): + return self.num_samples + + def __getitem__(self, index): + example = {} + example["prompt"] = self.prompt + example["index"] = index + return example + + +def main(args): + if args.report_to == "wandb" and args.hub_token is not None: + raise ValueError( + "You cannot use both --report_to=wandb and --hub_token due to a security risk of exposing your token." + " Please use `hf auth login` to authenticate with the Hub." + ) + + if torch.backends.mps.is_available() and args.mixed_precision == "bf16": + # due to pytorch#99272, MPS does not yet support bfloat16. + raise ValueError( + "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead." + ) + if args.do_fp8_training: + from torchao.float8 import Float8LinearConfig, convert_to_float8_training + + logging_dir = Path(args.output_dir, args.logging_dir) + + accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) + kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) + accelerator = Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + mixed_precision=args.mixed_precision, + log_with=args.report_to, + project_config=accelerator_project_config, + kwargs_handlers=[kwargs], + ) + + # Disable AMP for MPS. + if torch.backends.mps.is_available(): + accelerator.native_amp = False + + if args.report_to == "wandb": + if not is_wandb_available(): + raise ImportError("Make sure to install wandb if you want to use it for logging during training.") + + # Make one log on every process with the configuration for debugging. + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + ) + logger.info(accelerator.state, main_process_only=False) + if accelerator.is_local_main_process: + transformers.utils.logging.set_verbosity_warning() + diffusers.utils.logging.set_verbosity_info() + else: + transformers.utils.logging.set_verbosity_error() + diffusers.utils.logging.set_verbosity_error() + + # If passed along, set the training seed now. + if args.seed is not None: + set_seed(args.seed) + + # Generate class images if prior preservation is enabled. + if args.with_prior_preservation: + class_images_dir = Path(args.class_data_dir) + if not class_images_dir.exists(): + class_images_dir.mkdir(parents=True) + cur_class_images = len(list(class_images_dir.iterdir())) + + if cur_class_images < args.num_class_images: + has_supported_fp16_accelerator = torch.cuda.is_available() or torch.backends.mps.is_available() + torch_dtype = torch.float16 if has_supported_fp16_accelerator else torch.float32 + if args.prior_generation_precision == "fp32": + torch_dtype = torch.float32 + elif args.prior_generation_precision == "fp16": + torch_dtype = torch.float16 + elif args.prior_generation_precision == "bf16": + torch_dtype = torch.bfloat16 + + pipeline = Flux2KleinPipeline.from_pretrained( + args.pretrained_model_name_or_path, + torch_dtype=torch_dtype, + revision=args.revision, + variant=args.variant, + ) + pipeline.set_progress_bar_config(disable=True) + + num_new_images = args.num_class_images - cur_class_images + logger.info(f"Number of class images to sample: {num_new_images}.") + + sample_dataset = PromptDataset(args.class_prompt, num_new_images) + sample_dataloader = torch.utils.data.DataLoader(sample_dataset, batch_size=args.sample_batch_size) + + sample_dataloader = accelerator.prepare(sample_dataloader) + pipeline.to(accelerator.device) + + for example in tqdm( + sample_dataloader, desc="Generating class images", disable=not accelerator.is_local_main_process + ): + with torch.autocast(device_type=accelerator.device.type, dtype=torch_dtype): + images = pipeline(prompt=example["prompt"]).images + + for i, image in enumerate(images): + hash_image = insecure_hashlib.sha1(image.tobytes()).hexdigest() + image_filename = class_images_dir / f"{example['index'][i] + cur_class_images}-{hash_image}.jpg" + image.save(image_filename) + + del pipeline + free_memory() + + # Handle the repository creation + if accelerator.is_main_process: + if args.output_dir is not None: + os.makedirs(args.output_dir, exist_ok=True) + + if args.push_to_hub: + repo_id = create_repo( + repo_id=args.hub_model_id or Path(args.output_dir).name, + exist_ok=True, + ).repo_id + + # Load the tokenizers + tokenizer = Qwen2TokenizerFast.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="tokenizer", + revision=args.revision, + ) + + # For mixed precision training we cast all non-trainable weights (vae, text_encoder and transformer) to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + weight_dtype = torch.float32 + if accelerator.mixed_precision == "fp16": + weight_dtype = torch.float16 + elif accelerator.mixed_precision == "bf16": + weight_dtype = torch.bfloat16 + + # Load scheduler and models + noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="scheduler", + revision=args.revision, + ) + noise_scheduler_copy = copy.deepcopy(noise_scheduler) + vae = AutoencoderKLFlux2.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="vae", + revision=args.revision, + variant=args.variant, + ) + latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(accelerator.device) + latents_bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to( + accelerator.device + ) + + quantization_config = None + if args.bnb_quantization_config_path is not None: + with open(args.bnb_quantization_config_path, "r") as f: + config_kwargs = json.load(f) + if "load_in_4bit" in config_kwargs and config_kwargs["load_in_4bit"]: + config_kwargs["bnb_4bit_compute_dtype"] = weight_dtype + quantization_config = BitsAndBytesConfig(**config_kwargs) + + transformer = Flux2Transformer2DModel.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="transformer", + revision=args.revision, + variant=args.variant, + quantization_config=quantization_config, + torch_dtype=weight_dtype, + ) + if args.bnb_quantization_config_path is not None: + transformer = prepare_model_for_kbit_training(transformer, use_gradient_checkpointing=False) + + text_encoder = Qwen3ForCausalLM.from_pretrained( + args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision, variant=args.variant + ) + text_encoder.requires_grad_(False) + + # We only train the additional adapter LoRA layers + transformer.requires_grad_(False) + vae.requires_grad_(False) + + if args.enable_npu_flash_attention: + if is_torch_npu_available(): + logger.info("npu flash attention enabled.") + transformer.set_attention_backend("_native_npu") + else: + raise ValueError("npu flash attention requires torch_npu extensions and is supported only on npu device ") + + if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16: + # due to pytorch#99272, MPS does not yet support bfloat16. + raise ValueError( + "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead." + ) + + to_kwargs = {"dtype": weight_dtype, "device": accelerator.device} if not args.offload else {"dtype": weight_dtype} + # flux vae is stable in bf16 so load it in weight_dtype to reduce memory + vae.to(**to_kwargs) + # we never offload the transformer to CPU, so we can just use the accelerator device + transformer_to_kwargs = ( + {"device": accelerator.device} + if args.bnb_quantization_config_path is not None + else {"device": accelerator.device, "dtype": weight_dtype} + ) + + is_fsdp = getattr(accelerator.state, "fsdp_plugin", None) is not None + if not is_fsdp: + transformer.to(**transformer_to_kwargs) + + if args.do_fp8_training: + convert_to_float8_training( + transformer, module_filter_fn=module_filter_fn, config=Float8LinearConfig(pad_inner_dim=True) + ) + + text_encoder.to(**to_kwargs) + # Initialize a text encoding pipeline and keep it to CPU for now. + text_encoding_pipeline = Flux2KleinPipeline.from_pretrained( + args.pretrained_model_name_or_path, + vae=None, + transformer=None, + tokenizer=tokenizer, + text_encoder=text_encoder, + scheduler=None, + revision=args.revision, + ) + + if args.gradient_checkpointing: + transformer.enable_gradient_checkpointing() + + if args.lora_layers is not None: + target_modules = [layer.strip() for layer in args.lora_layers.split(",")] + else: + target_modules = ["to_k", "to_q", "to_v", "to_out.0"] + + # now we will add new LoRA weights the transformer layers + transformer_lora_config = LoraConfig( + r=args.rank, + lora_alpha=args.lora_alpha, + lora_dropout=args.lora_dropout, + init_lora_weights="gaussian", + target_modules=target_modules, + ) + transformer.add_adapter(transformer_lora_config) + + def unwrap_model(model): + model = accelerator.unwrap_model(model) + model = model._orig_mod if is_compiled_module(model) else model + return model + + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + def save_model_hook(models, weights, output_dir): + transformer_cls = type(unwrap_model(transformer)) + + # 1) Validate and pick the transformer model + modules_to_save: dict[str, Any] = {} + transformer_model = None + + for model in models: + if isinstance(unwrap_model(model), transformer_cls): + transformer_model = model + modules_to_save["transformer"] = model + else: + raise ValueError(f"unexpected save model: {model.__class__}") + + if transformer_model is None: + raise ValueError("No transformer model found in 'models'") + + # 2) Optionally gather FSDP state dict once + state_dict = accelerator.get_state_dict(model) if is_fsdp else None + + # 3) Only main process materializes the LoRA state dict + transformer_lora_layers_to_save = None + if accelerator.is_main_process: + peft_kwargs = {} + if is_fsdp: + peft_kwargs["state_dict"] = state_dict + + transformer_lora_layers_to_save = get_peft_model_state_dict( + unwrap_model(transformer_model) if is_fsdp else transformer_model, + **peft_kwargs, + ) + + if is_fsdp: + transformer_lora_layers_to_save = _to_cpu_contiguous(transformer_lora_layers_to_save) + + # make sure to pop weight so that corresponding model is not saved again + if weights: + weights.pop() + + Flux2KleinPipeline.save_lora_weights( + output_dir, + transformer_lora_layers=transformer_lora_layers_to_save, + **_collate_lora_metadata(modules_to_save), + ) + + def load_model_hook(models, input_dir): + transformer_ = None + + if not is_fsdp: + while len(models) > 0: + model = models.pop() + + if isinstance(unwrap_model(model), type(unwrap_model(transformer))): + transformer_ = unwrap_model(model) + else: + raise ValueError(f"unexpected save model: {model.__class__}") + else: + transformer_ = Flux2Transformer2DModel.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="transformer", + ) + transformer_.add_adapter(transformer_lora_config) + + lora_state_dict = Flux2KleinPipeline.lora_state_dict(input_dir) + + transformer_state_dict = { + f"{k.replace('transformer.', '')}": v for k, v in lora_state_dict.items() if k.startswith("transformer.") + } + transformer_state_dict = convert_unet_state_dict_to_peft(transformer_state_dict) + incompatible_keys = set_peft_model_state_dict(transformer_, transformer_state_dict, adapter_name="default") + if incompatible_keys is not None: + # check only for unexpected keys + unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None) + if unexpected_keys: + logger.warning( + f"Loading adapter weights from state_dict led to unexpected keys not found in the model: " + f" {unexpected_keys}. " + ) + + # Make sure the trainable params are in float32. This is again needed since the base models + # are in `weight_dtype`. More details: + # https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804 + if args.mixed_precision == "fp16": + models = [transformer_] + # only upcast trainable parameters (LoRA) into fp32 + cast_training_params(models) + + accelerator.register_save_state_pre_hook(save_model_hook) + accelerator.register_load_state_pre_hook(load_model_hook) + + # Enable TF32 for faster training on Ampere GPUs, + # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices + if args.allow_tf32 and torch.cuda.is_available(): + torch.backends.cuda.matmul.allow_tf32 = True + + if args.scale_lr: + args.learning_rate = ( + args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes + ) + + # Make sure the trainable params are in float32. + if args.mixed_precision == "fp16": + models = [transformer] + # only upcast trainable parameters (LoRA) into fp32 + cast_training_params(models, dtype=torch.float32) + + transformer_lora_parameters = list(filter(lambda p: p.requires_grad, transformer.parameters())) + + # Optimization parameters + transformer_parameters_with_lr = {"params": transformer_lora_parameters, "lr": args.learning_rate} + params_to_optimize = [transformer_parameters_with_lr] + + # Optimizer creation + if not (args.optimizer.lower() == "prodigy" or args.optimizer.lower() == "adamw"): + logger.warning( + f"Unsupported choice of optimizer: {args.optimizer}.Supported optimizers include [adamW, prodigy]." + "Defaulting to adamW" + ) + args.optimizer = "adamw" + + if args.use_8bit_adam and not args.optimizer.lower() == "adamw": + logger.warning( + f"use_8bit_adam is ignored when optimizer is not set to 'AdamW'. Optimizer was " + f"set to {args.optimizer.lower()}" + ) + + if args.optimizer.lower() == "adamw": + if args.use_8bit_adam: + try: + import bitsandbytes as bnb + except ImportError: + raise ImportError( + "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`." + ) + + optimizer_class = bnb.optim.AdamW8bit + else: + optimizer_class = torch.optim.AdamW + + optimizer = optimizer_class( + params_to_optimize, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + + if args.optimizer.lower() == "prodigy": + try: + import prodigyopt + except ImportError: + raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`") + + optimizer_class = prodigyopt.Prodigy + + if args.learning_rate <= 0.1: + logger.warning( + "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0" + ) + + optimizer = optimizer_class( + params_to_optimize, + betas=(args.adam_beta1, args.adam_beta2), + beta3=args.prodigy_beta3, + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + decouple=args.prodigy_decouple, + use_bias_correction=args.prodigy_use_bias_correction, + safeguard_warmup=args.prodigy_safeguard_warmup, + ) + + if args.aspect_ratio_buckets is not None: + buckets = parse_buckets_string(args.aspect_ratio_buckets) + else: + buckets = [(args.resolution, args.resolution)] + logger.info(f"Using parsed aspect ratio buckets: {buckets}") + + # Dataset and DataLoaders creation: + train_dataset = DreamBoothDataset( + instance_data_root=args.instance_data_dir, + instance_prompt=args.instance_prompt, + class_prompt=args.class_prompt, + class_data_root=args.class_data_dir if args.with_prior_preservation else None, + class_num=args.num_class_images, + size=args.resolution, + repeats=args.repeats, + center_crop=args.center_crop, + buckets=buckets, + ) + batch_sampler = BucketBatchSampler(train_dataset, batch_size=args.train_batch_size, drop_last=True) + train_dataloader = torch.utils.data.DataLoader( + train_dataset, + batch_sampler=batch_sampler, + collate_fn=lambda examples: collate_fn(examples, args.with_prior_preservation), + num_workers=args.dataloader_num_workers, + ) + + def compute_text_embeddings(prompt, text_encoding_pipeline): + with torch.no_grad(): + prompt_embeds, text_ids = text_encoding_pipeline.encode_prompt( + prompt=prompt, + max_sequence_length=args.max_sequence_length, + text_encoder_out_layers=args.text_encoder_out_layers, + ) + return prompt_embeds, text_ids + + # If no type of tuning is done on the text_encoder and custom instance prompts are NOT + # provided (i.e. the --instance_prompt is used for all images), we encode the instance prompt once to avoid + # the redundant encoding. + if not train_dataset.custom_instance_prompts: + with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload): + instance_prompt_hidden_states, instance_text_ids = compute_text_embeddings( + args.instance_prompt, text_encoding_pipeline + ) + + # Handle class prompt for prior-preservation. + if args.with_prior_preservation: + with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload): + class_prompt_hidden_states, class_text_ids = compute_text_embeddings( + args.class_prompt, text_encoding_pipeline + ) + validation_embeddings = {} + if args.validation_prompt is not None: + with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload): + (validation_embeddings["prompt_embeds"], validation_embeddings["text_ids"]) = compute_text_embeddings( + args.validation_prompt, text_encoding_pipeline + ) + + # Init FSDP for text encoder + if args.fsdp_text_encoder: + fsdp_kwargs = get_fsdp_kwargs_from_accelerator(accelerator) + text_encoder_fsdp = wrap_with_fsdp( + model=text_encoding_pipeline.text_encoder, + device=accelerator.device, + offload=args.offload, + limit_all_gathers=True, + use_orig_params=True, + fsdp_kwargs=fsdp_kwargs, + ) + + text_encoding_pipeline.text_encoder = text_encoder_fsdp + dist.barrier() + + # If custom instance prompts are NOT provided (i.e. the instance prompt is used for all images), + # pack the statically computed variables appropriately here. This is so that we don't + # have to pass them to the dataloader. + if not train_dataset.custom_instance_prompts: + prompt_embeds = instance_prompt_hidden_states + text_ids = instance_text_ids + if args.with_prior_preservation: + prompt_embeds = torch.cat([prompt_embeds, class_prompt_hidden_states], dim=0) + text_ids = torch.cat([text_ids, class_text_ids], dim=0) + + # if cache_latents is set to True, we encode images to latents and store them. + # Similar to pre-encoding in the case of a single instance prompt, if custom prompts are provided + # we encode them in advance as well. + precompute_latents = args.cache_latents or train_dataset.custom_instance_prompts + if precompute_latents: + prompt_embeds_cache = [] + text_ids_cache = [] + latents_cache = [] + for batch in tqdm(train_dataloader, desc="Caching latents"): + with torch.no_grad(): + if args.cache_latents: + with offload_models(vae, device=accelerator.device, offload=args.offload): + batch["pixel_values"] = batch["pixel_values"].to( + accelerator.device, non_blocking=True, dtype=vae.dtype + ) + latents_cache.append(vae.encode(batch["pixel_values"]).latent_dist) + if train_dataset.custom_instance_prompts: + if args.fsdp_text_encoder: + prompt_embeds, text_ids = compute_text_embeddings(batch["prompts"], text_encoding_pipeline) + else: + with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload): + prompt_embeds, text_ids = compute_text_embeddings(batch["prompts"], text_encoding_pipeline) + prompt_embeds_cache.append(prompt_embeds) + text_ids_cache.append(text_ids) + + # move back to cpu before deleting to ensure memory is freed see: https://github.com/huggingface/diffusers/issues/11376#issue-3008144624 + if args.cache_latents: + vae = vae.to("cpu") + del vae + + # move back to cpu before deleting to ensure memory is freed see: https://github.com/huggingface/diffusers/issues/11376#issue-3008144624 + text_encoding_pipeline = text_encoding_pipeline.to("cpu") + del text_encoder, tokenizer + free_memory() + + # Scheduler and math around the number of training steps. + # Check the PR https://github.com/huggingface/diffusers/pull/8312 for detailed explanation. + num_warmup_steps_for_scheduler = args.lr_warmup_steps * accelerator.num_processes + if args.max_train_steps is None: + len_train_dataloader_after_sharding = math.ceil(len(train_dataloader) / accelerator.num_processes) + num_update_steps_per_epoch = math.ceil(len_train_dataloader_after_sharding / args.gradient_accumulation_steps) + num_training_steps_for_scheduler = ( + args.num_train_epochs * accelerator.num_processes * num_update_steps_per_epoch + ) + else: + num_training_steps_for_scheduler = args.max_train_steps * accelerator.num_processes + + lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=optimizer, + num_warmup_steps=num_warmup_steps_for_scheduler, + num_training_steps=num_training_steps_for_scheduler, + num_cycles=args.lr_num_cycles, + power=args.lr_power, + ) + + # Prepare everything with our `accelerator`. + transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( + transformer, optimizer, train_dataloader, lr_scheduler + ) + + # We need to recalculate our total training steps as the size of the training dataloader may have changed. + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if args.max_train_steps is None: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + if num_training_steps_for_scheduler != args.max_train_steps: + logger.warning( + f"The length of the 'train_dataloader' after 'accelerator.prepare' ({len(train_dataloader)}) does not match " + f"the expected length ({len_train_dataloader_after_sharding}) when the learning rate scheduler was created. " + f"This inconsistency may result in the learning rate scheduler not functioning properly." + ) + # Afterwards we recalculate our number of training epochs + args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + + # We need to initialize the trackers we use, and also store our configuration. + # The trackers initializes automatically on the main process. + if accelerator.is_main_process: + tracker_name = "dreambooth-flux2-klein-lora" + args_cp = vars(args).copy() + args_cp["text_encoder_out_layers"] = str(args_cp["text_encoder_out_layers"]) + accelerator.init_trackers(tracker_name, config=args_cp) + + # Train! + total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps + + logger.info("***** Running training *****") + logger.info(f" Num examples = {len(train_dataset)}") + logger.info(f" Num batches each epoch = {len(train_dataloader)}") + logger.info(f" Num Epochs = {args.num_train_epochs}") + logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") + logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") + logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") + logger.info(f" Total optimization steps = {args.max_train_steps}") + global_step = 0 + first_epoch = 0 + + # Potentially load in the weights and states from a previous save + if args.resume_from_checkpoint: + if args.resume_from_checkpoint != "latest": + path = os.path.basename(args.resume_from_checkpoint) + else: + # Get the mos recent checkpoint + dirs = os.listdir(args.output_dir) + dirs = [d for d in dirs if d.startswith("checkpoint")] + dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) + path = dirs[-1] if len(dirs) > 0 else None + + if path is None: + accelerator.print( + f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run." + ) + args.resume_from_checkpoint = None + initial_global_step = 0 + else: + accelerator.print(f"Resuming from checkpoint {path}") + accelerator.load_state(os.path.join(args.output_dir, path)) + global_step = int(path.split("-")[1]) + + initial_global_step = global_step + first_epoch = global_step // num_update_steps_per_epoch + + else: + initial_global_step = 0 + + progress_bar = tqdm( + range(0, args.max_train_steps), + initial=initial_global_step, + desc="Steps", + # Only show the progress bar once on each machine. + disable=not accelerator.is_local_main_process, + ) + + def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): + sigmas = noise_scheduler_copy.sigmas.to(device=accelerator.device, dtype=dtype) + schedule_timesteps = noise_scheduler_copy.timesteps.to(accelerator.device) + timesteps = timesteps.to(accelerator.device) + step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < n_dim: + sigma = sigma.unsqueeze(-1) + return sigma + + for epoch in range(first_epoch, args.num_train_epochs): + transformer.train() + + for step, batch in enumerate(train_dataloader): + models_to_accumulate = [transformer] + prompts = batch["prompts"] + + with accelerator.accumulate(models_to_accumulate): + if train_dataset.custom_instance_prompts: + prompt_embeds = prompt_embeds_cache[step] + text_ids = text_ids_cache[step] + else: + num_repeat_elements = len(prompts) + prompt_embeds = prompt_embeds.repeat(num_repeat_elements, 1, 1) + text_ids = text_ids.repeat(num_repeat_elements, 1, 1) + + # Convert images to latent space + if args.cache_latents: + model_input = latents_cache[step].mode() + else: + with offload_models(vae, device=accelerator.device, offload=args.offload): + pixel_values = batch["pixel_values"].to(dtype=vae.dtype) + model_input = vae.encode(pixel_values).latent_dist.mode() + + model_input = Flux2KleinPipeline._patchify_latents(model_input) + model_input = (model_input - latents_bn_mean) / latents_bn_std + + model_input_ids = Flux2KleinPipeline._prepare_latent_ids(model_input).to(device=model_input.device) + # Sample noise that we'll add to the latents + noise = torch.randn_like(model_input) + bsz = model_input.shape[0] + + # Sample a random timestep for each image + # for weighting schemes where we sample timesteps non-uniformly + u = compute_density_for_timestep_sampling( + weighting_scheme=args.weighting_scheme, + batch_size=bsz, + logit_mean=args.logit_mean, + logit_std=args.logit_std, + mode_scale=args.mode_scale, + ) + indices = (u * noise_scheduler_copy.config.num_train_timesteps).long() + timesteps = noise_scheduler_copy.timesteps[indices].to(device=model_input.device) + + # Add noise according to flow matching. + # zt = (1 - texp) * x + texp * z1 + sigmas = get_sigmas(timesteps, n_dim=model_input.ndim, dtype=model_input.dtype) + noisy_model_input = (1.0 - sigmas) * model_input + sigmas * noise + + # [B, C, H, W] -> [B, H*W, C] + packed_noisy_model_input = Flux2KleinPipeline._pack_latents(noisy_model_input) + + # handle guidance + if unwrap_model(transformer).config.guidance_embeds: + guidance = torch.full([1], args.guidance_scale, device=accelerator.device) + guidance = guidance.expand(model_input.shape[0]) + else: + guidance = None + + # Predict the noise residual + model_pred = transformer( + hidden_states=packed_noisy_model_input, # (B, image_seq_len, C) + timestep=timesteps / 1000, + guidance=guidance, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, # B, text_seq_len, 4 + img_ids=model_input_ids, # B, image_seq_len, 4 + return_dict=False, + )[0] + model_pred = model_pred[:, : packed_noisy_model_input.size(1) :] + + model_pred = Flux2KleinPipeline._unpack_latents_with_ids(model_pred, model_input_ids) + + # these weighting schemes use a uniform timestep sampling + # and instead post-weight the loss + weighting = compute_loss_weighting_for_sd3(weighting_scheme=args.weighting_scheme, sigmas=sigmas) + + # flow matching loss + target = noise - model_input + + if args.with_prior_preservation: + # Chunk the noise and model_pred into two parts and compute the loss on each part separately. + model_pred, model_pred_prior = torch.chunk(model_pred, 2, dim=0) + target, target_prior = torch.chunk(target, 2, dim=0) + + # Compute prior loss + prior_loss = torch.mean( + (weighting.float() * (model_pred_prior.float() - target_prior.float()) ** 2).reshape( + target_prior.shape[0], -1 + ), + 1, + ) + prior_loss = prior_loss.mean() + + # Compute regular loss. + loss = torch.mean( + (weighting.float() * (model_pred.float() - target.float()) ** 2).reshape(target.shape[0], -1), + 1, + ) + loss = loss.mean() + + if args.with_prior_preservation: + # Add the prior loss to the instance loss. + loss = loss + args.prior_loss_weight * prior_loss + + accelerator.backward(loss) + if accelerator.sync_gradients: + params_to_clip = transformer.parameters() + accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm) + + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + # Checks if the accelerator has performed an optimization step behind the scenes + if accelerator.sync_gradients: + progress_bar.update(1) + global_step += 1 + + if accelerator.is_main_process or is_fsdp: + if global_step % args.checkpointing_steps == 0: + # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` + if args.checkpoints_total_limit is not None: + checkpoints = os.listdir(args.output_dir) + checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] + checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) + + # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints + if len(checkpoints) >= args.checkpoints_total_limit: + num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 + removing_checkpoints = checkpoints[0:num_to_remove] + + logger.info( + f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" + ) + logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") + + for removing_checkpoint in removing_checkpoints: + removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) + shutil.rmtree(removing_checkpoint) + + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + accelerator.save_state(save_path) + logger.info(f"Saved state to {save_path}") + + logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} + progress_bar.set_postfix(**logs) + accelerator.log(logs, step=global_step) + + if global_step >= args.max_train_steps: + break + + if accelerator.is_main_process: + if args.validation_prompt is not None and epoch % args.validation_epochs == 0: + # create pipeline + pipeline = Flux2KleinPipeline.from_pretrained( + args.pretrained_model_name_or_path, + transformer=unwrap_model(transformer), + revision=args.revision, + variant=args.variant, + torch_dtype=weight_dtype, + ) + images = log_validation( + pipeline=pipeline, + args=args, + accelerator=accelerator, + pipeline_args=validation_embeddings, + epoch=epoch, + torch_dtype=weight_dtype, + ) + + del pipeline + free_memory() + + # Save the lora layers + accelerator.wait_for_everyone() + + if is_fsdp: + transformer = unwrap_model(transformer) + state_dict = accelerator.get_state_dict(transformer) + if accelerator.is_main_process: + modules_to_save = {} + if is_fsdp: + if args.bnb_quantization_config_path is None: + if args.upcast_before_saving: + state_dict = { + k: v.to(torch.float32) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items() + } + else: + state_dict = { + k: v.to(weight_dtype) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items() + } + + transformer_lora_layers = get_peft_model_state_dict( + transformer, + state_dict=state_dict, + ) + transformer_lora_layers = { + k: v.detach().cpu().contiguous() if isinstance(v, torch.Tensor) else v + for k, v in transformer_lora_layers.items() + } + + else: + transformer = unwrap_model(transformer) + if args.bnb_quantization_config_path is None: + if args.upcast_before_saving: + transformer.to(torch.float32) + else: + transformer = transformer.to(weight_dtype) + transformer_lora_layers = get_peft_model_state_dict(transformer) + + modules_to_save["transformer"] = transformer + + Flux2KleinPipeline.save_lora_weights( + save_directory=args.output_dir, + transformer_lora_layers=transformer_lora_layers, + **_collate_lora_metadata(modules_to_save), + ) + + images = [] + run_validation = (args.validation_prompt and args.num_validation_images > 0) or (args.final_validation_prompt) + should_run_final_inference = not args.skip_final_inference and run_validation + if should_run_final_inference: + pipeline = Flux2KleinPipeline.from_pretrained( + args.pretrained_model_name_or_path, + revision=args.revision, + variant=args.variant, + torch_dtype=weight_dtype, + ) + # load attention processors + pipeline.load_lora_weights(args.output_dir) + + # run inference + images = [] + if args.validation_prompt and args.num_validation_images > 0: + images = log_validation( + pipeline=pipeline, + args=args, + accelerator=accelerator, + pipeline_args=validation_embeddings, + epoch=epoch, + is_final_validation=True, + torch_dtype=weight_dtype, + ) + images = None + del pipeline + free_memory() + + validation_prompt = args.validation_prompt if args.validation_prompt else args.final_validation_prompt + quant_training = None + if args.do_fp8_training: + quant_training = "FP8 TorchAO" + elif args.bnb_quantization_config_path: + quant_training = "BitsandBytes" + save_model_card( + (args.hub_model_id or Path(args.output_dir).name) if not args.push_to_hub else repo_id, + images=images, + base_model=args.pretrained_model_name_or_path, + instance_prompt=args.instance_prompt, + validation_prompt=validation_prompt, + repo_folder=args.output_dir, + quant_training=quant_training, + ) + + if args.push_to_hub: + upload_folder( + repo_id=repo_id, + folder_path=args.output_dir, + commit_message="End of training", + ignore_patterns=["step_*", "epoch_*"], + ) + + accelerator.end_training() + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/examples/dreambooth/train_dreambooth_lora_flux2_klein_img2img.py b/examples/dreambooth/train_dreambooth_lora_flux2_klein_img2img.py new file mode 100644 index 000000000000..0205f2e9e65f --- /dev/null +++ b/examples/dreambooth/train_dreambooth_lora_flux2_klein_img2img.py @@ -0,0 +1,1889 @@ +#!/usr/bin/env python +# coding=utf-8 +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# /// script +# dependencies = [ +# "diffusers @ git+https://github.com/huggingface/diffusers.git", +# "torch>=2.0.0", +# "accelerate>=0.31.0", +# "transformers>=4.41.2", +# "ftfy", +# "tensorboard", +# "Jinja2", +# "peft>=0.11.1", +# "sentencepiece", +# "torchvision", +# "datasets", +# "bitsandbytes", +# "prodigyopt", +# ] +# /// + +import argparse +import copy +import itertools +import json +import logging +import math +import os +import random +import shutil +from contextlib import nullcontext +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import transformers +from accelerate import Accelerator +from accelerate.logging import get_logger +from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration, set_seed +from huggingface_hub import create_repo, upload_folder +from peft import LoraConfig, prepare_model_for_kbit_training, set_peft_model_state_dict +from peft.utils import get_peft_model_state_dict +from PIL import Image +from PIL.ImageOps import exif_transpose +from torch.utils.data import Dataset +from torch.utils.data.sampler import BatchSampler +from torchvision import transforms +from torchvision.transforms import functional as TF +from tqdm.auto import tqdm +from transformers import Qwen2TokenizerFast, Qwen3ForCausalLM + +import diffusers +from diffusers import ( + AutoencoderKLFlux2, + BitsAndBytesConfig, + FlowMatchEulerDiscreteScheduler, + Flux2KleinPipeline, + Flux2Transformer2DModel, +) +from diffusers.optimization import get_scheduler +from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor +from diffusers.training_utils import ( + _collate_lora_metadata, + _to_cpu_contiguous, + cast_training_params, + compute_density_for_timestep_sampling, + compute_loss_weighting_for_sd3, + find_nearest_bucket, + free_memory, + get_fsdp_kwargs_from_accelerator, + offload_models, + parse_buckets_string, + wrap_with_fsdp, +) +from diffusers.utils import ( + check_min_version, + convert_unet_state_dict_to_peft, + is_wandb_available, + load_image, +) +from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card +from diffusers.utils.import_utils import is_torch_npu_available +from diffusers.utils.torch_utils import is_compiled_module + + +if getattr(torch, "distributed", None) is not None: + import torch.distributed as dist + +if is_wandb_available(): + import wandb + +# Will error if the minimal version of diffusers is not installed. Remove at your own risks. +check_min_version("0.38.0.dev0") + +logger = get_logger(__name__) + + +def save_model_card( + repo_id: str, + images=None, + base_model: str = None, + instance_prompt=None, + validation_prompt=None, + repo_folder=None, + fp8_training=False, +): + widget_dict = [] + if images is not None: + for i, image in enumerate(images): + image.save(os.path.join(repo_folder, f"image_{i}.png")) + widget_dict.append( + {"text": validation_prompt if validation_prompt else " ", "output": {"url": f"image_{i}.png"}} + ) + + model_description = f""" +# Flux.2 [Klein] DreamBooth LoRA - {repo_id} + + + +## Model description + +These are {repo_id} DreamBooth LoRA weights for {base_model}. + +The weights were trained using [DreamBooth](https://dreambooth.github.io/) with the [Flux2 diffusers trainer](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/README_flux2.md). + +FP8 training? {fp8_training} + +## Trigger words + +You should use `{instance_prompt}` to trigger the image generation. + +## Download model + +[Download the *.safetensors LoRA]({repo_id}/tree/main) in the Files & versions tab. + +## Use it with the [🧨 diffusers library](https://github.com/huggingface/diffusers) + +```py +from diffusers import AutoPipelineForText2Image +import torch +pipeline = AutoPipelineForText2Image.from_pretrained("black-forest-labs/FLUX.2", torch_dtype=torch.bfloat16).to('cuda') +pipeline.load_lora_weights('{repo_id}', weight_name='pytorch_lora_weights.safetensors') +image = pipeline('{validation_prompt if validation_prompt else instance_prompt}').images[0] +``` + +For more details, including weighting, merging and fusing LoRAs, check the [documentation on loading LoRAs in diffusers](https://huggingface.co/docs/diffusers/main/en/using-diffusers/loading_adapters) + +## License + +Please adhere to the licensing terms as described [here](https://huggingface.co/black-forest-labs/FLUX.2/blob/main/LICENSE.md). +""" + model_card = load_or_create_model_card( + repo_id_or_path=repo_id, + from_training=True, + license="other", + base_model=base_model, + prompt=instance_prompt, + model_description=model_description, + widget=widget_dict, + ) + tags = [ + "text-to-image", + "diffusers-training", + "diffusers", + "lora", + "flux2", + "flux2-diffusers", + "template:sd-lora", + ] + + model_card = populate_model_card(model_card, tags=tags) + model_card.save(os.path.join(repo_folder, "README.md")) + + +def log_validation( + pipeline, + args, + accelerator, + pipeline_args, + epoch, + torch_dtype, + is_final_validation=False, +): + args.num_validation_images = args.num_validation_images if args.num_validation_images else 1 + logger.info( + f"Running validation... \n Generating {args.num_validation_images} images with prompt:" + f" {args.validation_prompt}." + ) + pipeline = pipeline.to(dtype=torch_dtype) + pipeline.enable_model_cpu_offload() + pipeline.set_progress_bar_config(disable=True) + + # run inference + generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed is not None else None + autocast_ctx = torch.autocast(accelerator.device.type) if not is_final_validation else nullcontext() + + images = [] + for _ in range(args.num_validation_images): + with autocast_ctx: + image = pipeline( + image=pipeline_args["image"], + prompt_embeds=pipeline_args["prompt_embeds"], + negative_prompt_embeds=pipeline_args["negative_prompt_embeds"], + generator=generator, + ).images[0] + images.append(image) + + for tracker in accelerator.trackers: + phase_name = "test" if is_final_validation else "validation" + if tracker.name == "tensorboard": + np_images = np.stack([np.asarray(img) for img in images]) + tracker.writer.add_images(phase_name, np_images, epoch, dataformats="NHWC") + if tracker.name == "wandb": + tracker.log( + { + phase_name: [ + wandb.Image(image, caption=f"{i}: {args.validation_prompt}") for i, image in enumerate(images) + ] + } + ) + + del pipeline + free_memory() + + return images + + +def module_filter_fn(mod: torch.nn.Module, fqn: str): + # don't convert the output module + if fqn == "proj_out": + return False + # don't convert linear modules with weight dimensions not divisible by 16 + if isinstance(mod, torch.nn.Linear): + if mod.in_features % 16 != 0 or mod.out_features % 16 != 0: + return False + return True + + +def parse_args(input_args=None): + parser = argparse.ArgumentParser(description="Simple example of a training script.") + parser.add_argument( + "--pretrained_model_name_or_path", + type=str, + default=None, + required=True, + help="Path to pretrained model or model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--revision", + type=str, + default=None, + required=False, + help="Revision of pretrained model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--bnb_quantization_config_path", + type=str, + default=None, + help="Quantization config in a JSON file that will be used to define the bitsandbytes quant config of the DiT.", + ) + parser.add_argument( + "--do_fp8_training", + action="store_true", + help="if we are doing FP8 training.", + ) + parser.add_argument( + "--variant", + type=str, + default=None, + help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16", + ) + parser.add_argument( + "--dataset_name", + type=str, + default=None, + help=( + "The name of the Dataset (from the HuggingFace hub) containing the training data of instance images (could be your own, possibly private," + " dataset). It can also be a path pointing to a local copy of a dataset in your filesystem," + " or to a folder containing files that 🤗 Datasets can understand." + ), + ) + parser.add_argument( + "--dataset_config_name", + type=str, + default=None, + help="The config of the Dataset, leave as None if there's only one config.", + ) + parser.add_argument( + "--instance_data_dir", + type=str, + default=None, + help=("A folder containing the training data. "), + ) + + parser.add_argument( + "--cache_dir", + type=str, + default=None, + help="The directory where the downloaded models and datasets will be stored.", + ) + + parser.add_argument( + "--image_column", + type=str, + default="image", + help="The column of the dataset containing the target image. By " + "default, the standard Image Dataset maps out 'file_name' " + "to 'image'.", + ) + parser.add_argument( + "--cond_image_column", + type=str, + default=None, + help="Column in the dataset containing the condition image. Must be specified when performing I2I fine-tuning", + ) + parser.add_argument( + "--caption_column", + type=str, + default=None, + help="The column of the dataset containing the instance prompt for each image", + ) + + parser.add_argument("--repeats", type=int, default=1, help="How many times to repeat the training data.") + + parser.add_argument( + "--class_data_dir", + type=str, + default=None, + required=False, + help="A folder containing the training data of class images.", + ) + parser.add_argument( + "--instance_prompt", + type=str, + default=None, + required=False, + help="The prompt with identifier specifying the instance, e.g. 'photo of a TOK dog', 'in the style of TOK'", + ) + parser.add_argument( + "--max_sequence_length", + type=int, + default=512, + help="Maximum sequence length to use with with the T5 text encoder", + ) + parser.add_argument( + "--validation_prompt", + type=str, + default=None, + help="A prompt that is used during validation to verify that the model is learning.", + ) + parser.add_argument( + "--validation_image", + type=str, + default=None, + help="path to an image that is used during validation as the condition image to verify that the model is learning.", + ) + parser.add_argument( + "--skip_final_inference", + default=False, + action="store_true", + help="Whether to skip the final inference step with loaded lora weights upon training completion. This will run intermediate validation inference if `validation_prompt` is provided. Specify to reduce memory.", + ) + parser.add_argument( + "--final_validation_prompt", + type=str, + default=None, + help="A prompt that is used during a final validation to verify that the model is learning. Ignored if `--validation_prompt` is provided.", + ) + parser.add_argument( + "--num_validation_images", + type=int, + default=4, + help="Number of images that should be generated during validation with `validation_prompt`.", + ) + parser.add_argument( + "--validation_epochs", + type=int, + default=50, + help=( + "Run dreambooth validation every X epochs. Dreambooth validation consists of running the prompt" + " `args.validation_prompt` multiple times: `args.num_validation_images`." + ), + ) + parser.add_argument( + "--rank", + type=int, + default=4, + help=("The dimension of the LoRA update matrices."), + ) + parser.add_argument( + "--lora_alpha", + type=int, + default=4, + help="LoRA alpha to be used for additional scaling.", + ) + parser.add_argument("--lora_dropout", type=float, default=0.0, help="Dropout probability for LoRA layers") + + parser.add_argument( + "--output_dir", + type=str, + default="flux-dreambooth-lora", + help="The output directory where the model predictions and checkpoints will be written.", + ) + parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") + parser.add_argument( + "--resolution", + type=int, + default=512, + help=( + "The resolution for input images, all the images in the train/validation dataset will be resized to this" + " resolution" + ), + ) + parser.add_argument( + "--aspect_ratio_buckets", + type=str, + default=None, + help=( + "Aspect ratio buckets to use for training. Define as a string of 'h1,w1;h2,w2;...'. " + "e.g. '1024,1024;768,1360;1360,768;880,1168;1168,880;1248,832;832,1248'" + "Images will be resized and cropped to fit the nearest bucket. If provided, --resolution is ignored." + ), + ) + parser.add_argument( + "--center_crop", + default=False, + action="store_true", + help=( + "Whether to center crop the input images to the resolution. If not set, the images will be randomly" + " cropped. The images will be resized to the resolution first before cropping." + ), + ) + parser.add_argument( + "--random_flip", + action="store_true", + help="whether to randomly flip images horizontally", + ) + parser.add_argument( + "--train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader." + ) + parser.add_argument( + "--sample_batch_size", type=int, default=4, help="Batch size (per device) for sampling images." + ) + parser.add_argument("--num_train_epochs", type=int, default=1) + parser.add_argument( + "--max_train_steps", + type=int, + default=None, + help="Total number of training steps to perform. If provided, overrides num_train_epochs.", + ) + parser.add_argument( + "--checkpointing_steps", + type=int, + default=500, + help=( + "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final" + " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming" + " training using `--resume_from_checkpoint`." + ), + ) + parser.add_argument( + "--checkpoints_total_limit", + type=int, + default=None, + help=("Max number of checkpoints to store."), + ) + parser.add_argument( + "--resume_from_checkpoint", + type=str, + default=None, + help=( + "Whether training should be resumed from a previous checkpoint. Use a path saved by" + ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' + ), + ) + parser.add_argument( + "--gradient_accumulation_steps", + type=int, + default=1, + help="Number of updates steps to accumulate before performing a backward/update pass.", + ) + parser.add_argument( + "--gradient_checkpointing", + action="store_true", + help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.", + ) + parser.add_argument( + "--learning_rate", + type=float, + default=1e-4, + help="Initial learning rate (after the potential warmup period) to use.", + ) + + parser.add_argument( + "--guidance_scale", + type=float, + default=3.5, + help="the FLUX.1 dev variant is a guidance distilled model", + ) + + parser.add_argument( + "--scale_lr", + action="store_true", + default=False, + help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.", + ) + parser.add_argument( + "--lr_scheduler", + type=str, + default="constant", + help=( + 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' + ' "constant", "constant_with_warmup"]' + ), + ) + parser.add_argument( + "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler." + ) + parser.add_argument( + "--lr_num_cycles", + type=int, + default=1, + help="Number of hard resets of the lr in cosine_with_restarts scheduler.", + ) + parser.add_argument("--lr_power", type=float, default=1.0, help="Power factor of the polynomial scheduler.") + parser.add_argument( + "--dataloader_num_workers", + type=int, + default=0, + help=( + "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process." + ), + ) + parser.add_argument( + "--weighting_scheme", + type=str, + default="none", + choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "none"], + help=('We default to the "none" weighting scheme for uniform sampling and uniform loss'), + ) + parser.add_argument( + "--logit_mean", type=float, default=0.0, help="mean to use when using the `'logit_normal'` weighting scheme." + ) + parser.add_argument( + "--logit_std", type=float, default=1.0, help="std to use when using the `'logit_normal'` weighting scheme." + ) + parser.add_argument( + "--mode_scale", + type=float, + default=1.29, + help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.", + ) + parser.add_argument( + "--optimizer", + type=str, + default="AdamW", + help=('The optimizer type to use. Choose between ["AdamW", "prodigy"]'), + ) + + parser.add_argument( + "--use_8bit_adam", + action="store_true", + help="Whether or not to use 8-bit Adam from bitsandbytes. Ignored if optimizer is not set to AdamW", + ) + + parser.add_argument( + "--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam and Prodigy optimizers." + ) + parser.add_argument( + "--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam and Prodigy optimizers." + ) + parser.add_argument( + "--prodigy_beta3", + type=float, + default=None, + help="coefficients for computing the Prodigy stepsize using running averages. If set to None, " + "uses the value of square root of beta2. Ignored if optimizer is adamW", + ) + parser.add_argument("--prodigy_decouple", type=bool, default=True, help="Use AdamW style decoupled weight decay") + parser.add_argument("--adam_weight_decay", type=float, default=1e-04, help="Weight decay to use for unet params") + parser.add_argument( + "--adam_weight_decay_text_encoder", type=float, default=1e-03, help="Weight decay to use for text_encoder" + ) + + parser.add_argument( + "--lora_layers", + type=str, + default=None, + help=( + 'The transformer modules to apply LoRA training on. Please specify the layers in a comma separated. E.g. - "to_k,to_q,to_v,to_out.0" will result in lora training of attention layers only' + ), + ) + + parser.add_argument( + "--adam_epsilon", + type=float, + default=1e-08, + help="Epsilon value for the Adam optimizer and Prodigy optimizers.", + ) + + parser.add_argument( + "--prodigy_use_bias_correction", + type=bool, + default=True, + help="Turn on Adam's bias correction. True by default. Ignored if optimizer is adamW", + ) + parser.add_argument( + "--prodigy_safeguard_warmup", + type=bool, + default=True, + help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage. True by default. " + "Ignored if optimizer is adamW", + ) + parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") + parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.") + parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.") + parser.add_argument( + "--hub_model_id", + type=str, + default=None, + help="The name of the repository to keep in sync with the local `output_dir`.", + ) + parser.add_argument( + "--logging_dir", + type=str, + default="logs", + help=( + "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to" + " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***." + ), + ) + parser.add_argument( + "--allow_tf32", + action="store_true", + help=( + "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" + " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" + ), + ) + parser.add_argument( + "--cache_latents", + action="store_true", + default=False, + help="Cache the VAE latents", + ) + parser.add_argument( + "--report_to", + type=str, + default="tensorboard", + help=( + 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' + ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' + ), + ) + parser.add_argument( + "--mixed_precision", + type=str, + default=None, + choices=["no", "fp16", "bf16"], + help=( + "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" + " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" + " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." + ), + ) + parser.add_argument( + "--upcast_before_saving", + action="store_true", + default=False, + help=( + "Whether to upcast the trained transformer layers to float32 before saving (at the end of training). " + "Defaults to precision dtype used for training to save memory" + ), + ) + parser.add_argument( + "--offload", + action="store_true", + help="Whether to offload the VAE and the text encoder to CPU when they are not used.", + ) + + parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") + parser.add_argument("--enable_npu_flash_attention", action="store_true", help="Enabla Flash Attention for NPU") + parser.add_argument("--fsdp_text_encoder", action="store_true", help="Use FSDP for text encoder") + + if input_args is not None: + args = parser.parse_args(input_args) + else: + args = parser.parse_args() + + if args.cond_image_column is None: + raise ValueError( + "you must provide --cond_image_column for image-to-image training. Otherwise please see Flux2 text-to-image training example." + ) + else: + assert args.image_column is not None + assert args.caption_column is not None + + if args.dataset_name is None and args.instance_data_dir is None: + raise ValueError("Specify either `--dataset_name` or `--instance_data_dir`") + + if args.dataset_name is not None and args.instance_data_dir is not None: + raise ValueError("Specify only one of `--dataset_name` or `--instance_data_dir`") + + env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) + if env_local_rank != -1 and env_local_rank != args.local_rank: + args.local_rank = env_local_rank + + return args + + +class DreamBoothDataset(Dataset): + """ + A dataset to prepare the instance and class images with the prompts for fine-tuning the model. + It pre-processes the images. + """ + + def __init__( + self, + instance_data_root, + instance_prompt, + size=1024, + repeats=1, + center_crop=False, + buckets=None, + ): + self.size = size + self.center_crop = center_crop + + self.instance_prompt = instance_prompt + self.custom_instance_prompts = None + + self.buckets = buckets + + # if --dataset_name is provided or a metadata jsonl file is provided in the local --instance_data directory, + # we load the training data using load_dataset + if args.dataset_name is not None: + try: + from datasets import load_dataset + except ImportError: + raise ImportError( + "You are trying to load your data using the datasets library. If you wish to train using custom " + "captions please install the datasets library: `pip install datasets`. If you wish to load a " + "local folder containing images only, specify --instance_data_dir instead." + ) + # Downloading and loading a dataset from the hub. + # See more about loading custom images at + # https://huggingface.co/docs/datasets/v2.0.0/en/dataset_script + dataset = load_dataset( + args.dataset_name, + args.dataset_config_name, + cache_dir=args.cache_dir, + ) + # Preprocessing the datasets. + column_names = dataset["train"].column_names + + # 6. Get the column names for input/target. + if args.cond_image_column is not None and args.cond_image_column not in column_names: + raise ValueError( + f"`--cond_image_column` value '{args.cond_image_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}" + ) + if args.image_column is None: + image_column = column_names[0] + logger.info(f"image column defaulting to {image_column}") + else: + image_column = args.image_column + if image_column not in column_names: + raise ValueError( + f"`--image_column` value '{args.image_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}" + ) + instance_images = dataset["train"][image_column] + cond_images = None + cond_image_column = args.cond_image_column + if cond_image_column is not None: + cond_images = [dataset["train"][i][cond_image_column] for i in range(len(dataset["train"]))] + assert len(instance_images) == len(cond_images) + + if args.caption_column is None: + logger.info( + "No caption column provided, defaulting to instance_prompt for all images. If your dataset " + "contains captions/prompts for the images, make sure to specify the " + "column as --caption_column" + ) + self.custom_instance_prompts = None + else: + if args.caption_column not in column_names: + raise ValueError( + f"`--caption_column` value '{args.caption_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}" + ) + custom_instance_prompts = dataset["train"][args.caption_column] + # create final list of captions according to --repeats + self.custom_instance_prompts = [] + for caption in custom_instance_prompts: + self.custom_instance_prompts.extend(itertools.repeat(caption, repeats)) + else: + self.instance_data_root = Path(instance_data_root) + if not self.instance_data_root.exists(): + raise ValueError("Instance images root doesn't exists.") + + instance_images = [Image.open(path) for path in list(Path(instance_data_root).iterdir())] + self.custom_instance_prompts = None + + self.instance_images = [] + self.cond_images = [] + for i, img in enumerate(instance_images): + self.instance_images.extend(itertools.repeat(img, repeats)) + if args.dataset_name is not None and cond_images is not None: + self.cond_images.extend(itertools.repeat(cond_images[i], repeats)) + + self.pixel_values = [] + self.cond_pixel_values = [] + for i, image in enumerate(self.instance_images): + image = exif_transpose(image) + if not image.mode == "RGB": + image = image.convert("RGB") + dest_image = None + if self.cond_images: # todo: take care of max area for buckets + dest_image = self.cond_images[i] + image_width, image_height = dest_image.size + if image_width * image_height > 1024 * 1024: + dest_image = Flux2ImageProcessor._resize_to_target_area(dest_image, 1024 * 1024) + image_width, image_height = dest_image.size + + multiple_of = 2 ** (4 - 1) # 2 ** (len(vae.config.block_out_channels) - 1), temp! + image_width = (image_width // multiple_of) * multiple_of + image_height = (image_height // multiple_of) * multiple_of + image_processor = Flux2ImageProcessor() + dest_image = image_processor.preprocess( + dest_image, height=image_height, width=image_width, resize_mode="crop" + ) + # Convert back to PIL + dest_image = dest_image.squeeze(0) + if dest_image.min() < 0: + dest_image = (dest_image + 1) / 2 + dest_image = (torch.clamp(dest_image, 0, 1) * 255).byte().cpu() + + if dest_image.shape[0] == 1: + # Gray scale image + dest_image = Image.fromarray(dest_image.squeeze().numpy(), mode="L") + else: + # RGB scale image: (C, H, W) -> (H, W, C) + dest_image = TF.to_pil_image(dest_image) + + dest_image = exif_transpose(dest_image) + if not dest_image.mode == "RGB": + dest_image = dest_image.convert("RGB") + + width, height = image.size + + # Find the closest bucket + bucket_idx = find_nearest_bucket(height, width, self.buckets) + target_height, target_width = self.buckets[bucket_idx] + self.size = (target_height, target_width) + + # based on the bucket assignment, define the transformations + image, dest_image = self.paired_transform( + image, + dest_image=dest_image, + size=self.size, + center_crop=args.center_crop, + random_flip=args.random_flip, + ) + self.pixel_values.append((image, bucket_idx)) + if dest_image is not None: + self.cond_pixel_values.append((dest_image, bucket_idx)) + + self.num_instance_images = len(self.instance_images) + self._length = self.num_instance_images + + self.image_transforms = transforms.Compose( + [ + transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR), + transforms.CenterCrop(size) if center_crop else transforms.RandomCrop(size), + transforms.ToTensor(), + transforms.Normalize([0.5], [0.5]), + ] + ) + + def __len__(self): + return self._length + + def __getitem__(self, index): + example = {} + instance_image, bucket_idx = self.pixel_values[index % self.num_instance_images] + example["instance_images"] = instance_image + example["bucket_idx"] = bucket_idx + if self.cond_pixel_values: + dest_image, _ = self.cond_pixel_values[index % self.num_instance_images] + example["cond_images"] = dest_image + + if self.custom_instance_prompts: + caption = self.custom_instance_prompts[index % self.num_instance_images] + if caption: + example["instance_prompt"] = caption + else: + example["instance_prompt"] = self.instance_prompt + + else: # custom prompts were provided, but length does not match size of image dataset + example["instance_prompt"] = self.instance_prompt + + return example + + def paired_transform(self, image, dest_image=None, size=(224, 224), center_crop=False, random_flip=False): + # 1. Resize (deterministic) + resize = transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR) + image = resize(image) + if dest_image is not None: + dest_image = resize(dest_image) + + # 2. Crop: either center or SAME random crop + if center_crop: + crop = transforms.CenterCrop(size) + image = crop(image) + if dest_image is not None: + dest_image = crop(dest_image) + else: + # get_params returns (i, j, h, w) + i, j, h, w = transforms.RandomCrop.get_params(image, output_size=size) + image = TF.crop(image, i, j, h, w) + if dest_image is not None: + dest_image = TF.crop(dest_image, i, j, h, w) + + # 3. Random horizontal flip with the SAME coin flip + if random_flip: + do_flip = random.random() < 0.5 + if do_flip: + image = TF.hflip(image) + if dest_image is not None: + dest_image = TF.hflip(dest_image) + + # 4. ToTensor + Normalize (deterministic) + to_tensor = transforms.ToTensor() + normalize = transforms.Normalize([0.5], [0.5]) + image = normalize(to_tensor(image)) + if dest_image is not None: + dest_image = normalize(to_tensor(dest_image)) + + return (image, dest_image) if dest_image is not None else (image, None) + + +def collate_fn(examples): + pixel_values = [example["instance_images"] for example in examples] + prompts = [example["instance_prompt"] for example in examples] + + pixel_values = torch.stack(pixel_values) + pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float() + + batch = {"pixel_values": pixel_values, "prompts": prompts} + if any("cond_images" in example for example in examples): + cond_pixel_values = [example["cond_images"] for example in examples] + cond_pixel_values = torch.stack(cond_pixel_values) + cond_pixel_values = cond_pixel_values.to(memory_format=torch.contiguous_format).float() + batch.update({"cond_pixel_values": cond_pixel_values}) + return batch + + +class BucketBatchSampler(BatchSampler): + def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False): + if not isinstance(batch_size, int) or batch_size <= 0: + raise ValueError("batch_size should be a positive integer value, but got batch_size={}".format(batch_size)) + if not isinstance(drop_last, bool): + raise ValueError("drop_last should be a boolean value, but got drop_last={}".format(drop_last)) + + self.dataset = dataset + self.batch_size = batch_size + self.drop_last = drop_last + + # Group indices by bucket + self.bucket_indices = [[] for _ in range(len(self.dataset.buckets))] + for idx, (_, bucket_idx) in enumerate(self.dataset.pixel_values): + self.bucket_indices[bucket_idx].append(idx) + + self.sampler_len = 0 + self.batches = [] + + # Pre-generate batches for each bucket + for indices_in_bucket in self.bucket_indices: + # Shuffle indices within the bucket + random.shuffle(indices_in_bucket) + # Create batches + for i in range(0, len(indices_in_bucket), self.batch_size): + batch = indices_in_bucket[i : i + self.batch_size] + if len(batch) < self.batch_size and self.drop_last: + continue # Skip partial batch if drop_last is True + self.batches.append(batch) + self.sampler_len += 1 # Count the number of batches + + def __iter__(self): + # Shuffle the order of the batches each epoch + random.shuffle(self.batches) + for batch in self.batches: + yield batch + + def __len__(self): + return self.sampler_len + + +class PromptDataset(Dataset): + "A simple dataset to prepare the prompts to generate class images on multiple GPUs." + + def __init__(self, prompt, num_samples): + self.prompt = prompt + self.num_samples = num_samples + + def __len__(self): + return self.num_samples + + def __getitem__(self, index): + example = {} + example["prompt"] = self.prompt + example["index"] = index + return example + + +def main(args): + if args.report_to == "wandb" and args.hub_token is not None: + raise ValueError( + "You cannot use both --report_to=wandb and --hub_token due to a security risk of exposing your token." + " Please use `hf auth login` to authenticate with the Hub." + ) + + if torch.backends.mps.is_available() and args.mixed_precision == "bf16": + # due to pytorch#99272, MPS does not yet support bfloat16. + raise ValueError( + "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead." + ) + if args.do_fp8_training: + from torchao.float8 import Float8LinearConfig, convert_to_float8_training + + logging_dir = Path(args.output_dir, args.logging_dir) + + accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) + kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) + accelerator = Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + mixed_precision=args.mixed_precision, + log_with=args.report_to, + project_config=accelerator_project_config, + kwargs_handlers=[kwargs], + ) + + # Disable AMP for MPS. + if torch.backends.mps.is_available(): + accelerator.native_amp = False + + if args.report_to == "wandb": + if not is_wandb_available(): + raise ImportError("Make sure to install wandb if you want to use it for logging during training.") + + # Make one log on every process with the configuration for debugging. + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + ) + logger.info(accelerator.state, main_process_only=False) + if accelerator.is_local_main_process: + transformers.utils.logging.set_verbosity_warning() + diffusers.utils.logging.set_verbosity_info() + else: + transformers.utils.logging.set_verbosity_error() + diffusers.utils.logging.set_verbosity_error() + + # If passed along, set the training seed now. + if args.seed is not None: + set_seed(args.seed) + + # Handle the repository creation + if accelerator.is_main_process: + if args.output_dir is not None: + os.makedirs(args.output_dir, exist_ok=True) + + if args.push_to_hub: + repo_id = create_repo( + repo_id=args.hub_model_id or Path(args.output_dir).name, + exist_ok=True, + ).repo_id + + # Load the tokenizers + tokenizer = Qwen2TokenizerFast.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="tokenizer", + revision=args.revision, + ) + + # For mixed precision training we cast all non-trainable weights (vae, text_encoder and transformer) to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + weight_dtype = torch.float32 + if accelerator.mixed_precision == "fp16": + weight_dtype = torch.float16 + elif accelerator.mixed_precision == "bf16": + weight_dtype = torch.bfloat16 + + # Load scheduler and models + noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="scheduler", + revision=args.revision, + ) + noise_scheduler_copy = copy.deepcopy(noise_scheduler) + vae = AutoencoderKLFlux2.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="vae", + revision=args.revision, + variant=args.variant, + ) + latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(accelerator.device) + latents_bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to( + accelerator.device + ) + + quantization_config = None + if args.bnb_quantization_config_path is not None: + with open(args.bnb_quantization_config_path, "r") as f: + config_kwargs = json.load(f) + if "load_in_4bit" in config_kwargs and config_kwargs["load_in_4bit"]: + config_kwargs["bnb_4bit_compute_dtype"] = weight_dtype + quantization_config = BitsAndBytesConfig(**config_kwargs) + + transformer = Flux2Transformer2DModel.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="transformer", + revision=args.revision, + variant=args.variant, + quantization_config=quantization_config, + torch_dtype=weight_dtype, + ) + if args.bnb_quantization_config_path is not None: + transformer = prepare_model_for_kbit_training(transformer, use_gradient_checkpointing=False) + + text_encoder = Qwen3ForCausalLM.from_pretrained( + args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision, variant=args.variant + ) + text_encoder.requires_grad_(False) + + # We only train the additional adapter LoRA layers + transformer.requires_grad_(False) + vae.requires_grad_(False) + + if args.enable_npu_flash_attention: + if is_torch_npu_available(): + logger.info("npu flash attention enabled.") + transformer.set_attention_backend("_native_npu") + else: + raise ValueError("npu flash attention requires torch_npu extensions and is supported only on npu device ") + + if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16: + # due to pytorch#99272, MPS does not yet support bfloat16. + raise ValueError( + "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead." + ) + + to_kwargs = {"dtype": weight_dtype, "device": accelerator.device} if not args.offload else {"dtype": weight_dtype} + # flux vae is stable in bf16 so load it in weight_dtype to reduce memory + vae.to(**to_kwargs) + # we never offload the transformer to CPU, so we can just use the accelerator device + transformer_to_kwargs = ( + {"device": accelerator.device} + if args.bnb_quantization_config_path is not None + else {"device": accelerator.device, "dtype": weight_dtype} + ) + + is_fsdp = getattr(accelerator.state, "fsdp_plugin", None) is not None + if not is_fsdp: + transformer.to(**transformer_to_kwargs) + + if args.do_fp8_training: + convert_to_float8_training( + transformer, module_filter_fn=module_filter_fn, config=Float8LinearConfig(pad_inner_dim=True) + ) + + text_encoder.to(**to_kwargs) + # Initialize a text encoding pipeline and keep it to CPU for now. + text_encoding_pipeline = Flux2KleinPipeline.from_pretrained( + args.pretrained_model_name_or_path, + vae=None, + transformer=None, + tokenizer=tokenizer, + text_encoder=text_encoder, + scheduler=None, + revision=args.revision, + ) + + if args.gradient_checkpointing: + transformer.enable_gradient_checkpointing() + + if args.lora_layers is not None: + target_modules = [layer.strip() for layer in args.lora_layers.split(",")] + else: + target_modules = ["to_k", "to_q", "to_v", "to_out.0"] + + # now we will add new LoRA weights the transformer layers + transformer_lora_config = LoraConfig( + r=args.rank, + lora_alpha=args.lora_alpha, + lora_dropout=args.lora_dropout, + init_lora_weights="gaussian", + target_modules=target_modules, + ) + transformer.add_adapter(transformer_lora_config) + + def unwrap_model(model): + model = accelerator.unwrap_model(model) + model = model._orig_mod if is_compiled_module(model) else model + return model + + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + def save_model_hook(models, weights, output_dir): + transformer_cls = type(unwrap_model(transformer)) + + # 1) Validate and pick the transformer model + modules_to_save: dict[str, Any] = {} + transformer_model = None + + for model in models: + if isinstance(unwrap_model(model), transformer_cls): + transformer_model = model + modules_to_save["transformer"] = model + else: + raise ValueError(f"unexpected save model: {model.__class__}") + + if transformer_model is None: + raise ValueError("No transformer model found in 'models'") + + # 2) Optionally gather FSDP state dict once + state_dict = accelerator.get_state_dict(model) if is_fsdp else None + + # 3) Only main process materializes the LoRA state dict + transformer_lora_layers_to_save = None + if accelerator.is_main_process: + peft_kwargs = {} + if is_fsdp: + peft_kwargs["state_dict"] = state_dict + + transformer_lora_layers_to_save = get_peft_model_state_dict( + unwrap_model(transformer_model) if is_fsdp else transformer_model, + **peft_kwargs, + ) + + if is_fsdp: + transformer_lora_layers_to_save = _to_cpu_contiguous(transformer_lora_layers_to_save) + + # make sure to pop weight so that corresponding model is not saved again + if weights: + weights.pop() + + Flux2KleinPipeline.save_lora_weights( + output_dir, + transformer_lora_layers=transformer_lora_layers_to_save, + **_collate_lora_metadata(modules_to_save), + ) + + def load_model_hook(models, input_dir): + transformer_ = None + + if not is_fsdp: + while len(models) > 0: + model = models.pop() + + if isinstance(unwrap_model(model), type(unwrap_model(transformer))): + transformer_ = unwrap_model(model) + else: + raise ValueError(f"unexpected save model: {model.__class__}") + else: + transformer_ = Flux2Transformer2DModel.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="transformer", + ) + transformer_.add_adapter(transformer_lora_config) + + lora_state_dict = Flux2KleinPipeline.lora_state_dict(input_dir) + + transformer_state_dict = { + f"{k.replace('transformer.', '')}": v for k, v in lora_state_dict.items() if k.startswith("transformer.") + } + transformer_state_dict = convert_unet_state_dict_to_peft(transformer_state_dict) + incompatible_keys = set_peft_model_state_dict(transformer_, transformer_state_dict, adapter_name="default") + if incompatible_keys is not None: + # check only for unexpected keys + unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None) + if unexpected_keys: + logger.warning( + f"Loading adapter weights from state_dict led to unexpected keys not found in the model: " + f" {unexpected_keys}. " + ) + + # Make sure the trainable params are in float32. This is again needed since the base models + # are in `weight_dtype`. More details: + # https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804 + if args.mixed_precision == "fp16": + models = [transformer_] + # only upcast trainable parameters (LoRA) into fp32 + cast_training_params(models) + + accelerator.register_save_state_pre_hook(save_model_hook) + accelerator.register_load_state_pre_hook(load_model_hook) + + # Enable TF32 for faster training on Ampere GPUs, + # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices + if args.allow_tf32 and torch.cuda.is_available(): + torch.backends.cuda.matmul.allow_tf32 = True + + if args.scale_lr: + args.learning_rate = ( + args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes + ) + + # Make sure the trainable params are in float32. + if args.mixed_precision == "fp16": + models = [transformer] + # only upcast trainable parameters (LoRA) into fp32 + cast_training_params(models, dtype=torch.float32) + + transformer_lora_parameters = list(filter(lambda p: p.requires_grad, transformer.parameters())) + + # Optimization parameters + transformer_parameters_with_lr = {"params": transformer_lora_parameters, "lr": args.learning_rate} + params_to_optimize = [transformer_parameters_with_lr] + + # Optimizer creation + if not (args.optimizer.lower() == "prodigy" or args.optimizer.lower() == "adamw"): + logger.warning( + f"Unsupported choice of optimizer: {args.optimizer}.Supported optimizers include [adamW, prodigy]." + "Defaulting to adamW" + ) + args.optimizer = "adamw" + + if args.use_8bit_adam and not args.optimizer.lower() == "adamw": + logger.warning( + f"use_8bit_adam is ignored when optimizer is not set to 'AdamW'. Optimizer was " + f"set to {args.optimizer.lower()}" + ) + + if args.optimizer.lower() == "adamw": + if args.use_8bit_adam: + try: + import bitsandbytes as bnb + except ImportError: + raise ImportError( + "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`." + ) + + optimizer_class = bnb.optim.AdamW8bit + else: + optimizer_class = torch.optim.AdamW + + optimizer = optimizer_class( + params_to_optimize, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + + if args.optimizer.lower() == "prodigy": + try: + import prodigyopt + except ImportError: + raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`") + + optimizer_class = prodigyopt.Prodigy + + if args.learning_rate <= 0.1: + logger.warning( + "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0" + ) + + optimizer = optimizer_class( + params_to_optimize, + betas=(args.adam_beta1, args.adam_beta2), + beta3=args.prodigy_beta3, + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + decouple=args.prodigy_decouple, + use_bias_correction=args.prodigy_use_bias_correction, + safeguard_warmup=args.prodigy_safeguard_warmup, + ) + + if args.aspect_ratio_buckets is not None: + buckets = parse_buckets_string(args.aspect_ratio_buckets) + else: + buckets = [(args.resolution, args.resolution)] + logger.info(f"Using parsed aspect ratio buckets: {buckets}") + + # Dataset and DataLoaders creation: + train_dataset = DreamBoothDataset( + instance_data_root=args.instance_data_dir, + instance_prompt=args.instance_prompt, + size=args.resolution, + repeats=args.repeats, + center_crop=args.center_crop, + buckets=buckets, + ) + batch_sampler = BucketBatchSampler(train_dataset, batch_size=args.train_batch_size, drop_last=True) + train_dataloader = torch.utils.data.DataLoader( + train_dataset, + batch_sampler=batch_sampler, + collate_fn=lambda examples: collate_fn(examples), + num_workers=args.dataloader_num_workers, + ) + + def compute_text_embeddings(prompt, text_encoding_pipeline): + with torch.no_grad(): + prompt_embeds, text_ids = text_encoding_pipeline.encode_prompt( + prompt=prompt, max_sequence_length=args.max_sequence_length + ) + return prompt_embeds, text_ids + + # If no type of tuning is done on the text_encoder and custom instance prompts are NOT + # provided (i.e. the --instance_prompt is used for all images), we encode the instance prompt once to avoid + # the redundant encoding. + if not train_dataset.custom_instance_prompts: + with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload): + instance_prompt_hidden_states, instance_text_ids = compute_text_embeddings( + args.instance_prompt, text_encoding_pipeline + ) + + if args.validation_prompt is not None: + validation_image = load_image(args.validation_image).convert("RGB") + validation_kwargs = {"image": validation_image} + with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload): + validation_kwargs["prompt_embeds"], _text_ids = compute_text_embeddings( + args.validation_prompt, text_encoding_pipeline + ) + validation_kwargs["negative_prompt_embeds"], _text_ids = compute_text_embeddings( + "", text_encoding_pipeline + ) + + # Init FSDP for text encoder + if args.fsdp_text_encoder: + fsdp_kwargs = get_fsdp_kwargs_from_accelerator(accelerator) + text_encoder_fsdp = wrap_with_fsdp( + model=text_encoding_pipeline.text_encoder, + device=accelerator.device, + offload=args.offload, + limit_all_gathers=True, + use_orig_params=True, + fsdp_kwargs=fsdp_kwargs, + ) + + text_encoding_pipeline.text_encoder = text_encoder_fsdp + dist.barrier() + + # If custom instance prompts are NOT provided (i.e. the instance prompt is used for all images), + # pack the statically computed variables appropriately here. This is so that we don't + # have to pass them to the dataloader. + if not train_dataset.custom_instance_prompts: + prompt_embeds = instance_prompt_hidden_states + text_ids = instance_text_ids + + # if cache_latents is set to True, we encode images to latents and store them. + # Similar to pre-encoding in the case of a single instance prompt, if custom prompts are provided + # we encode them in advance as well. + precompute_latents = args.cache_latents or train_dataset.custom_instance_prompts + if precompute_latents: + prompt_embeds_cache = [] + text_ids_cache = [] + latents_cache = [] + cond_latents_cache = [] + for batch in tqdm(train_dataloader, desc="Caching latents"): + with torch.no_grad(): + if args.cache_latents: + with offload_models(vae, device=accelerator.device, offload=args.offload): + batch["pixel_values"] = batch["pixel_values"].to( + accelerator.device, non_blocking=True, dtype=vae.dtype + ) + latents_cache.append(vae.encode(batch["pixel_values"]).latent_dist) + batch["cond_pixel_values"] = batch["cond_pixel_values"].to( + accelerator.device, non_blocking=True, dtype=vae.dtype + ) + cond_latents_cache.append(vae.encode(batch["cond_pixel_values"]).latent_dist) + if train_dataset.custom_instance_prompts: + if args.fsdp_text_encoder: + prompt_embeds, text_ids = compute_text_embeddings(batch["prompts"], text_encoding_pipeline) + else: + with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload): + prompt_embeds, text_ids = compute_text_embeddings(batch["prompts"], text_encoding_pipeline) + prompt_embeds_cache.append(prompt_embeds) + text_ids_cache.append(text_ids) + + # move back to cpu before deleting to ensure memory is freed see: https://github.com/huggingface/diffusers/issues/11376#issue-3008144624 + if args.cache_latents: + vae = vae.to("cpu") + del vae + + # move back to cpu before deleting to ensure memory is freed see: https://github.com/huggingface/diffusers/issues/11376#issue-3008144624 + text_encoding_pipeline = text_encoding_pipeline.to("cpu") + del text_encoder, tokenizer + free_memory() + + # Scheduler and math around the number of training steps. + # Check the PR https://github.com/huggingface/diffusers/pull/8312 for detailed explanation. + num_warmup_steps_for_scheduler = args.lr_warmup_steps * accelerator.num_processes + if args.max_train_steps is None: + len_train_dataloader_after_sharding = math.ceil(len(train_dataloader) / accelerator.num_processes) + num_update_steps_per_epoch = math.ceil(len_train_dataloader_after_sharding / args.gradient_accumulation_steps) + num_training_steps_for_scheduler = ( + args.num_train_epochs * accelerator.num_processes * num_update_steps_per_epoch + ) + else: + num_training_steps_for_scheduler = args.max_train_steps * accelerator.num_processes + + lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=optimizer, + num_warmup_steps=num_warmup_steps_for_scheduler, + num_training_steps=num_training_steps_for_scheduler, + num_cycles=args.lr_num_cycles, + power=args.lr_power, + ) + + # Prepare everything with our `accelerator`. + transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( + transformer, optimizer, train_dataloader, lr_scheduler + ) + + # We need to recalculate our total training steps as the size of the training dataloader may have changed. + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if args.max_train_steps is None: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + if num_training_steps_for_scheduler != args.max_train_steps: + logger.warning( + f"The length of the 'train_dataloader' after 'accelerator.prepare' ({len(train_dataloader)}) does not match " + f"the expected length ({len_train_dataloader_after_sharding}) when the learning rate scheduler was created. " + f"This inconsistency may result in the learning rate scheduler not functioning properly." + ) + # Afterwards we recalculate our number of training epochs + args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + + # We need to initialize the trackers we use, and also store our configuration. + # The trackers initializes automatically on the main process. + if accelerator.is_main_process: + tracker_name = "dreambooth-flux2-image2img-lora" + accelerator.init_trackers(tracker_name, config=vars(args)) + + # Train! + total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps + + logger.info("***** Running training *****") + logger.info(f" Num examples = {len(train_dataset)}") + logger.info(f" Num batches each epoch = {len(train_dataloader)}") + logger.info(f" Num Epochs = {args.num_train_epochs}") + logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") + logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") + logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") + logger.info(f" Total optimization steps = {args.max_train_steps}") + global_step = 0 + first_epoch = 0 + + # Potentially load in the weights and states from a previous save + if args.resume_from_checkpoint: + if args.resume_from_checkpoint != "latest": + path = os.path.basename(args.resume_from_checkpoint) + else: + # Get the mos recent checkpoint + dirs = os.listdir(args.output_dir) + dirs = [d for d in dirs if d.startswith("checkpoint")] + dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) + path = dirs[-1] if len(dirs) > 0 else None + + if path is None: + accelerator.print( + f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run." + ) + args.resume_from_checkpoint = None + initial_global_step = 0 + else: + accelerator.print(f"Resuming from checkpoint {path}") + accelerator.load_state(os.path.join(args.output_dir, path)) + global_step = int(path.split("-")[1]) + + initial_global_step = global_step + first_epoch = global_step // num_update_steps_per_epoch + + else: + initial_global_step = 0 + + progress_bar = tqdm( + range(0, args.max_train_steps), + initial=initial_global_step, + desc="Steps", + # Only show the progress bar once on each machine. + disable=not accelerator.is_local_main_process, + ) + + def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): + sigmas = noise_scheduler_copy.sigmas.to(device=accelerator.device, dtype=dtype) + schedule_timesteps = noise_scheduler_copy.timesteps.to(accelerator.device) + timesteps = timesteps.to(accelerator.device) + step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < n_dim: + sigma = sigma.unsqueeze(-1) + return sigma + + for epoch in range(first_epoch, args.num_train_epochs): + transformer.train() + + for step, batch in enumerate(train_dataloader): + models_to_accumulate = [transformer] + prompts = batch["prompts"] + + with accelerator.accumulate(models_to_accumulate): + if train_dataset.custom_instance_prompts: + prompt_embeds = prompt_embeds_cache[step] + text_ids = text_ids_cache[step] + else: + num_repeat_elements = len(prompts) + prompt_embeds = prompt_embeds.repeat(num_repeat_elements, 1, 1) + text_ids = text_ids.repeat(num_repeat_elements, 1, 1) + + # Convert images to latent space + if args.cache_latents: + model_input = latents_cache[step].mode() + cond_model_input = cond_latents_cache[step].mode() + else: + with offload_models(vae, device=accelerator.device, offload=args.offload): + pixel_values = batch["pixel_values"].to(dtype=vae.dtype) + cond_pixel_values = batch["cond_pixel_values"].to(dtype=vae.dtype) + + model_input = vae.encode(pixel_values).latent_dist.mode() + cond_model_input = vae.encode(cond_pixel_values).latent_dist.mode() + + model_input = Flux2KleinPipeline._patchify_latents(model_input) + model_input = (model_input - latents_bn_mean) / latents_bn_std + + cond_model_input = Flux2KleinPipeline._patchify_latents(cond_model_input) + cond_model_input = (cond_model_input - latents_bn_mean) / latents_bn_std + + model_input_ids = Flux2KleinPipeline._prepare_latent_ids(model_input).to(device=model_input.device) + cond_model_input_list = [cond_model_input[i].unsqueeze(0) for i in range(cond_model_input.shape[0])] + cond_model_input_ids = Flux2KleinPipeline._prepare_image_ids(cond_model_input_list).to( + device=cond_model_input.device + ) + cond_model_input_ids = cond_model_input_ids.view( + cond_model_input.shape[0], -1, model_input_ids.shape[-1] + ) + + # Sample noise that we'll add to the latents + noise = torch.randn_like(model_input) + bsz = model_input.shape[0] + + # Sample a random timestep for each image + # for weighting schemes where we sample timesteps non-uniformly + u = compute_density_for_timestep_sampling( + weighting_scheme=args.weighting_scheme, + batch_size=bsz, + logit_mean=args.logit_mean, + logit_std=args.logit_std, + mode_scale=args.mode_scale, + ) + indices = (u * noise_scheduler_copy.config.num_train_timesteps).long() + timesteps = noise_scheduler_copy.timesteps[indices].to(device=model_input.device) + + # Add noise according to flow matching. + # zt = (1 - texp) * x + texp * z1 + sigmas = get_sigmas(timesteps, n_dim=model_input.ndim, dtype=model_input.dtype) + noisy_model_input = (1.0 - sigmas) * model_input + sigmas * noise + + # [B, C, H, W] -> [B, H*W, C] + # concatenate the model inputs with the cond inputs + packed_noisy_model_input = Flux2KleinPipeline._pack_latents(noisy_model_input) + packed_cond_model_input = Flux2KleinPipeline._pack_latents(cond_model_input) + orig_input_shape = packed_noisy_model_input.shape + orig_input_ids_shape = model_input_ids.shape + + # concatenate the model inputs with the cond inputs + packed_noisy_model_input = torch.cat([packed_noisy_model_input, packed_cond_model_input], dim=1) + model_input_ids = torch.cat([model_input_ids, cond_model_input_ids], dim=1) + + # handle guidance + if unwrap_model(transformer).config.guidance_embeds: + guidance = torch.full([1], args.guidance_scale, device=accelerator.device) + guidance = guidance.expand(model_input.shape[0]) + else: + guidance = None + + # Predict the noise residual + model_pred = transformer( + hidden_states=packed_noisy_model_input, # (B, image_seq_len, C) + timestep=timesteps / 1000, + guidance=guidance, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, # B, text_seq_len, 4 + img_ids=model_input_ids, # B, image_seq_len, 4 + return_dict=False, + )[0] + # pruning the condition information + model_pred = model_pred[:, : orig_input_shape[1], :] + model_input_ids = model_input_ids[:, : orig_input_ids_shape[1], :] + + model_pred = Flux2KleinPipeline._unpack_latents_with_ids(model_pred, model_input_ids) + + # these weighting schemes use a uniform timestep sampling + # and instead post-weight the loss + weighting = compute_loss_weighting_for_sd3(weighting_scheme=args.weighting_scheme, sigmas=sigmas) + + # flow matching loss + target = noise - model_input + + # Compute regular loss. + loss = torch.mean( + (weighting.float() * (model_pred.float() - target.float()) ** 2).reshape(target.shape[0], -1), + 1, + ) + loss = loss.mean() + + accelerator.backward(loss) + if accelerator.sync_gradients: + params_to_clip = transformer.parameters() + accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm) + + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + # Checks if the accelerator has performed an optimization step behind the scenes + if accelerator.sync_gradients: + progress_bar.update(1) + global_step += 1 + + if accelerator.is_main_process or is_fsdp: + if global_step % args.checkpointing_steps == 0: + # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` + if args.checkpoints_total_limit is not None: + checkpoints = os.listdir(args.output_dir) + checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] + checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) + + # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints + if len(checkpoints) >= args.checkpoints_total_limit: + num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 + removing_checkpoints = checkpoints[0:num_to_remove] + + logger.info( + f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" + ) + logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") + + for removing_checkpoint in removing_checkpoints: + removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) + shutil.rmtree(removing_checkpoint) + + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + accelerator.save_state(save_path) + logger.info(f"Saved state to {save_path}") + + logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} + progress_bar.set_postfix(**logs) + accelerator.log(logs, step=global_step) + + if global_step >= args.max_train_steps: + break + + if accelerator.is_main_process: + if args.validation_prompt is not None and epoch % args.validation_epochs == 0: + # create pipeline + pipeline = Flux2KleinPipeline.from_pretrained( + args.pretrained_model_name_or_path, + text_encoder=None, + tokenizer=None, + transformer=unwrap_model(transformer), + revision=args.revision, + variant=args.variant, + torch_dtype=weight_dtype, + ) + images = log_validation( + pipeline=pipeline, + args=args, + accelerator=accelerator, + pipeline_args=validation_kwargs, + epoch=epoch, + torch_dtype=weight_dtype, + ) + + del pipeline + free_memory() + + # Save the lora layers + accelerator.wait_for_everyone() + + if is_fsdp: + transformer = unwrap_model(transformer) + state_dict = accelerator.get_state_dict(transformer) + if accelerator.is_main_process: + modules_to_save = {} + if is_fsdp: + if args.bnb_quantization_config_path is None: + if args.upcast_before_saving: + state_dict = { + k: v.to(torch.float32) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items() + } + else: + state_dict = { + k: v.to(weight_dtype) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items() + } + + transformer_lora_layers = get_peft_model_state_dict( + transformer, + state_dict=state_dict, + ) + transformer_lora_layers = { + k: v.detach().cpu().contiguous() if isinstance(v, torch.Tensor) else v + for k, v in transformer_lora_layers.items() + } + + else: + transformer = unwrap_model(transformer) + if args.bnb_quantization_config_path is None: + if args.upcast_before_saving: + transformer.to(torch.float32) + else: + transformer = transformer.to(weight_dtype) + transformer_lora_layers = get_peft_model_state_dict(transformer) + + modules_to_save["transformer"] = transformer + + Flux2KleinPipeline.save_lora_weights( + save_directory=args.output_dir, + transformer_lora_layers=transformer_lora_layers, + **_collate_lora_metadata(modules_to_save), + ) + + images = [] + run_validation = (args.validation_prompt and args.num_validation_images > 0) or (args.final_validation_prompt) + should_run_final_inference = not args.skip_final_inference and run_validation + if should_run_final_inference: + pipeline = Flux2KleinPipeline.from_pretrained( + args.pretrained_model_name_or_path, + revision=args.revision, + variant=args.variant, + torch_dtype=weight_dtype, + ) + # load attention processors + pipeline.load_lora_weights(args.output_dir) + + # run inference + images = [] + if args.validation_prompt and args.num_validation_images > 0: + images = log_validation( + pipeline=pipeline, + args=args, + accelerator=accelerator, + pipeline_args=validation_kwargs, + epoch=epoch, + is_final_validation=True, + torch_dtype=weight_dtype, + ) + del pipeline + free_memory() + + validation_prompt = args.validation_prompt if args.validation_prompt else args.final_validation_prompt + save_model_card( + (args.hub_model_id or Path(args.output_dir).name) if not args.push_to_hub else repo_id, + images=images, + base_model=args.pretrained_model_name_or_path, + instance_prompt=args.instance_prompt, + validation_prompt=validation_prompt, + repo_folder=args.output_dir, + fp8_training=args.do_fp8_training, + ) + + if args.push_to_hub: + upload_folder( + repo_id=repo_id, + folder_path=args.output_dir, + commit_message="End of training", + ignore_patterns=["step_*", "epoch_*"], + ) + + accelerator.end_training() + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/examples/dreambooth/train_dreambooth_lora_flux_kontext.py b/examples/dreambooth/train_dreambooth_lora_flux_kontext.py index 1a6757810a80..dee65761e92b 100644 --- a/examples/dreambooth/train_dreambooth_lora_flux_kontext.py +++ b/examples/dreambooth/train_dreambooth_lora_flux_kontext.py @@ -92,7 +92,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_hidream.py b/examples/dreambooth/train_dreambooth_lora_hidream.py index 3abc7afcad2c..bd2fb8db2d21 100644 --- a/examples/dreambooth/train_dreambooth_lora_hidream.py +++ b/examples/dreambooth/train_dreambooth_lora_hidream.py @@ -75,7 +75,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_lumina2.py b/examples/dreambooth/train_dreambooth_lora_lumina2.py index a13c579718c7..48eba4c5041d 100644 --- a/examples/dreambooth/train_dreambooth_lora_lumina2.py +++ b/examples/dreambooth/train_dreambooth_lora_lumina2.py @@ -73,7 +73,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_qwen_image.py b/examples/dreambooth/train_dreambooth_lora_qwen_image.py index 53b01bf0cfc8..a1e2fa0f6052 100644 --- a/examples/dreambooth/train_dreambooth_lora_qwen_image.py +++ b/examples/dreambooth/train_dreambooth_lora_qwen_image.py @@ -93,7 +93,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) @@ -1467,7 +1467,8 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): else: num_repeat_elements = len(prompts) prompt_embeds = prompt_embeds.repeat(num_repeat_elements, 1, 1) - prompt_embeds_mask = prompt_embeds_mask.repeat(num_repeat_elements, 1) + if prompt_embeds_mask is not None: + prompt_embeds_mask = prompt_embeds_mask.repeat(num_repeat_elements, 1) # Convert images to latent space if args.cache_latents: model_input = latents_cache[step].sample() @@ -1513,14 +1514,12 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): height=model_input.shape[3], width=model_input.shape[4], ) - print(f"{prompt_embeds_mask.sum(dim=1).tolist()=}") model_pred = transformer( hidden_states=packed_noisy_model_input, encoder_hidden_states=prompt_embeds, encoder_hidden_states_mask=prompt_embeds_mask, timestep=timesteps / 1000, img_shapes=img_shapes, - txt_seq_lens=prompt_embeds_mask.sum(dim=1).tolist(), return_dict=False, )[0] model_pred = QwenImagePipeline._unpack_latents( diff --git a/examples/dreambooth/train_dreambooth_lora_sana.py b/examples/dreambooth/train_dreambooth_lora_sana.py index 0afc31cf8a9a..3b295163b73d 100644 --- a/examples/dreambooth/train_dreambooth_lora_sana.py +++ b/examples/dreambooth/train_dreambooth_lora_sana.py @@ -91,7 +91,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_sd3.py b/examples/dreambooth/train_dreambooth_lora_sd3.py index d6770c805d25..4f49ef4bd801 100644 --- a/examples/dreambooth/train_dreambooth_lora_sd3.py +++ b/examples/dreambooth/train_dreambooth_lora_sd3.py @@ -73,7 +73,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_sdxl.py b/examples/dreambooth/train_dreambooth_lora_sdxl.py index 51bac5d59667..502ce1a3f1ec 100644 --- a/examples/dreambooth/train_dreambooth_lora_sdxl.py +++ b/examples/dreambooth/train_dreambooth_lora_sdxl.py @@ -80,7 +80,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_z_image.py b/examples/dreambooth/train_dreambooth_lora_z_image.py new file mode 100644 index 000000000000..623ae4d2aca3 --- /dev/null +++ b/examples/dreambooth/train_dreambooth_lora_z_image.py @@ -0,0 +1,1912 @@ +#!/usr/bin/env python +# coding=utf-8 +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# /// script +# dependencies = [ +# "diffusers @ git+https://github.com/huggingface/diffusers.git", +# "torch>=2.0.0", +# "accelerate>=0.31.0", +# "transformers>=4.41.2", +# "ftfy", +# "tensorboard", +# "Jinja2", +# "peft>=0.11.1", +# "sentencepiece", +# "torchvision", +# "datasets", +# "bitsandbytes", +# "prodigyopt", +# ] +# /// + +import argparse +import copy +import itertools +import json +import logging +import math +import os +import random +import shutil +import warnings +from contextlib import nullcontext +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import transformers +from accelerate import Accelerator +from accelerate.logging import get_logger +from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration, set_seed +from huggingface_hub import create_repo, upload_folder +from huggingface_hub.utils import insecure_hashlib +from peft import LoraConfig, prepare_model_for_kbit_training, set_peft_model_state_dict +from peft.utils import get_peft_model_state_dict +from PIL import Image +from PIL.ImageOps import exif_transpose +from torch.utils.data import Dataset +from torch.utils.data.sampler import BatchSampler +from torchvision import transforms +from torchvision.transforms import functional as TF +from tqdm.auto import tqdm +from transformers import Qwen2Tokenizer, Qwen3Model + +import diffusers +from diffusers import ( + AutoencoderKL, + BitsAndBytesConfig, + FlowMatchEulerDiscreteScheduler, + ZImagePipeline, + ZImageTransformer2DModel, +) +from diffusers.optimization import get_scheduler +from diffusers.training_utils import ( + _collate_lora_metadata, + _to_cpu_contiguous, + cast_training_params, + compute_density_for_timestep_sampling, + compute_loss_weighting_for_sd3, + find_nearest_bucket, + free_memory, + get_fsdp_kwargs_from_accelerator, + offload_models, + parse_buckets_string, + wrap_with_fsdp, +) +from diffusers.utils import ( + check_min_version, + convert_unet_state_dict_to_peft, + is_wandb_available, +) +from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card +from diffusers.utils.import_utils import is_torch_npu_available +from diffusers.utils.torch_utils import is_compiled_module + + +if getattr(torch, "distributed", None) is not None: + import torch.distributed as dist + +if is_wandb_available(): + import wandb + +# Will error if the minimal version of diffusers is not installed. Remove at your own risks. +check_min_version("0.38.0.dev0") + +logger = get_logger(__name__) + + +def save_model_card( + repo_id: str, + images=None, + base_model: str = None, + instance_prompt=None, + validation_prompt=None, + repo_folder=None, + quant_training=None, +): + widget_dict = [] + if images is not None: + for i, image in enumerate(images): + image.save(os.path.join(repo_folder, f"image_{i}.png")) + widget_dict.append( + {"text": validation_prompt if validation_prompt else " ", "output": {"url": f"image_{i}.png"}} + ) + + model_description = f""" +# Z Image DreamBooth LoRA - {repo_id} + + + +## Model description + +These are {repo_id} DreamBooth LoRA weights for {base_model}. + +The weights were trained using [DreamBooth](https://dreambooth.github.io/) with the [Z Image diffusers trainer](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/README_z_image.md). + +Quant training? {quant_training} + +## Trigger words + +You should use `{instance_prompt}` to trigger the image generation. + +## Download model + +[Download the *.safetensors LoRA]({repo_id}/tree/main) in the Files & versions tab. + +## Use it with the [🧨 diffusers library](https://github.com/huggingface/diffusers) + +```py +from diffusers import AutoPipelineForText2Image +import torch +pipeline = AutoPipelineForText2Image.from_pretrained("Tongyi-MAI/Z-Image", torch_dtype=torch.bfloat16).to('cuda') +pipeline.load_lora_weights('{repo_id}', weight_name='pytorch_lora_weights.safetensors') +image = pipeline('{validation_prompt if validation_prompt else instance_prompt}').images[0] +``` + +For more details, including weighting, merging and fusing LoRAs, check the [documentation on loading LoRAs in diffusers](https://huggingface.co/docs/diffusers/main/en/using-diffusers/loading_adapters) + +## License + +Apace License 2.0 +""" + model_card = load_or_create_model_card( + repo_id_or_path=repo_id, + from_training=True, + license="apache-2.0", + base_model=base_model, + prompt=instance_prompt, + model_description=model_description, + widget=widget_dict, + ) + tags = [ + "text-to-image", + "diffusers-training", + "diffusers", + "lora", + "z-image", + "template:sd-lora", + ] + + model_card = populate_model_card(model_card, tags=tags) + model_card.save(os.path.join(repo_folder, "README.md")) + + +def log_validation( + pipeline, + args, + accelerator, + pipeline_args, + epoch, + torch_dtype, + is_final_validation=False, +): + args.num_validation_images = args.num_validation_images if args.num_validation_images else 1 + logger.info( + f"Running validation... \n Generating {args.num_validation_images} images with prompt:" + f" {args.validation_prompt}." + ) + pipeline = pipeline.to(dtype=torch_dtype) + pipeline.enable_model_cpu_offload() + pipeline.set_progress_bar_config(disable=True) + + # run inference + generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed is not None else None + autocast_ctx = torch.autocast(accelerator.device.type) if not is_final_validation else nullcontext() + + images = [] + for _ in range(args.num_validation_images): + with autocast_ctx: + image = pipeline( + prompt=args.validation_prompt, + prompt_embeds=pipeline_args["prompt_embeds"], + generator=generator, + ).images[0] + images.append(image) + + for tracker in accelerator.trackers: + phase_name = "test" if is_final_validation else "validation" + if tracker.name == "tensorboard": + np_images = np.stack([np.asarray(img) for img in images]) + tracker.writer.add_images(phase_name, np_images, epoch, dataformats="NHWC") + if tracker.name == "wandb": + tracker.log( + { + phase_name: [ + wandb.Image(image, caption=f"{i}: {args.validation_prompt}") for i, image in enumerate(images) + ] + } + ) + + del pipeline + free_memory() + + return images + + +def module_filter_fn(mod: torch.nn.Module, fqn: str): + # don't convert the output module + if fqn == "proj_out": + return False + # don't convert linear modules with weight dimensions not divisible by 16 + if isinstance(mod, torch.nn.Linear): + if mod.in_features % 16 != 0 or mod.out_features % 16 != 0: + return False + return True + + +def parse_args(input_args=None): + parser = argparse.ArgumentParser(description="Simple example of a training script.") + parser.add_argument( + "--pretrained_model_name_or_path", + type=str, + default=None, + required=True, + help="Path to pretrained model or model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--revision", + type=str, + default=None, + required=False, + help="Revision of pretrained model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--bnb_quantization_config_path", + type=str, + default=None, + help="Quantization config in a JSON file that will be used to define the bitsandbytes quant config of the DiT.", + ) + parser.add_argument( + "--do_fp8_training", + action="store_true", + help="if we are doing FP8 training.", + ) + parser.add_argument( + "--variant", + type=str, + default=None, + help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16", + ) + parser.add_argument( + "--dataset_name", + type=str, + default=None, + help=( + "The name of the Dataset (from the HuggingFace hub) containing the training data of instance images (could be your own, possibly private," + " dataset). It can also be a path pointing to a local copy of a dataset in your filesystem," + " or to a folder containing files that 🤗 Datasets can understand." + ), + ) + parser.add_argument( + "--dataset_config_name", + type=str, + default=None, + help="The config of the Dataset, leave as None if there's only one config.", + ) + parser.add_argument( + "--instance_data_dir", + type=str, + default=None, + help=("A folder containing the training data. "), + ) + + parser.add_argument( + "--cache_dir", + type=str, + default=None, + help="The directory where the downloaded models and datasets will be stored.", + ) + + parser.add_argument( + "--image_column", + type=str, + default="image", + help="The column of the dataset containing the target image. By " + "default, the standard Image Dataset maps out 'file_name' " + "to 'image'.", + ) + parser.add_argument( + "--caption_column", + type=str, + default=None, + help="The column of the dataset containing the instance prompt for each image", + ) + + parser.add_argument("--repeats", type=int, default=1, help="How many times to repeat the training data.") + + parser.add_argument( + "--class_data_dir", + type=str, + default=None, + required=False, + help="A folder containing the training data of class images.", + ) + parser.add_argument( + "--instance_prompt", + type=str, + default=None, + required=True, + help="The prompt with identifier specifying the instance, e.g. 'photo of a TOK dog', 'in the style of TOK'", + ) + parser.add_argument( + "--class_prompt", + type=str, + default=None, + help="The prompt to specify images in the same class as provided instance images.", + ) + parser.add_argument( + "--max_sequence_length", + type=int, + default=512, + help="Maximum sequence length to use with with the T5 text encoder", + ) + + parser.add_argument( + "--validation_prompt", + type=str, + default=None, + help="A prompt that is used during validation to verify that the model is learning.", + ) + parser.add_argument( + "--skip_final_inference", + default=False, + action="store_true", + help="Whether to skip the final inference step with loaded lora weights upon training completion. This will run intermediate validation inference if `validation_prompt` is provided. Specify to reduce memory.", + ) + parser.add_argument( + "--final_validation_prompt", + type=str, + default=None, + help="A prompt that is used during a final validation to verify that the model is learning. Ignored if `--validation_prompt` is provided.", + ) + parser.add_argument( + "--num_validation_images", + type=int, + default=4, + help="Number of images that should be generated during validation with `validation_prompt`.", + ) + parser.add_argument( + "--validation_epochs", + type=int, + default=50, + help=( + "Run dreambooth validation every X epochs. Dreambooth validation consists of running the prompt" + " `args.validation_prompt` multiple times: `args.num_validation_images`." + ), + ) + parser.add_argument( + "--rank", + type=int, + default=4, + help=("The dimension of the LoRA update matrices."), + ) + parser.add_argument( + "--lora_alpha", + type=int, + default=4, + help="LoRA alpha to be used for additional scaling.", + ) + parser.add_argument("--lora_dropout", type=float, default=0.0, help="Dropout probability for LoRA layers") + + parser.add_argument( + "--with_prior_preservation", + default=False, + action="store_true", + help="Flag to add prior preservation loss.", + ) + parser.add_argument("--prior_loss_weight", type=float, default=1.0, help="The weight of prior preservation loss.") + parser.add_argument( + "--num_class_images", + type=int, + default=100, + help=( + "Minimal class images for prior preservation loss. If there are not enough images already present in" + " class_data_dir, additional images will be sampled with class_prompt." + ), + ) + parser.add_argument( + "--output_dir", + type=str, + default="z-image-dreambooth-lora", + help="The output directory where the model predictions and checkpoints will be written.", + ) + parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") + parser.add_argument( + "--resolution", + type=int, + default=512, + help=( + "The resolution for input images, all the images in the train/validation dataset will be resized to this" + " resolution" + ), + ) + parser.add_argument( + "--aspect_ratio_buckets", + type=str, + default=None, + help=( + "Aspect ratio buckets to use for training. Define as a string of 'h1,w1;h2,w2;...'. " + "e.g. '1024,1024;768,1360;1360,768;880,1168;1168,880;1248,832;832,1248'" + "Images will be resized and cropped to fit the nearest bucket. If provided, --resolution is ignored." + ), + ) + parser.add_argument( + "--center_crop", + default=False, + action="store_true", + help=( + "Whether to center crop the input images to the resolution. If not set, the images will be randomly" + " cropped. The images will be resized to the resolution first before cropping." + ), + ) + parser.add_argument( + "--random_flip", + action="store_true", + help="whether to randomly flip images horizontally", + ) + parser.add_argument( + "--train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader." + ) + parser.add_argument( + "--sample_batch_size", type=int, default=4, help="Batch size (per device) for sampling images." + ) + parser.add_argument("--num_train_epochs", type=int, default=1) + parser.add_argument( + "--max_train_steps", + type=int, + default=None, + help="Total number of training steps to perform. If provided, overrides num_train_epochs.", + ) + parser.add_argument( + "--checkpointing_steps", + type=int, + default=500, + help=( + "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final" + " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming" + " training using `--resume_from_checkpoint`." + ), + ) + parser.add_argument( + "--checkpoints_total_limit", + type=int, + default=None, + help=("Max number of checkpoints to store."), + ) + parser.add_argument( + "--resume_from_checkpoint", + type=str, + default=None, + help=( + "Whether training should be resumed from a previous checkpoint. Use a path saved by" + ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' + ), + ) + parser.add_argument( + "--gradient_accumulation_steps", + type=int, + default=1, + help="Number of updates steps to accumulate before performing a backward/update pass.", + ) + parser.add_argument( + "--gradient_checkpointing", + action="store_true", + help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.", + ) + parser.add_argument( + "--learning_rate", + type=float, + default=1e-4, + help="Initial learning rate (after the potential warmup period) to use.", + ) + + parser.add_argument( + "--guidance_scale", + type=float, + default=3.5, + help="the FLUX.1 dev variant is a guidance distilled model", + ) + + parser.add_argument( + "--text_encoder_lr", + type=float, + default=5e-6, + help="Text encoder learning rate to use.", + ) + parser.add_argument( + "--scale_lr", + action="store_true", + default=False, + help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.", + ) + parser.add_argument( + "--lr_scheduler", + type=str, + default="constant", + help=( + 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' + ' "constant", "constant_with_warmup"]' + ), + ) + parser.add_argument( + "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler." + ) + parser.add_argument( + "--lr_num_cycles", + type=int, + default=1, + help="Number of hard resets of the lr in cosine_with_restarts scheduler.", + ) + parser.add_argument("--lr_power", type=float, default=1.0, help="Power factor of the polynomial scheduler.") + parser.add_argument( + "--dataloader_num_workers", + type=int, + default=0, + help=( + "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process." + ), + ) + parser.add_argument( + "--weighting_scheme", + type=str, + default="none", + choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "none"], + help=('We default to the "none" weighting scheme for uniform sampling and uniform loss'), + ) + parser.add_argument( + "--logit_mean", type=float, default=0.0, help="mean to use when using the `'logit_normal'` weighting scheme." + ) + parser.add_argument( + "--logit_std", type=float, default=1.0, help="std to use when using the `'logit_normal'` weighting scheme." + ) + parser.add_argument( + "--mode_scale", + type=float, + default=1.29, + help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.", + ) + parser.add_argument( + "--optimizer", + type=str, + default="AdamW", + help=('The optimizer type to use. Choose between ["AdamW", "prodigy"]'), + ) + + parser.add_argument( + "--use_8bit_adam", + action="store_true", + help="Whether or not to use 8-bit Adam from bitsandbytes. Ignored if optimizer is not set to AdamW", + ) + + parser.add_argument( + "--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam and Prodigy optimizers." + ) + parser.add_argument( + "--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam and Prodigy optimizers." + ) + parser.add_argument( + "--prodigy_beta3", + type=float, + default=None, + help="coefficients for computing the Prodigy stepsize using running averages. If set to None, " + "uses the value of square root of beta2. Ignored if optimizer is adamW", + ) + parser.add_argument("--prodigy_decouple", type=bool, default=True, help="Use AdamW style decoupled weight decay") + parser.add_argument("--adam_weight_decay", type=float, default=1e-04, help="Weight decay to use for unet params") + parser.add_argument( + "--adam_weight_decay_text_encoder", type=float, default=1e-03, help="Weight decay to use for text_encoder" + ) + + parser.add_argument( + "--lora_layers", + type=str, + default=None, + help=( + 'The transformer modules to apply LoRA training on. Please specify the layers in a comma separated. E.g. - "to_k,to_q,to_v,to_out.0" will result in lora training of attention layers only' + ), + ) + + parser.add_argument( + "--adam_epsilon", + type=float, + default=1e-08, + help="Epsilon value for the Adam optimizer and Prodigy optimizers.", + ) + + parser.add_argument( + "--prodigy_use_bias_correction", + type=bool, + default=True, + help="Turn on Adam's bias correction. True by default. Ignored if optimizer is adamW", + ) + parser.add_argument( + "--prodigy_safeguard_warmup", + type=bool, + default=True, + help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage. True by default. " + "Ignored if optimizer is adamW", + ) + parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") + parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.") + parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.") + parser.add_argument( + "--hub_model_id", + type=str, + default=None, + help="The name of the repository to keep in sync with the local `output_dir`.", + ) + parser.add_argument( + "--logging_dir", + type=str, + default="logs", + help=( + "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to" + " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***." + ), + ) + parser.add_argument( + "--allow_tf32", + action="store_true", + help=( + "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" + " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" + ), + ) + parser.add_argument( + "--cache_latents", + action="store_true", + default=False, + help="Cache the VAE latents", + ) + parser.add_argument( + "--report_to", + type=str, + default="tensorboard", + help=( + 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' + ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' + ), + ) + parser.add_argument( + "--mixed_precision", + type=str, + default=None, + choices=["no", "fp16", "bf16"], + help=( + "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" + " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" + " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." + ), + ) + parser.add_argument( + "--upcast_before_saving", + action="store_true", + default=False, + help=( + "Whether to upcast the trained transformer layers to float32 before saving (at the end of training). " + "Defaults to precision dtype used for training to save memory" + ), + ) + parser.add_argument( + "--offload", + action="store_true", + help="Whether to offload the VAE and the text encoder to CPU when they are not used.", + ) + parser.add_argument( + "--prior_generation_precision", + type=str, + default=None, + choices=["no", "fp32", "fp16", "bf16"], + help=( + "Choose prior generation precision between fp32, fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" + " 1.10.and an Nvidia Ampere GPU. Default to fp16 if a GPU is available else fp32." + ), + ) + parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") + parser.add_argument("--enable_npu_flash_attention", action="store_true", help="Enabla Flash Attention for NPU") + parser.add_argument("--fsdp_text_encoder", action="store_true", help="Use FSDP for text encoder") + + if input_args is not None: + args = parser.parse_args(input_args) + else: + args = parser.parse_args() + + if args.dataset_name is None and args.instance_data_dir is None: + raise ValueError("Specify either `--dataset_name` or `--instance_data_dir`") + + if args.dataset_name is not None and args.instance_data_dir is not None: + raise ValueError("Specify only one of `--dataset_name` or `--instance_data_dir`") + if args.do_fp8_training and args.bnb_quantization_config_path: + raise ValueError("Both `do_fp8_training` and `bnb_quantization_config_path` cannot be passed.") + + env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) + if env_local_rank != -1 and env_local_rank != args.local_rank: + args.local_rank = env_local_rank + + if args.with_prior_preservation: + if args.class_data_dir is None: + raise ValueError("You must specify a data directory for class images.") + if args.class_prompt is None: + raise ValueError("You must specify prompt for class images.") + else: + # logger is not available yet + if args.class_data_dir is not None: + warnings.warn("You need not use --class_data_dir without --with_prior_preservation.") + if args.class_prompt is not None: + warnings.warn("You need not use --class_prompt without --with_prior_preservation.") + + return args + + +class DreamBoothDataset(Dataset): + """ + A dataset to prepare the instance and class images with the prompts for fine-tuning the model. + It pre-processes the images. + """ + + def __init__( + self, + instance_data_root, + instance_prompt, + class_prompt, + class_data_root=None, + class_num=None, + size=1024, + repeats=1, + center_crop=False, + buckets=None, + ): + self.size = size + self.center_crop = center_crop + + self.instance_prompt = instance_prompt + self.custom_instance_prompts = None + self.class_prompt = class_prompt + + self.buckets = buckets + + # if --dataset_name is provided or a metadata jsonl file is provided in the local --instance_data directory, + # we load the training data using load_dataset + if args.dataset_name is not None: + try: + from datasets import load_dataset + except ImportError: + raise ImportError( + "You are trying to load your data using the datasets library. If you wish to train using custom " + "captions please install the datasets library: `pip install datasets`. If you wish to load a " + "local folder containing images only, specify --instance_data_dir instead." + ) + # Downloading and loading a dataset from the hub. + # See more about loading custom images at + # https://huggingface.co/docs/datasets/v2.0.0/en/dataset_script + dataset = load_dataset( + args.dataset_name, + args.dataset_config_name, + cache_dir=args.cache_dir, + ) + # Preprocessing the datasets. + column_names = dataset["train"].column_names + + # 6. Get the column names for input/target. + if args.image_column is None: + image_column = column_names[0] + logger.info(f"image column defaulting to {image_column}") + else: + image_column = args.image_column + if image_column not in column_names: + raise ValueError( + f"`--image_column` value '{args.image_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}" + ) + instance_images = dataset["train"][image_column] + + if args.caption_column is None: + logger.info( + "No caption column provided, defaulting to instance_prompt for all images. If your dataset " + "contains captions/prompts for the images, make sure to specify the " + "column as --caption_column" + ) + self.custom_instance_prompts = None + else: + if args.caption_column not in column_names: + raise ValueError( + f"`--caption_column` value '{args.caption_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}" + ) + custom_instance_prompts = dataset["train"][args.caption_column] + # create final list of captions according to --repeats + self.custom_instance_prompts = [] + for caption in custom_instance_prompts: + self.custom_instance_prompts.extend(itertools.repeat(caption, repeats)) + else: + self.instance_data_root = Path(instance_data_root) + if not self.instance_data_root.exists(): + raise ValueError("Instance images root doesn't exists.") + + instance_images = [Image.open(path) for path in list(Path(instance_data_root).iterdir())] + self.custom_instance_prompts = None + + self.instance_images = [] + for img in instance_images: + self.instance_images.extend(itertools.repeat(img, repeats)) + + self.pixel_values = [] + for i, image in enumerate(self.instance_images): + image = exif_transpose(image) + if not image.mode == "RGB": + image = image.convert("RGB") + + width, height = image.size + + # Find the closest bucket + bucket_idx = find_nearest_bucket(height, width, self.buckets) + target_height, target_width = self.buckets[bucket_idx] + self.size = (target_height, target_width) + + # based on the bucket assignment, define the transformations + image = self.train_transform( + image, + size=self.size, + center_crop=args.center_crop, + random_flip=args.random_flip, + ) + self.pixel_values.append((image, bucket_idx)) + + self.num_instance_images = len(self.instance_images) + self._length = self.num_instance_images + + if class_data_root is not None: + self.class_data_root = Path(class_data_root) + self.class_data_root.mkdir(parents=True, exist_ok=True) + self.class_images_path = list(self.class_data_root.iterdir()) + if class_num is not None: + self.num_class_images = min(len(self.class_images_path), class_num) + else: + self.num_class_images = len(self.class_images_path) + self._length = max(self.num_class_images, self.num_instance_images) + else: + self.class_data_root = None + + self.image_transforms = transforms.Compose( + [ + transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR), + transforms.CenterCrop(size) if center_crop else transforms.RandomCrop(size), + transforms.ToTensor(), + transforms.Normalize([0.5], [0.5]), + ] + ) + + def __len__(self): + return self._length + + def __getitem__(self, index): + example = {} + instance_image, bucket_idx = self.pixel_values[index % self.num_instance_images] + example["instance_images"] = instance_image + example["bucket_idx"] = bucket_idx + if self.custom_instance_prompts: + caption = self.custom_instance_prompts[index % self.num_instance_images] + if caption: + example["instance_prompt"] = caption + else: + example["instance_prompt"] = self.instance_prompt + + else: # custom prompts were provided, but length does not match size of image dataset + example["instance_prompt"] = self.instance_prompt + + if self.class_data_root: + class_image = Image.open(self.class_images_path[index % self.num_class_images]) + class_image = exif_transpose(class_image) + + if not class_image.mode == "RGB": + class_image = class_image.convert("RGB") + example["class_images"] = self.image_transforms(class_image) + example["class_prompt"] = self.class_prompt + + return example + + def train_transform(self, image, size=(224, 224), center_crop=False, random_flip=False): + # 1. Resize (deterministic) + resize = transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR) + image = resize(image) + + # 2. Crop: either center or SAME random crop + if center_crop: + crop = transforms.CenterCrop(size) + image = crop(image) + else: + # get_params returns (i, j, h, w) + i, j, h, w = transforms.RandomCrop.get_params(image, output_size=size) + image = TF.crop(image, i, j, h, w) + + # 3. Random horizontal flip with the SAME coin flip + if random_flip: + do_flip = random.random() < 0.5 + if do_flip: + image = TF.hflip(image) + + # 4. ToTensor + Normalize (deterministic) + to_tensor = transforms.ToTensor() + normalize = transforms.Normalize([0.5], [0.5]) + image = normalize(to_tensor(image)) + + return image + + +def collate_fn(examples, with_prior_preservation=False): + pixel_values = [example["instance_images"] for example in examples] + prompts = [example["instance_prompt"] for example in examples] + + # Concat class and instance examples for prior preservation. + # We do this to avoid doing two forward passes. + if with_prior_preservation: + pixel_values += [example["class_images"] for example in examples] + prompts += [example["class_prompt"] for example in examples] + + pixel_values = torch.stack(pixel_values) + pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float() + + batch = {"pixel_values": pixel_values, "prompts": prompts} + return batch + + +class BucketBatchSampler(BatchSampler): + def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False): + if not isinstance(batch_size, int) or batch_size <= 0: + raise ValueError("batch_size should be a positive integer value, but got batch_size={}".format(batch_size)) + if not isinstance(drop_last, bool): + raise ValueError("drop_last should be a boolean value, but got drop_last={}".format(drop_last)) + + self.dataset = dataset + self.batch_size = batch_size + self.drop_last = drop_last + + # Group indices by bucket + self.bucket_indices = [[] for _ in range(len(self.dataset.buckets))] + for idx, (_, bucket_idx) in enumerate(self.dataset.pixel_values): + self.bucket_indices[bucket_idx].append(idx) + + self.sampler_len = 0 + self.batches = [] + + # Pre-generate batches for each bucket + for indices_in_bucket in self.bucket_indices: + # Shuffle indices within the bucket + random.shuffle(indices_in_bucket) + # Create batches + for i in range(0, len(indices_in_bucket), self.batch_size): + batch = indices_in_bucket[i : i + self.batch_size] + if len(batch) < self.batch_size and self.drop_last: + continue # Skip partial batch if drop_last is True + self.batches.append(batch) + self.sampler_len += 1 # Count the number of batches + + def __iter__(self): + # Shuffle the order of the batches each epoch + random.shuffle(self.batches) + for batch in self.batches: + yield batch + + def __len__(self): + return self.sampler_len + + +class PromptDataset(Dataset): + "A simple dataset to prepare the prompts to generate class images on multiple GPUs." + + def __init__(self, prompt, num_samples): + self.prompt = prompt + self.num_samples = num_samples + + def __len__(self): + return self.num_samples + + def __getitem__(self, index): + example = {} + example["prompt"] = self.prompt + example["index"] = index + return example + + +def main(args): + if args.report_to == "wandb" and args.hub_token is not None: + raise ValueError( + "You cannot use both --report_to=wandb and --hub_token due to a security risk of exposing your token." + " Please use `hf auth login` to authenticate with the Hub." + ) + + if torch.backends.mps.is_available() and args.mixed_precision == "bf16": + # due to pytorch#99272, MPS does not yet support bfloat16. + raise ValueError( + "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead." + ) + if args.do_fp8_training: + from torchao.float8 import Float8LinearConfig, convert_to_float8_training + + logging_dir = Path(args.output_dir, args.logging_dir) + + accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) + kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) + accelerator = Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + mixed_precision=args.mixed_precision, + log_with=args.report_to, + project_config=accelerator_project_config, + kwargs_handlers=[kwargs], + ) + + # Disable AMP for MPS. + if torch.backends.mps.is_available(): + accelerator.native_amp = False + + if args.report_to == "wandb": + if not is_wandb_available(): + raise ImportError("Make sure to install wandb if you want to use it for logging during training.") + + # Make one log on every process with the configuration for debugging. + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + ) + logger.info(accelerator.state, main_process_only=False) + if accelerator.is_local_main_process: + transformers.utils.logging.set_verbosity_warning() + diffusers.utils.logging.set_verbosity_info() + else: + transformers.utils.logging.set_verbosity_error() + diffusers.utils.logging.set_verbosity_error() + + # If passed along, set the training seed now. + if args.seed is not None: + set_seed(args.seed) + + # Generate class images if prior preservation is enabled. + if args.with_prior_preservation: + class_images_dir = Path(args.class_data_dir) + if not class_images_dir.exists(): + class_images_dir.mkdir(parents=True) + cur_class_images = len(list(class_images_dir.iterdir())) + + if cur_class_images < args.num_class_images: + has_supported_fp16_accelerator = torch.cuda.is_available() or torch.backends.mps.is_available() + torch_dtype = torch.float16 if has_supported_fp16_accelerator else torch.float32 + if args.prior_generation_precision == "fp32": + torch_dtype = torch.float32 + elif args.prior_generation_precision == "fp16": + torch_dtype = torch.float16 + elif args.prior_generation_precision == "bf16": + torch_dtype = torch.bfloat16 + + pipeline = ZImagePipeline.from_pretrained( + args.pretrained_model_name_or_path, + torch_dtype=torch_dtype, + revision=args.revision, + variant=args.variant, + ) + pipeline.set_progress_bar_config(disable=True) + + num_new_images = args.num_class_images - cur_class_images + logger.info(f"Number of class images to sample: {num_new_images}.") + + sample_dataset = PromptDataset(args.class_prompt, num_new_images) + sample_dataloader = torch.utils.data.DataLoader(sample_dataset, batch_size=args.sample_batch_size) + + sample_dataloader = accelerator.prepare(sample_dataloader) + pipeline.to(accelerator.device) + + for example in tqdm( + sample_dataloader, desc="Generating class images", disable=not accelerator.is_local_main_process + ): + with torch.autocast(device_type=accelerator.device.type, dtype=torch_dtype): + images = pipeline(prompt=example["prompt"]).images + + for i, image in enumerate(images): + hash_image = insecure_hashlib.sha1(image.tobytes()).hexdigest() + image_filename = class_images_dir / f"{example['index'][i] + cur_class_images}-{hash_image}.jpg" + image.save(image_filename) + + del pipeline + free_memory() + + # Handle the repository creation + if accelerator.is_main_process: + if args.output_dir is not None: + os.makedirs(args.output_dir, exist_ok=True) + + if args.push_to_hub: + repo_id = create_repo( + repo_id=args.hub_model_id or Path(args.output_dir).name, + exist_ok=True, + ).repo_id + + # Load the tokenizers + tokenizer = Qwen2Tokenizer.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="tokenizer", + revision=args.revision, + ) + + # For mixed precision training we cast all non-trainable weights (vae, text_encoder and transformer) to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + weight_dtype = torch.float32 + if accelerator.mixed_precision == "fp16": + weight_dtype = torch.float16 + elif accelerator.mixed_precision == "bf16": + weight_dtype = torch.bfloat16 + + # Load scheduler and models + noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="scheduler", + revision=args.revision, + ) + noise_scheduler_copy = copy.deepcopy(noise_scheduler) + vae = AutoencoderKL.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="vae", + revision=args.revision, + variant=args.variant, + ) + vae_config_shift_factor = vae.config.shift_factor + vae_config_scaling_factor = vae.config.scaling_factor + + quantization_config = None + if args.bnb_quantization_config_path is not None: + with open(args.bnb_quantization_config_path, "r") as f: + config_kwargs = json.load(f) + if "load_in_4bit" in config_kwargs and config_kwargs["load_in_4bit"]: + config_kwargs["bnb_4bit_compute_dtype"] = weight_dtype + quantization_config = BitsAndBytesConfig(**config_kwargs) + + transformer = ZImageTransformer2DModel.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="transformer", + revision=args.revision, + variant=args.variant, + quantization_config=quantization_config, + torch_dtype=weight_dtype, + ) + if args.bnb_quantization_config_path is not None: + transformer = prepare_model_for_kbit_training(transformer, use_gradient_checkpointing=False) + + text_encoder = Qwen3Model.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="text_encoder", + revision=args.revision, + variant=args.variant, + ) + text_encoder.requires_grad_(False) + + # We only train the additional adapter LoRA layers + transformer.requires_grad_(False) + vae.requires_grad_(False) + + if args.enable_npu_flash_attention: + if is_torch_npu_available(): + logger.info("npu flash attention enabled.") + transformer.set_attention_backend("_native_npu") + else: + raise ValueError("npu flash attention requires torch_npu extensions and is supported only on npu device ") + + if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16: + # due to pytorch#99272, MPS does not yet support bfloat16. + raise ValueError( + "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead." + ) + + to_kwargs = {"dtype": weight_dtype, "device": accelerator.device} if not args.offload else {"dtype": weight_dtype} + vae.to(**to_kwargs) + # we never offload the transformer to CPU, so we can just use the accelerator device + transformer_to_kwargs = ( + {"device": accelerator.device} + if args.bnb_quantization_config_path is not None + else {"device": accelerator.device, "dtype": weight_dtype} + ) + + is_fsdp = getattr(accelerator.state, "fsdp_plugin", None) is not None + if not is_fsdp: + transformer.to(**transformer_to_kwargs) + + if args.do_fp8_training: + convert_to_float8_training( + transformer, module_filter_fn=module_filter_fn, config=Float8LinearConfig(pad_inner_dim=True) + ) + + text_encoder.to(**to_kwargs) + # Initialize a text encoding pipeline and keep it to CPU for now. + text_encoding_pipeline = ZImagePipeline.from_pretrained( + args.pretrained_model_name_or_path, + vae=None, + transformer=None, + tokenizer=tokenizer, + text_encoder=text_encoder, + scheduler=None, + revision=args.revision, + ) + + if args.gradient_checkpointing: + transformer.enable_gradient_checkpointing() + + if args.lora_layers is not None: + target_modules = [layer.strip() for layer in args.lora_layers.split(",")] + else: + target_modules = ["to_k", "to_q", "to_v", "to_out.0"] + + # now we will add new LoRA weights the transformer layers + transformer_lora_config = LoraConfig( + r=args.rank, + lora_alpha=args.lora_alpha, + lora_dropout=args.lora_dropout, + init_lora_weights="gaussian", + target_modules=target_modules, + ) + transformer.add_adapter(transformer_lora_config) + + def unwrap_model(model): + model = accelerator.unwrap_model(model) + model = model._orig_mod if is_compiled_module(model) else model + return model + + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + def save_model_hook(models, weights, output_dir): + transformer_cls = type(unwrap_model(transformer)) + + # 1) Validate and pick the transformer model + modules_to_save: dict[str, Any] = {} + transformer_model = None + + for model in models: + if isinstance(unwrap_model(model), transformer_cls): + transformer_model = model + modules_to_save["transformer"] = model + else: + raise ValueError(f"unexpected save model: {model.__class__}") + + if transformer_model is None: + raise ValueError("No transformer model found in 'models'") + + # 2) Optionally gather FSDP state dict once + state_dict = accelerator.get_state_dict(model) if is_fsdp else None + + # 3) Only main process materializes the LoRA state dict + transformer_lora_layers_to_save = None + if accelerator.is_main_process: + peft_kwargs = {} + if is_fsdp: + peft_kwargs["state_dict"] = state_dict + + transformer_lora_layers_to_save = get_peft_model_state_dict( + unwrap_model(transformer_model) if is_fsdp else transformer_model, + **peft_kwargs, + ) + + if is_fsdp: + transformer_lora_layers_to_save = _to_cpu_contiguous(transformer_lora_layers_to_save) + + # make sure to pop weight so that corresponding model is not saved again + if weights: + weights.pop() + + ZImagePipeline.save_lora_weights( + output_dir, + transformer_lora_layers=transformer_lora_layers_to_save, + **_collate_lora_metadata(modules_to_save), + ) + + def load_model_hook(models, input_dir): + transformer_ = None + + if not is_fsdp: + while len(models) > 0: + model = models.pop() + + if isinstance(unwrap_model(model), type(unwrap_model(transformer))): + transformer_ = unwrap_model(model) + else: + raise ValueError(f"unexpected save model: {model.__class__}") + else: + transformer_ = ZImageTransformer2DModel.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="transformer", + ) + transformer_.add_adapter(transformer_lora_config) + + lora_state_dict = ZImagePipeline.lora_state_dict(input_dir) + + transformer_state_dict = { + f"{k.replace('transformer.', '')}": v for k, v in lora_state_dict.items() if k.startswith("transformer.") + } + transformer_state_dict = convert_unet_state_dict_to_peft(transformer_state_dict) + incompatible_keys = set_peft_model_state_dict(transformer_, transformer_state_dict, adapter_name="default") + if incompatible_keys is not None: + # check only for unexpected keys + unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None) + if unexpected_keys: + logger.warning( + f"Loading adapter weights from state_dict led to unexpected keys not found in the model: " + f" {unexpected_keys}. " + ) + + # Make sure the trainable params are in float32. This is again needed since the base models + # are in `weight_dtype`. More details: + # https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804 + if args.mixed_precision == "fp16": + models = [transformer_] + # only upcast trainable parameters (LoRA) into fp32 + cast_training_params(models) + + accelerator.register_save_state_pre_hook(save_model_hook) + accelerator.register_load_state_pre_hook(load_model_hook) + + # Enable TF32 for faster training on Ampere GPUs, + # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices + if args.allow_tf32 and torch.cuda.is_available(): + torch.backends.cuda.matmul.allow_tf32 = True + + if args.scale_lr: + args.learning_rate = ( + args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes + ) + + # Make sure the trainable params are in float32. + if args.mixed_precision == "fp16": + models = [transformer] + # only upcast trainable parameters (LoRA) into fp32 + cast_training_params(models, dtype=torch.float32) + + transformer_lora_parameters = list(filter(lambda p: p.requires_grad, transformer.parameters())) + + # Optimization parameters + transformer_parameters_with_lr = {"params": transformer_lora_parameters, "lr": args.learning_rate} + params_to_optimize = [transformer_parameters_with_lr] + + # Optimizer creation + if not (args.optimizer.lower() == "prodigy" or args.optimizer.lower() == "adamw"): + logger.warning( + f"Unsupported choice of optimizer: {args.optimizer}.Supported optimizers include [adamW, prodigy]." + "Defaulting to adamW" + ) + args.optimizer = "adamw" + + if args.use_8bit_adam and not args.optimizer.lower() == "adamw": + logger.warning( + f"use_8bit_adam is ignored when optimizer is not set to 'AdamW'. Optimizer was " + f"set to {args.optimizer.lower()}" + ) + + if args.optimizer.lower() == "adamw": + if args.use_8bit_adam: + try: + import bitsandbytes as bnb + except ImportError: + raise ImportError( + "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`." + ) + + optimizer_class = bnb.optim.AdamW8bit + else: + optimizer_class = torch.optim.AdamW + + optimizer = optimizer_class( + params_to_optimize, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + + if args.optimizer.lower() == "prodigy": + try: + import prodigyopt + except ImportError: + raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`") + + optimizer_class = prodigyopt.Prodigy + + if args.learning_rate <= 0.1: + logger.warning( + "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0" + ) + + optimizer = optimizer_class( + params_to_optimize, + betas=(args.adam_beta1, args.adam_beta2), + beta3=args.prodigy_beta3, + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + decouple=args.prodigy_decouple, + use_bias_correction=args.prodigy_use_bias_correction, + safeguard_warmup=args.prodigy_safeguard_warmup, + ) + + if args.aspect_ratio_buckets is not None: + buckets = parse_buckets_string(args.aspect_ratio_buckets) + else: + buckets = [(args.resolution, args.resolution)] + logger.info(f"Using parsed aspect ratio buckets: {buckets}") + + # Dataset and DataLoaders creation: + train_dataset = DreamBoothDataset( + instance_data_root=args.instance_data_dir, + instance_prompt=args.instance_prompt, + class_prompt=args.class_prompt, + class_data_root=args.class_data_dir if args.with_prior_preservation else None, + class_num=args.num_class_images, + size=args.resolution, + repeats=args.repeats, + center_crop=args.center_crop, + buckets=buckets, + ) + batch_sampler = BucketBatchSampler(train_dataset, batch_size=args.train_batch_size, drop_last=True) + train_dataloader = torch.utils.data.DataLoader( + train_dataset, + batch_sampler=batch_sampler, + collate_fn=lambda examples: collate_fn(examples, args.with_prior_preservation), + num_workers=args.dataloader_num_workers, + ) + + def compute_text_embeddings(prompt, text_encoding_pipeline): + with torch.no_grad(): + prompt_embeds, _ = text_encoding_pipeline.encode_prompt( + prompt=prompt, + max_sequence_length=args.max_sequence_length, + ) + return prompt_embeds + + # If no type of tuning is done on the text_encoder and custom instance prompts are NOT + # provided (i.e. the --instance_prompt is used for all images), we encode the instance prompt once to avoid + # the redundant encoding. + if not train_dataset.custom_instance_prompts: + with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload): + instance_prompt_hidden_states = compute_text_embeddings(args.instance_prompt, text_encoding_pipeline) + + # Handle class prompt for prior-preservation. + if args.with_prior_preservation: + with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload): + class_prompt_hidden_states = compute_text_embeddings(args.class_prompt, text_encoding_pipeline) + validation_embeddings = {} + if args.validation_prompt is not None: + with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload): + validation_embeddings["prompt_embeds"] = compute_text_embeddings( + args.validation_prompt, text_encoding_pipeline + ) + + # Init FSDP for text encoder + if args.fsdp_text_encoder: + fsdp_kwargs = get_fsdp_kwargs_from_accelerator(accelerator) + text_encoder_fsdp = wrap_with_fsdp( + model=text_encoding_pipeline.text_encoder, + device=accelerator.device, + offload=args.offload, + limit_all_gathers=True, + use_orig_params=True, + fsdp_kwargs=fsdp_kwargs, + ) + + text_encoding_pipeline.text_encoder = text_encoder_fsdp + dist.barrier() + + # If custom instance prompts are NOT provided (i.e. the instance prompt is used for all images), + # pack the statically computed variables appropriately here. This is so that we don't + # have to pass them to the dataloader. + if not train_dataset.custom_instance_prompts: + prompt_embeds = instance_prompt_hidden_states + if args.with_prior_preservation: + prompt_embeds = torch.cat([prompt_embeds, class_prompt_hidden_states], dim=0) + + # if cache_latents is set to True, we encode images to latents and store them. + # Similar to pre-encoding in the case of a single instance prompt, if custom prompts are provided + # we encode them in advance as well. + precompute_latents = args.cache_latents or train_dataset.custom_instance_prompts + if precompute_latents: + prompt_embeds_cache = [] + latents_cache = [] + for batch in tqdm(train_dataloader, desc="Caching latents"): + with torch.no_grad(): + if args.cache_latents: + with offload_models(vae, device=accelerator.device, offload=args.offload): + batch["pixel_values"] = batch["pixel_values"].to( + accelerator.device, non_blocking=True, dtype=vae.dtype + ) + latents_cache.append(vae.encode(batch["pixel_values"]).latent_dist) + if train_dataset.custom_instance_prompts: + if args.fsdp_text_encoder: + prompt_embeds = compute_text_embeddings(batch["prompts"], text_encoding_pipeline) + else: + with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload): + prompt_embeds = compute_text_embeddings(batch["prompts"], text_encoding_pipeline) + prompt_embeds_cache.append(prompt_embeds) + + # move back to cpu before deleting to ensure memory is freed see: https://github.com/huggingface/diffusers/issues/11376#issue-3008144624 + if args.cache_latents: + vae = vae.to("cpu") + del vae + + # move back to cpu before deleting to ensure memory is freed see: https://github.com/huggingface/diffusers/issues/11376#issue-3008144624 + text_encoding_pipeline = text_encoding_pipeline.to("cpu") + del text_encoder, tokenizer + free_memory() + + # Scheduler and math around the number of training steps. + # Check the PR https://github.com/huggingface/diffusers/pull/8312 for detailed explanation. + num_warmup_steps_for_scheduler = args.lr_warmup_steps * accelerator.num_processes + if args.max_train_steps is None: + len_train_dataloader_after_sharding = math.ceil(len(train_dataloader) / accelerator.num_processes) + num_update_steps_per_epoch = math.ceil(len_train_dataloader_after_sharding / args.gradient_accumulation_steps) + num_training_steps_for_scheduler = ( + args.num_train_epochs * accelerator.num_processes * num_update_steps_per_epoch + ) + else: + num_training_steps_for_scheduler = args.max_train_steps * accelerator.num_processes + + lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=optimizer, + num_warmup_steps=num_warmup_steps_for_scheduler, + num_training_steps=num_training_steps_for_scheduler, + num_cycles=args.lr_num_cycles, + power=args.lr_power, + ) + + # Prepare everything with our `accelerator`. + transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( + transformer, optimizer, train_dataloader, lr_scheduler + ) + + # We need to recalculate our total training steps as the size of the training dataloader may have changed. + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if args.max_train_steps is None: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + if num_training_steps_for_scheduler != args.max_train_steps: + logger.warning( + f"The length of the 'train_dataloader' after 'accelerator.prepare' ({len(train_dataloader)}) does not match " + f"the expected length ({len_train_dataloader_after_sharding}) when the learning rate scheduler was created. " + f"This inconsistency may result in the learning rate scheduler not functioning properly." + ) + # Afterwards we recalculate our number of training epochs + args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + + # We need to initialize the trackers we use, and also store our configuration. + # The trackers initializes automatically on the main process. + if accelerator.is_main_process: + tracker_name = "dreambooth-z-image-lora" + args_cp = vars(args).copy() + accelerator.init_trackers(tracker_name, config=args_cp) + + # Train! + total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps + + logger.info("***** Running training *****") + logger.info(f" Num examples = {len(train_dataset)}") + logger.info(f" Num batches each epoch = {len(train_dataloader)}") + logger.info(f" Num Epochs = {args.num_train_epochs}") + logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") + logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") + logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") + logger.info(f" Total optimization steps = {args.max_train_steps}") + global_step = 0 + first_epoch = 0 + + # Potentially load in the weights and states from a previous save + if args.resume_from_checkpoint: + if args.resume_from_checkpoint != "latest": + path = os.path.basename(args.resume_from_checkpoint) + else: + # Get the mos recent checkpoint + dirs = os.listdir(args.output_dir) + dirs = [d for d in dirs if d.startswith("checkpoint")] + dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) + path = dirs[-1] if len(dirs) > 0 else None + + if path is None: + accelerator.print( + f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run." + ) + args.resume_from_checkpoint = None + initial_global_step = 0 + else: + accelerator.print(f"Resuming from checkpoint {path}") + accelerator.load_state(os.path.join(args.output_dir, path)) + global_step = int(path.split("-")[1]) + + initial_global_step = global_step + first_epoch = global_step // num_update_steps_per_epoch + + else: + initial_global_step = 0 + + progress_bar = tqdm( + range(0, args.max_train_steps), + initial=initial_global_step, + desc="Steps", + # Only show the progress bar once on each machine. + disable=not accelerator.is_local_main_process, + ) + + def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): + sigmas = noise_scheduler_copy.sigmas.to(device=accelerator.device, dtype=dtype) + schedule_timesteps = noise_scheduler_copy.timesteps.to(accelerator.device) + timesteps = timesteps.to(accelerator.device) + step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < n_dim: + sigma = sigma.unsqueeze(-1) + return sigma + + for epoch in range(first_epoch, args.num_train_epochs): + transformer.train() + + for step, batch in enumerate(train_dataloader): + models_to_accumulate = [transformer] + prompts = batch["prompts"] + + with accelerator.accumulate(models_to_accumulate): + if train_dataset.custom_instance_prompts: + prompt_embeds = prompt_embeds_cache[step] + else: + num_repeat_elements = len(prompts) + prompt_embeds = [pe for pe in prompt_embeds for _ in range(num_repeat_elements)] + + # Convert images to latent space + if args.cache_latents: + model_input = latents_cache[step].mode() + else: + with offload_models(vae, device=accelerator.device, offload=args.offload): + pixel_values = batch["pixel_values"].to(dtype=vae.dtype) + model_input = vae.encode(pixel_values).latent_dist.mode() + + model_input = (model_input - vae_config_shift_factor) * vae_config_scaling_factor + # Sample noise that we'll add to the latents + noise = torch.randn_like(model_input) + bsz = model_input.shape[0] + + # Sample a random timestep for each image + # for weighting schemes where we sample timesteps non-uniformly + u = compute_density_for_timestep_sampling( + weighting_scheme=args.weighting_scheme, + batch_size=bsz, + logit_mean=args.logit_mean, + logit_std=args.logit_std, + mode_scale=args.mode_scale, + ) + indices = (u * noise_scheduler_copy.config.num_train_timesteps).long() + timesteps = noise_scheduler_copy.timesteps[indices].to(device=model_input.device) + + # Add noise according to flow matching. + # zt = (1 - texp) * x + texp * z1 + sigmas = get_sigmas(timesteps, n_dim=model_input.ndim, dtype=model_input.dtype) + noisy_model_input = (1.0 - sigmas) * model_input + sigmas * noise + + timestep_normalized = (1000 - timesteps) / 1000 + + noisy_model_input_5d = noisy_model_input.unsqueeze(2) # (B, C, H, W) -> (B, C, 1, H, W) + noisy_model_input_list = list(noisy_model_input_5d.unbind(dim=0)) # List of (C, 1, H, W) + + model_pred_list = transformer( + noisy_model_input_list, + timestep_normalized, + prompt_embeds, # This is a List[torch.Tensor] for Z-Image + return_dict=False, + )[0] + model_pred = torch.stack(model_pred_list, dim=0) # (B, C, 1, H, W) + model_pred = model_pred.squeeze(2) # (B, C, H, W) + model_pred = -model_pred # z-Image negates the prediction + + # these weighting schemes use a uniform timestep sampling + # and instead post-weight the loss + weighting = compute_loss_weighting_for_sd3(weighting_scheme=args.weighting_scheme, sigmas=sigmas) + + # flow matching loss + target = noise - model_input + + if args.with_prior_preservation: + # Chunk the noise and model_pred into two parts and compute the loss on each part separately. + model_pred, model_pred_prior = torch.chunk(model_pred, 2, dim=0) + target, target_prior = torch.chunk(target, 2, dim=0) + + # Compute prior loss + prior_loss = torch.mean( + (weighting.float() * (model_pred_prior.float() - target_prior.float()) ** 2).reshape( + target_prior.shape[0], -1 + ), + 1, + ) + prior_loss = prior_loss.mean() + + # Compute regular loss. + loss = torch.mean( + (weighting.float() * (model_pred.float() - target.float()) ** 2).reshape(target.shape[0], -1), + 1, + ) + loss = loss.mean() + + if args.with_prior_preservation: + # Add the prior loss to the instance loss. + loss = loss + args.prior_loss_weight * prior_loss + + accelerator.backward(loss) + if accelerator.sync_gradients: + params_to_clip = transformer.parameters() + accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm) + + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + # Checks if the accelerator has performed an optimization step behind the scenes + if accelerator.sync_gradients: + progress_bar.update(1) + global_step += 1 + + if accelerator.is_main_process or is_fsdp: + if global_step % args.checkpointing_steps == 0: + # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` + if args.checkpoints_total_limit is not None: + checkpoints = os.listdir(args.output_dir) + checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] + checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) + + # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints + if len(checkpoints) >= args.checkpoints_total_limit: + num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 + removing_checkpoints = checkpoints[0:num_to_remove] + + logger.info( + f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" + ) + logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") + + for removing_checkpoint in removing_checkpoints: + removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) + shutil.rmtree(removing_checkpoint) + + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + accelerator.save_state(save_path) + logger.info(f"Saved state to {save_path}") + + logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} + progress_bar.set_postfix(**logs) + accelerator.log(logs, step=global_step) + + if global_step >= args.max_train_steps: + break + + if accelerator.is_main_process: + if args.validation_prompt is not None and epoch % args.validation_epochs == 0: + # create pipeline + pipeline = ZImagePipeline.from_pretrained( + args.pretrained_model_name_or_path, + transformer=unwrap_model(transformer), + revision=args.revision, + variant=args.variant, + torch_dtype=weight_dtype, + ) + images = log_validation( + pipeline=pipeline, + args=args, + accelerator=accelerator, + pipeline_args=validation_embeddings, + epoch=epoch, + torch_dtype=weight_dtype, + ) + + del pipeline + free_memory() + + # Save the lora layers + accelerator.wait_for_everyone() + + if is_fsdp: + transformer = unwrap_model(transformer) + state_dict = accelerator.get_state_dict(transformer) + if accelerator.is_main_process: + modules_to_save = {} + if is_fsdp: + if args.bnb_quantization_config_path is None: + if args.upcast_before_saving: + state_dict = { + k: v.to(torch.float32) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items() + } + else: + state_dict = { + k: v.to(weight_dtype) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items() + } + + transformer_lora_layers = get_peft_model_state_dict( + transformer, + state_dict=state_dict, + ) + transformer_lora_layers = { + k: v.detach().cpu().contiguous() if isinstance(v, torch.Tensor) else v + for k, v in transformer_lora_layers.items() + } + + else: + transformer = unwrap_model(transformer) + if args.bnb_quantization_config_path is None: + if args.upcast_before_saving: + transformer.to(torch.float32) + else: + transformer = transformer.to(weight_dtype) + transformer_lora_layers = get_peft_model_state_dict(transformer) + + modules_to_save["transformer"] = transformer + + ZImagePipeline.save_lora_weights( + save_directory=args.output_dir, + transformer_lora_layers=transformer_lora_layers, + **_collate_lora_metadata(modules_to_save), + ) + + images = [] + run_validation = (args.validation_prompt and args.num_validation_images > 0) or (args.final_validation_prompt) + should_run_final_inference = not args.skip_final_inference and run_validation + if should_run_final_inference: + pipeline = ZImagePipeline.from_pretrained( + args.pretrained_model_name_or_path, + revision=args.revision, + variant=args.variant, + torch_dtype=weight_dtype, + ) + # load attention processors + pipeline.load_lora_weights(args.output_dir) + + # run inference + images = [] + if args.validation_prompt and args.num_validation_images > 0: + images = log_validation( + pipeline=pipeline, + args=args, + accelerator=accelerator, + pipeline_args=validation_embeddings, + epoch=epoch, + is_final_validation=True, + torch_dtype=weight_dtype, + ) + images = None + del pipeline + free_memory() + + validation_prompt = args.validation_prompt if args.validation_prompt else args.final_validation_prompt + quant_training = None + if args.do_fp8_training: + quant_training = "FP8 TorchAO" + elif args.bnb_quantization_config_path: + quant_training = "BitsandBytes" + save_model_card( + (args.hub_model_id or Path(args.output_dir).name) if not args.push_to_hub else repo_id, + images=images, + base_model=args.pretrained_model_name_or_path, + instance_prompt=args.instance_prompt, + validation_prompt=validation_prompt, + repo_folder=args.output_dir, + quant_training=quant_training, + ) + + if args.push_to_hub: + upload_folder( + repo_id=repo_id, + folder_path=args.output_dir, + commit_message="End of training", + ignore_patterns=["step_*", "epoch_*"], + ) + + accelerator.end_training() + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/examples/dreambooth/train_dreambooth_sd3.py b/examples/dreambooth/train_dreambooth_sd3.py index e43e3178202a..98e7d2d66cbc 100644 --- a/examples/dreambooth/train_dreambooth_sd3.py +++ b/examples/dreambooth/train_dreambooth_sd3.py @@ -64,7 +64,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/flux-control/train_control_flux.py b/examples/flux-control/train_control_flux.py index 1e3be74464be..c5f93fa2e987 100644 --- a/examples/flux-control/train_control_flux.py +++ b/examples/flux-control/train_control_flux.py @@ -55,7 +55,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/flux-control/train_control_lora_flux.py b/examples/flux-control/train_control_lora_flux.py index 3185f1b2ea6a..f5d3c822b3ef 100644 --- a/examples/flux-control/train_control_lora_flux.py +++ b/examples/flux-control/train_control_lora_flux.py @@ -58,7 +58,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/instruct_pix2pix/train_instruct_pix2pix.py b/examples/instruct_pix2pix/train_instruct_pix2pix.py index 1bfe7aed30cb..55297b334cb9 100644 --- a/examples/instruct_pix2pix/train_instruct_pix2pix.py +++ b/examples/instruct_pix2pix/train_instruct_pix2pix.py @@ -58,7 +58,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__, log_level="INFO") diff --git a/examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py b/examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py index 9a5b23a8e623..5df0e22fe1cc 100644 --- a/examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py +++ b/examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py @@ -60,7 +60,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__, log_level="INFO") diff --git a/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py b/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py index 158c3a6f0994..9b6cb0523d67 100644 --- a/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py +++ b/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py @@ -53,7 +53,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__, log_level="INFO") diff --git a/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_decoder.py b/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_decoder.py index 30094f54827f..869b81ff5d33 100644 --- a/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_decoder.py +++ b/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_decoder.py @@ -46,7 +46,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__, log_level="INFO") diff --git a/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_prior.py b/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_prior.py index 9c0a4c38504e..8600269dd0fe 100644 --- a/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_prior.py +++ b/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_prior.py @@ -46,7 +46,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__, log_level="INFO") diff --git a/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py b/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py index caa8d96ef3ec..6cce862f95a5 100644 --- a/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py +++ b/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py @@ -52,7 +52,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__, log_level="INFO") diff --git a/examples/research_projects/anytext/anytext.py b/examples/research_projects/anytext/anytext.py index 7ae6ae57c22a..1e925e4fafaa 100644 --- a/examples/research_projects/anytext/anytext.py +++ b/examples/research_projects/anytext/anytext.py @@ -540,7 +540,7 @@ def __init__( max_length=77, freeze=True, use_fp16=False, - variant: Optional[str] = None, + variant: str | None = None, ): super().__init__() self.tokenizer = CLIPTokenizer.from_pretrained("tolgacangoz/anytext", subfolder="tokenizer") @@ -1018,7 +1018,7 @@ def insert_spaces(self, string, nSpace): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -1938,7 +1938,7 @@ def __call__( height: Optional[int] = None, width: Optional[int] = None, num_inference_steps: int = 50, - mode: Optional[str] = "generate", + mode: str | None = "generate", draw_pos: Optional[Union[str, torch.Tensor]] = None, ori_image: Optional[Union[str, torch.Tensor]] = None, timesteps: List[int] = None, @@ -1953,7 +1953,7 @@ def __call__( negative_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, controlnet_conditioning_scale: Union[float, List[float]] = 1.0, diff --git a/examples/research_projects/anytext/anytext_controlnet.py b/examples/research_projects/anytext/anytext_controlnet.py index 879d48fc8496..60a998dd27df 100644 --- a/examples/research_projects/anytext/anytext_controlnet.py +++ b/examples/research_projects/anytext/anytext_controlnet.py @@ -185,7 +185,7 @@ def __init__( "CrossAttnDownBlock2D", "DownBlock2D", ), - mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn", + mid_block_type: str | None = "UNetMidBlock2DCrossAttn", only_cross_attention: Union[bool, Tuple[bool]] = False, block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280), layers_per_block: int = 2, @@ -197,12 +197,12 @@ def __init__( cross_attention_dim: int = 1280, transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1, encoder_hid_dim: Optional[int] = None, - encoder_hid_dim_type: Optional[str] = None, + encoder_hid_dim_type: str | None = None, attention_head_dim: Union[int, Tuple[int, ...]] = 8, num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None, use_linear_projection: bool = False, - class_embed_type: Optional[str] = None, - addition_embed_type: Optional[str] = None, + class_embed_type: str | None = None, + addition_embed_type: str | None = None, addition_time_embed_dim: Optional[int] = None, num_class_embeds: Optional[int] = None, upcast_attention: bool = False, diff --git a/examples/research_projects/autoencoder_rae/README.md b/examples/research_projects/autoencoder_rae/README.md new file mode 100644 index 000000000000..9ade979090d9 --- /dev/null +++ b/examples/research_projects/autoencoder_rae/README.md @@ -0,0 +1,66 @@ +# Training AutoencoderRAE + +This example trains the decoder of `AutoencoderRAE` (stage-1 style), while keeping the representation encoder frozen. + +It follows the same high-level training recipe as the official RAE stage-1 setup: +- frozen encoder +- train decoder +- pixel reconstruction loss +- optional encoder feature consistency loss + +## Quickstart + +### Resume or finetune from pretrained weights + +```bash +accelerate launch examples/research_projects/autoencoder_rae/train_autoencoder_rae.py \ + --pretrained_model_name_or_path nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08 \ + --train_data_dir /path/to/imagenet_like_folder \ + --output_dir /tmp/autoencoder-rae \ + --resolution 256 \ + --train_batch_size 8 \ + --learning_rate 1e-4 \ + --num_train_epochs 10 \ + --report_to wandb \ + --reconstruction_loss_type l1 \ + --use_encoder_loss \ + --encoder_loss_weight 0.1 +``` + +### Train from scratch with a pretrained encoder +The following command launches RAE training with "facebook/dinov2-with-registers-base" as the base. + +```bash +accelerate launch examples/research_projects/autoencoder_rae/train_autoencoder_rae.py \ + --train_data_dir /path/to/imagenet_like_folder \ + --output_dir /tmp/autoencoder-rae \ + --resolution 256 \ + --encoder_type dinov2 \ + --encoder_name_or_path facebook/dinov2-with-registers-base \ + --encoder_input_size 224 \ + --patch_size 16 \ + --image_size 256 \ + --decoder_hidden_size 1152 \ + --decoder_num_hidden_layers 28 \ + --decoder_num_attention_heads 16 \ + --decoder_intermediate_size 4096 \ + --train_batch_size 8 \ + --learning_rate 1e-4 \ + --num_train_epochs 10 \ + --report_to wandb \ + --reconstruction_loss_type l1 \ + --use_encoder_loss \ + --encoder_loss_weight 0.1 +``` + +Note: stage-1 reconstruction loss assumes matching target/output spatial size, so `--resolution` must equal `--image_size`. + +Dataset format is expected to be `ImageFolder`-compatible: + +```text +train_data_dir/ + class_a/ + img_0001.jpg + class_b/ + img_0002.jpg +``` diff --git a/examples/research_projects/autoencoder_rae/train_autoencoder_rae.py b/examples/research_projects/autoencoder_rae/train_autoencoder_rae.py new file mode 100644 index 000000000000..ea02c674bc0c --- /dev/null +++ b/examples/research_projects/autoencoder_rae/train_autoencoder_rae.py @@ -0,0 +1,405 @@ +#!/usr/bin/env python +# coding=utf-8 +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import logging +import math +import os +from pathlib import Path + +import torch +import torch.nn.functional as F +from accelerate import Accelerator +from accelerate.logging import get_logger +from accelerate.utils import ProjectConfiguration, set_seed +from torch.utils.data import DataLoader +from torchvision import transforms +from torchvision.datasets import ImageFolder +from tqdm.auto import tqdm + +from diffusers import AutoencoderRAE +from diffusers.optimization import get_scheduler + + +logger = get_logger(__name__) + + +def parse_args(): + parser = argparse.ArgumentParser(description="Train a stage-1 Representation Autoencoder (RAE) decoder.") + parser.add_argument( + "--train_data_dir", + type=str, + required=True, + help="Path to an ImageFolder-style dataset root.", + ) + parser.add_argument( + "--output_dir", type=str, default="autoencoder-rae", help="Directory to save checkpoints/model." + ) + parser.add_argument("--logging_dir", type=str, default="logs", help="Accelerate logging directory.") + parser.add_argument("--seed", type=int, default=42) + + parser.add_argument("--resolution", type=int, default=256) + parser.add_argument("--center_crop", action="store_true") + parser.add_argument("--random_flip", action="store_true") + + parser.add_argument("--train_batch_size", type=int, default=8) + parser.add_argument("--dataloader_num_workers", type=int, default=4) + parser.add_argument("--num_train_epochs", type=int, default=10) + parser.add_argument("--max_train_steps", type=int, default=None) + parser.add_argument("--gradient_accumulation_steps", type=int, default=1) + parser.add_argument("--max_grad_norm", type=float, default=1.0) + + parser.add_argument("--learning_rate", type=float, default=1e-4) + parser.add_argument("--adam_beta1", type=float, default=0.9) + parser.add_argument("--adam_beta2", type=float, default=0.999) + parser.add_argument("--adam_weight_decay", type=float, default=1e-2) + parser.add_argument("--adam_epsilon", type=float, default=1e-8) + parser.add_argument("--lr_scheduler", type=str, default="cosine") + parser.add_argument("--lr_warmup_steps", type=int, default=500) + + parser.add_argument("--checkpointing_steps", type=int, default=1000) + parser.add_argument("--validation_steps", type=int, default=500) + + parser.add_argument( + "--pretrained_model_name_or_path", + type=str, + default=None, + help="Path to a pretrained AutoencoderRAE model (or HF Hub id) to resume training from.", + ) + parser.add_argument( + "--encoder_name_or_path", + type=str, + default=None, + help=( + "HF Hub id or local path of the pretrained encoder (e.g. 'facebook/dinov2-with-registers-base'). " + "When --pretrained_model_name_or_path is not set, the encoder weights are loaded from this path " + "into a freshly constructed AutoencoderRAE. Ignored when --pretrained_model_name_or_path is set." + ), + ) + + parser.add_argument("--encoder_type", type=str, choices=["dinov2", "siglip2", "mae"], default="dinov2") + parser.add_argument("--encoder_hidden_size", type=int, default=768) + parser.add_argument("--encoder_patch_size", type=int, default=14) + parser.add_argument("--encoder_num_hidden_layers", type=int, default=12) + parser.add_argument("--encoder_input_size", type=int, default=224) + parser.add_argument("--patch_size", type=int, default=16) + parser.add_argument("--image_size", type=int, default=256) + parser.add_argument("--num_channels", type=int, default=3) + + parser.add_argument("--decoder_hidden_size", type=int, default=1152) + parser.add_argument("--decoder_num_hidden_layers", type=int, default=28) + parser.add_argument("--decoder_num_attention_heads", type=int, default=16) + parser.add_argument("--decoder_intermediate_size", type=int, default=4096) + + parser.add_argument("--noise_tau", type=float, default=0.0) + parser.add_argument("--scaling_factor", type=float, default=1.0) + parser.add_argument("--reshape_to_2d", action=argparse.BooleanOptionalAction, default=True) + + parser.add_argument( + "--reconstruction_loss_type", + type=str, + choices=["l1", "mse"], + default="l1", + help="Pixel reconstruction loss.", + ) + parser.add_argument( + "--encoder_loss_weight", + type=float, + default=0.0, + help="Weight for encoder feature consistency loss in the training loop.", + ) + parser.add_argument( + "--use_encoder_loss", + action="store_true", + help="Enable encoder feature consistency loss term in the training loop.", + ) + parser.add_argument("--report_to", type=str, default="tensorboard") + + return parser.parse_args() + + +def build_transforms(args): + image_transforms = [ + transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BICUBIC), + transforms.CenterCrop(args.resolution) if args.center_crop else transforms.RandomCrop(args.resolution), + ] + if args.random_flip: + image_transforms.append(transforms.RandomHorizontalFlip()) + image_transforms.append(transforms.ToTensor()) + return transforms.Compose(image_transforms) + + +def compute_losses( + model, pixel_values, reconstruction_loss_type: str, use_encoder_loss: bool, encoder_loss_weight: float +): + decoded = model(pixel_values).sample + + if decoded.shape[-2:] != pixel_values.shape[-2:]: + raise ValueError( + "Training requires matching reconstruction and target sizes, got " + f"decoded={tuple(decoded.shape[-2:])}, target={tuple(pixel_values.shape[-2:])}." + ) + + if reconstruction_loss_type == "l1": + reconstruction_loss = F.l1_loss(decoded.float(), pixel_values.float()) + else: + reconstruction_loss = F.mse_loss(decoded.float(), pixel_values.float()) + + encoder_loss = torch.zeros_like(reconstruction_loss) + if use_encoder_loss and encoder_loss_weight > 0: + base_model = model.module if hasattr(model, "module") else model + target_encoder_input = base_model._resize_and_normalize(pixel_values) + reconstructed_encoder_input = base_model._resize_and_normalize(decoded) + + encoder_forward_kwargs = {"model": base_model.encoder} + if base_model.config.encoder_type == "mae": + encoder_forward_kwargs["patch_size"] = base_model.config.encoder_patch_size + with torch.no_grad(): + target_tokens = base_model._encoder_forward_fn(images=target_encoder_input, **encoder_forward_kwargs) + reconstructed_tokens = base_model._encoder_forward_fn( + images=reconstructed_encoder_input, **encoder_forward_kwargs + ) + encoder_loss = F.mse_loss(reconstructed_tokens.float(), target_tokens.float()) + + loss = reconstruction_loss + float(encoder_loss_weight) * encoder_loss + return decoded, loss, reconstruction_loss, encoder_loss + + +def _strip_final_layernorm_affine(state_dict, prefix=""): + """Remove final layernorm weight/bias so the model keeps its default init (identity).""" + keys_to_strip = {f"{prefix}weight", f"{prefix}bias"} + return {k: v for k, v in state_dict.items() if k not in keys_to_strip} + + +def _load_pretrained_encoder_weights(model, encoder_type, encoder_name_or_path): + """Load pretrained HF transformers encoder weights into the model's encoder.""" + if encoder_type == "dinov2": + from transformers import Dinov2WithRegistersModel + + hf_encoder = Dinov2WithRegistersModel.from_pretrained(encoder_name_or_path) + state_dict = hf_encoder.state_dict() + state_dict = _strip_final_layernorm_affine(state_dict, prefix="layernorm.") + elif encoder_type == "siglip2": + from transformers import SiglipModel + + hf_encoder = SiglipModel.from_pretrained(encoder_name_or_path).vision_model + state_dict = {f"vision_model.{k}": v for k, v in hf_encoder.state_dict().items()} + state_dict = _strip_final_layernorm_affine(state_dict, prefix="vision_model.post_layernorm.") + elif encoder_type == "mae": + from transformers import ViTMAEForPreTraining + + hf_encoder = ViTMAEForPreTraining.from_pretrained(encoder_name_or_path).vit + state_dict = hf_encoder.state_dict() + state_dict = _strip_final_layernorm_affine(state_dict, prefix="layernorm.") + else: + raise ValueError(f"Unknown encoder_type: {encoder_type}") + + model.encoder.load_state_dict(state_dict, strict=False) + + +def main(): + args = parse_args() + if args.resolution != args.image_size: + raise ValueError( + f"`--resolution` ({args.resolution}) must match `--image_size` ({args.image_size}) " + "for stage-1 reconstruction loss." + ) + + logging_dir = Path(args.output_dir, args.logging_dir) + accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) + accelerator = Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + project_config=accelerator_project_config, + log_with=args.report_to, + ) + + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + ) + logger.info(accelerator.state, main_process_only=False) + + if args.seed is not None: + set_seed(args.seed) + + if accelerator.is_main_process: + os.makedirs(args.output_dir, exist_ok=True) + accelerator.wait_for_everyone() + + dataset = ImageFolder(args.train_data_dir, transform=build_transforms(args)) + + def collate_fn(examples): + pixel_values = torch.stack([example[0] for example in examples]).float() + return {"pixel_values": pixel_values} + + train_dataloader = DataLoader( + dataset, + shuffle=True, + collate_fn=collate_fn, + batch_size=args.train_batch_size, + num_workers=args.dataloader_num_workers, + pin_memory=True, + drop_last=True, + ) + + if args.pretrained_model_name_or_path is not None: + model = AutoencoderRAE.from_pretrained(args.pretrained_model_name_or_path) + logger.info(f"Loaded pretrained AutoencoderRAE from {args.pretrained_model_name_or_path}") + else: + model = AutoencoderRAE( + encoder_type=args.encoder_type, + encoder_hidden_size=args.encoder_hidden_size, + encoder_patch_size=args.encoder_patch_size, + encoder_num_hidden_layers=args.encoder_num_hidden_layers, + decoder_hidden_size=args.decoder_hidden_size, + decoder_num_hidden_layers=args.decoder_num_hidden_layers, + decoder_num_attention_heads=args.decoder_num_attention_heads, + decoder_intermediate_size=args.decoder_intermediate_size, + patch_size=args.patch_size, + encoder_input_size=args.encoder_input_size, + image_size=args.image_size, + num_channels=args.num_channels, + noise_tau=args.noise_tau, + reshape_to_2d=args.reshape_to_2d, + use_encoder_loss=args.use_encoder_loss, + scaling_factor=args.scaling_factor, + ) + if args.encoder_name_or_path is not None: + _load_pretrained_encoder_weights(model, args.encoder_type, args.encoder_name_or_path) + logger.info(f"Loaded pretrained encoder weights from {args.encoder_name_or_path}") + model.encoder.requires_grad_(False) + model.decoder.requires_grad_(True) + model.train() + + optimizer = torch.optim.AdamW( + (p for p in model.parameters() if p.requires_grad), + lr=args.learning_rate, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + + overrode_max_train_steps = False + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if args.max_train_steps is None: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + overrode_max_train_steps = True + + lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=optimizer, + num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, + num_training_steps=args.max_train_steps * accelerator.num_processes, + ) + + model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( + model, optimizer, train_dataloader, lr_scheduler + ) + + if overrode_max_train_steps: + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + + if accelerator.is_main_process: + accelerator.init_trackers("train_autoencoder_rae", config=vars(args)) + + total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps + logger.info("***** Running training *****") + logger.info(f" Num examples = {len(dataset)}") + logger.info(f" Num Epochs = {args.num_train_epochs}") + logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") + logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") + logger.info(f" Total optimization steps = {args.max_train_steps}") + + progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process) + progress_bar.set_description("Steps") + global_step = 0 + + for epoch in range(args.num_train_epochs): + for step, batch in enumerate(train_dataloader): + with accelerator.accumulate(model): + pixel_values = batch["pixel_values"] + + _, loss, reconstruction_loss, encoder_loss = compute_losses( + model, + pixel_values, + reconstruction_loss_type=args.reconstruction_loss_type, + use_encoder_loss=args.use_encoder_loss, + encoder_loss_weight=args.encoder_loss_weight, + ) + + accelerator.backward(loss) + if accelerator.sync_gradients: + accelerator.clip_grad_norm_(model.parameters(), args.max_grad_norm) + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + if accelerator.sync_gradients: + progress_bar.update(1) + global_step += 1 + + logs = { + "loss": loss.detach().item(), + "reconstruction_loss": reconstruction_loss.detach().item(), + "encoder_loss": encoder_loss.detach().item(), + "lr": lr_scheduler.get_last_lr()[0], + } + progress_bar.set_postfix(**logs) + accelerator.log(logs, step=global_step) + + if global_step % args.validation_steps == 0: + with torch.no_grad(): + _, val_loss, val_reconstruction_loss, val_encoder_loss = compute_losses( + model, + pixel_values, + reconstruction_loss_type=args.reconstruction_loss_type, + use_encoder_loss=args.use_encoder_loss, + encoder_loss_weight=args.encoder_loss_weight, + ) + accelerator.log( + { + "val/loss": val_loss.detach().item(), + "val/reconstruction_loss": val_reconstruction_loss.detach().item(), + "val/encoder_loss": val_encoder_loss.detach().item(), + }, + step=global_step, + ) + + if global_step % args.checkpointing_steps == 0: + if accelerator.is_main_process: + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + unwrapped_model = accelerator.unwrap_model(model) + unwrapped_model.save_pretrained(save_path) + logger.info(f"Saved checkpoint to {save_path}") + + if global_step >= args.max_train_steps: + break + + if global_step >= args.max_train_steps: + break + + accelerator.wait_for_everyone() + if accelerator.is_main_process: + unwrapped_model = accelerator.unwrap_model(model) + unwrapped_model.save_pretrained(args.output_dir) + accelerator.end_training() + + +if __name__ == "__main__": + main() diff --git a/examples/research_projects/lpl/README.md b/examples/research_projects/lpl/README.md new file mode 100644 index 000000000000..a69fead50893 --- /dev/null +++ b/examples/research_projects/lpl/README.md @@ -0,0 +1,157 @@ +# Latent Perceptual Loss (LPL) for Stable Diffusion XL + +This directory contains an implementation of Latent Perceptual Loss (LPL) for training Stable Diffusion XL models, based on the paper: [Boosting Latent Diffusion with Perceptual Objectives](https://huggingface.co/papers/2411.04873) (Berrada et al., 2025). LPL is a perceptual loss that operates in the latent space of a VAE, helping to improve the quality and consistency of generated images by bridging the disconnect between the diffusion model and the autoencoder decoder. The implementation is based on the reference implementation provided by Tariq Berrada. + +## Overview + +LPL addresses a key limitation in latent diffusion models (LDMs): the disconnect between the diffusion model training and the autoencoder decoder. While LDMs train in the latent space, they don't receive direct feedback about how well their outputs decode into high-quality images. This can lead to: + +- Loss of fine details in generated images +- Inconsistent image quality +- Structural artifacts +- Reduced sharpness and realism + +LPL works by comparing intermediate features from the VAE decoder between the predicted and target latents. This helps the model learn better perceptual features and can lead to: + +- Improved image quality and consistency (6-20% FID improvement) +- Better preservation of fine details +- More stable training, especially at high noise levels +- Better handling of structural information +- Sharper and more realistic textures + +## Implementation Details + +The LPL implementation follows the paper's methodology and includes several key features: + +1. **Feature Extraction**: Extracts intermediate features from the VAE decoder, including: + - Middle block features + - Up block features (configurable number of blocks) + - Proper gradient checkpointing for memory efficiency + - Features are extracted only for timesteps below the threshold (high SNR) + +2. **Feature Normalization**: Multiple normalization options as validated in the paper: + - `default`: Normalize each feature map independently + - `shared`: Cross-normalize features using target statistics (recommended) + - `batch`: Batch-wise normalization + +3. **Outlier Handling**: Optional removal of outliers in feature maps using: + - Quantile-based filtering (2% quantiles) + - Morphological operations (opening/closing) + - Adaptive thresholding based on standard deviation + +4. **Loss Types**: + - MSE loss (default) + - L1 loss + - Optional power law weighting (2^(-i) for layer i) + +## Usage + +To use LPL in your training, add the following arguments to your training command: + +```bash +python examples/research_projects/lpl/train_sdxl_lpl.py \ + --use_lpl \ + --lpl_weight 1.0 \ # Weight for LPL loss (1.0-2.0 recommended) + --lpl_t_threshold 200 \ # Apply LPL only for timesteps < threshold (high SNR) + --lpl_loss_type mse \ # Loss type: "mse" or "l1" + --lpl_norm_type shared \ # Normalization type: "default", "shared" (recommended), or "batch" + --lpl_pow_law \ # Use power law weighting for layers + --lpl_num_blocks 4 \ # Number of up blocks to use (1-4) + --lpl_remove_outliers \ # Remove outliers in feature maps + --lpl_scale \ # Scale LPL loss by noise level weights + --lpl_start 0 \ # Step to start applying LPL + # ... other training arguments ... +``` + +### Key Parameters + +- `lpl_weight`: Controls the strength of the LPL loss relative to the main diffusion loss. Higher values (1.0-2.0) improve quality but may slow training. +- `lpl_t_threshold`: LPL is only applied for timesteps below this threshold (high SNR). Lower values (100-200) focus on more important timesteps. +- `lpl_loss_type`: Choose between MSE (default) and L1 loss. MSE is recommended for most cases. +- `lpl_norm_type`: Feature normalization strategy. "shared" is recommended as it showed best results in the paper. +- `lpl_pow_law`: Whether to use power law weighting (2^(-i) for layer i). Recommended for better feature balance. +- `lpl_num_blocks`: Number of up blocks to use for feature extraction (1-4). More blocks capture more features but use more memory. +- `lpl_remove_outliers`: Whether to remove outliers in feature maps. Recommended for stable training. +- `lpl_scale`: Whether to scale LPL loss by noise level weights. Helps focus on more important timesteps. +- `lpl_start`: Training step to start applying LPL. Can be used to warm up training. + +## Recommendations + +1. **Starting Point** (based on paper results): + ```bash + --use_lpl \ + --lpl_weight 1.0 \ + --lpl_t_threshold 200 \ + --lpl_loss_type mse \ + --lpl_norm_type shared \ + --lpl_pow_law \ + --lpl_num_blocks 4 \ + --lpl_remove_outliers \ + --lpl_scale + ``` + +2. **Memory Efficiency**: + - Use `--gradient_checkpointing` for memory efficiency (enabled by default) + - Reduce `lpl_num_blocks` if memory is constrained (2-3 blocks still give good results) + - Consider using `--lpl_scale` to focus on more important timesteps + - Features are extracted only for timesteps below threshold to save memory + +3. **Quality vs Speed**: + - Higher `lpl_weight` (1.0-2.0) for better quality + - Lower `lpl_t_threshold` (100-200) for faster training + - Use `lpl_remove_outliers` for more stable training + - `lpl_norm_type shared` provides best quality/speed trade-off + +## Technical Details + +### Feature Extraction + +The LPL implementation extracts features from the VAE decoder in the following order: +1. Middle block output +2. Up block outputs (configurable number of blocks) + +Each feature map is processed with: +1. Optional outlier removal (2% quantiles, morphological operations) +2. Feature normalization (shared statistics recommended) +3. Loss calculation (MSE or L1) +4. Optional power law weighting (2^(-i) for layer i) + +### Loss Calculation + +For each feature map: +1. Features are normalized according to the chosen strategy +2. Loss is calculated between normalized features +3. Outliers are masked out (if enabled) +4. Loss is weighted by layer depth (if power law enabled) +5. Final loss is averaged across all layers + +### Memory Considerations + +- Gradient checkpointing is used by default +- Features are extracted only for timesteps below the threshold +- Outlier removal is done in-place to save memory +- Feature normalization is done efficiently using vectorized operations +- Memory usage scales linearly with number of blocks used + +## Results + +Based on the paper's findings, LPL provides: +- 6-20% improvement in FID scores +- Better preservation of fine details +- More realistic textures and structures +- Improved consistency across different resolutions +- Better performance on both small and large datasets + +## Citation + +If you use this implementation in your research, please cite: + +```bibtex +@inproceedings{berrada2025boosting, + title={Boosting Latent Diffusion with Perceptual Objectives}, + author={Tariq Berrada and Pietro Astolfi and Melissa Hall and Marton Havasi and Yohann Benchetrit and Adriana Romero-Soriano and Karteek Alahari and Michal Drozdzal and Jakob Verbeek}, + booktitle={The Thirteenth International Conference on Learning Representations}, + year={2025}, + url={https://openreview.net/forum?id=y4DtzADzd1} +} +``` diff --git a/examples/research_projects/lpl/lpl_loss.py b/examples/research_projects/lpl/lpl_loss.py new file mode 100644 index 000000000000..de14a4d8d5aa --- /dev/null +++ b/examples/research_projects/lpl/lpl_loss.py @@ -0,0 +1,215 @@ +# Copyright 2025 Berrada et al. + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def normalize_tensor(in_feat, eps=1e-10): + norm_factor = torch.sqrt(torch.sum(in_feat**2, dim=1, keepdim=True)) + return in_feat / (norm_factor + eps) + + +def cross_normalize(input, target, eps=1e-10): + norm_factor = torch.sqrt(torch.sum(target**2, dim=1, keepdim=True)) + return input / (norm_factor + eps), target / (norm_factor + eps) + + +def remove_outliers(feat, down_f=1, opening=5, closing=3, m=100, quant=0.02): + opening = int(np.ceil(opening / down_f)) + closing = int(np.ceil(closing / down_f)) + if opening == 2: + opening = 3 + if closing == 2: + closing = 1 + + # replace quantile with kth value here. + feat_flat = feat.flatten(-2, -1) + k1, k2 = int(feat_flat.shape[-1] * quant), int(feat_flat.shape[-1] * (1 - quant)) + q1 = feat_flat.kthvalue(k1, dim=-1).values[..., None, None] + q2 = feat_flat.kthvalue(k2, dim=-1).values[..., None, None] + + m = 2 * feat_flat.std(-1)[..., None, None].detach() + mask = (q1 - m < feat) * (feat < q2 + m) + + # dilate the mask. + mask = nn.MaxPool2d(kernel_size=closing, stride=1, padding=(closing - 1) // 2)(mask.float()) # closing + mask = (-nn.MaxPool2d(kernel_size=opening, stride=1, padding=(opening - 1) // 2)(-mask)).bool() # opening + feat = feat * mask + return mask, feat + + +class LatentPerceptualLoss(nn.Module): + def __init__( + self, + vae, + loss_type="mse", + grad_ckpt=True, + pow_law=False, + norm_type="default", + num_mid_blocks=4, + feature_type="feature", + remove_outliers=True, + ): + super().__init__() + self.vae = vae + self.decoder = self.vae.decoder + # Store scaling factors as tensors on the correct device + device = next(self.vae.parameters()).device + + # Get scaling factors with proper defaults and handle None values + scale_factor = getattr(self.vae.config, "scaling_factor", None) + shift_factor = getattr(self.vae.config, "shift_factor", None) + + # Convert to tensors with proper defaults + self.scale = torch.tensor(1.0 if scale_factor is None else scale_factor, device=device) + self.shift = torch.tensor(0.0 if shift_factor is None else shift_factor, device=device) + + self.gradient_checkpointing = grad_ckpt + self.pow_law = pow_law + self.norm_type = norm_type.lower() + self.outlier_mask = remove_outliers + self.last_feature_stats = [] # Store feature statistics for logging + + assert feature_type in ["feature", "image"] + self.feature_type = feature_type + + assert self.norm_type in ["default", "shared", "batch"] + assert num_mid_blocks >= 0 and num_mid_blocks <= 4 + self.n_blocks = num_mid_blocks + + assert loss_type in ["mse", "l1"] + if loss_type == "mse": + self.loss_fn = nn.MSELoss(reduction="none") + elif loss_type == "l1": + self.loss_fn = nn.L1Loss(reduction="none") + + def get_features(self, z, latent_embeds=None, disable_grads=False): + with torch.set_grad_enabled(not disable_grads): + if self.gradient_checkpointing and not disable_grads: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + features = [] + upscale_dtype = next(iter(self.decoder.up_blocks.parameters())).dtype + sample = z + sample = self.decoder.conv_in(sample) + + # middle + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(self.decoder.mid_block), + sample, + latent_embeds, + use_reentrant=False, + ) + sample = sample.to(upscale_dtype) + features.append(sample) + + # up + for up_block in self.decoder.up_blocks[: self.n_blocks]: + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(up_block), + sample, + latent_embeds, + use_reentrant=False, + ) + features.append(sample) + return features + else: + features = [] + upscale_dtype = next(iter(self.decoder.up_blocks.parameters())).dtype + sample = z + sample = self.decoder.conv_in(sample) + + # middle + sample = self.decoder.mid_block(sample, latent_embeds) + sample = sample.to(upscale_dtype) + features.append(sample) + + # up + for up_block in self.decoder.up_blocks[: self.n_blocks]: + sample = up_block(sample, latent_embeds) + features.append(sample) + return features + + def get_loss(self, input, target, get_hist=False): + if self.feature_type == "feature": + inp_f = self.get_features(self.shift + input / self.scale) + tar_f = self.get_features(self.shift + target / self.scale, disable_grads=True) + losses = [] + self.last_feature_stats = [] # Reset feature stats + + for i, (x, y) in enumerate(zip(inp_f, tar_f, strict=False)): + my = torch.ones_like(y).bool() + outlier_ratio = 0.0 + + if self.outlier_mask: + with torch.no_grad(): + if i == 2: + my, y = remove_outliers(y, down_f=2) + outlier_ratio = 1.0 - my.float().mean().item() + elif i in [3, 4, 5]: + my, y = remove_outliers(y, down_f=1) + outlier_ratio = 1.0 - my.float().mean().item() + + # Store feature statistics before normalization + with torch.no_grad(): + stats = { + "mean": y.mean().item(), + "std": y.std().item(), + "outlier_ratio": outlier_ratio, + } + self.last_feature_stats.append(stats) + + # normalize feature tensors + if self.norm_type == "default": + x = normalize_tensor(x) + y = normalize_tensor(y) + elif self.norm_type == "shared": + x, y = cross_normalize(x, y, eps=1e-6) + + term_loss = self.loss_fn(x, y) * my + # reduce loss term + loss_f = 2 ** (-min(i, 3)) if self.pow_law else 1.0 + term_loss = term_loss.sum((2, 3)) * loss_f / my.sum((2, 3)) + losses.append(term_loss.mean((1,))) + + if get_hist: + return losses + else: + loss = sum(losses) + return loss / len(inp_f) + elif self.feature_type == "image": + inp_f = self.vae.decode(input / self.scale).sample + tar_f = self.vae.decode(target / self.scale).sample + return F.mse_loss(inp_f, tar_f) + + def get_first_conv(self, z): + sample = self.decoder.conv_in(z) + return sample + + def get_first_block(self, z): + sample = self.decoder.conv_in(z) + sample = self.decoder.mid_block(sample) + for resnet in self.decoder.up_blocks[0].resnets: + sample = resnet(sample, None) + return sample + + def get_first_layer(self, input, target, target_layer="conv"): + if target_layer == "conv": + feat_in = self.get_first_conv(input) + with torch.no_grad(): + feat_tar = self.get_first_conv(target) + else: + feat_in = self.get_first_block(input) + with torch.no_grad(): + feat_tar = self.get_first_block(target) + + feat_in, feat_tar = cross_normalize(feat_in, feat_tar) + + return F.mse_loss(feat_in, feat_tar, reduction="mean") diff --git a/examples/research_projects/lpl/train_sdxl_lpl.py b/examples/research_projects/lpl/train_sdxl_lpl.py new file mode 100644 index 000000000000..4c472c8871c0 --- /dev/null +++ b/examples/research_projects/lpl/train_sdxl_lpl.py @@ -0,0 +1,1622 @@ +#!/usr/bin/env python +# coding=utf-8 +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""LPL training script for Stable Diffusion XL for text2image.""" + +import argparse +import functools +import gc +import logging +import math +import os +import random +import re +import shutil +from contextlib import nullcontext +from pathlib import Path +from typing import Dict, List, Tuple + +import accelerate +import datasets +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint +import transformers +from accelerate import Accelerator +from accelerate.logging import get_logger +from accelerate.utils import DistributedType, ProjectConfiguration, set_seed +from datasets import concatenate_datasets, load_dataset +from huggingface_hub import create_repo, upload_folder +from lpl_loss import LatentPerceptualLoss +from packaging import version +from torchvision import transforms +from torchvision.transforms.functional import crop +from tqdm.auto import tqdm +from transformers import AutoTokenizer, PretrainedConfig + +import diffusers +from diffusers import AutoencoderKL, DDPMScheduler, StableDiffusionXLPipeline, UNet2DConditionModel +from diffusers.optimization import get_scheduler +from diffusers.training_utils import EMAModel, compute_snr +from diffusers.utils import check_min_version, is_wandb_available +from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card +from diffusers.utils.import_utils import is_torch_npu_available, is_xformers_available +from diffusers.utils.torch_utils import is_compiled_module + + +# Will error if the minimal version of diffusers is not installed. Remove at your own risks. +check_min_version("0.33.0.dev0") + +logger = get_logger(__name__) +if is_torch_npu_available(): + import torch_npu + + torch.npu.config.allow_internal_format = False + +DATASET_NAME_MAPPING = { + "lambdalabs/naruto-blip-captions": ("image", "text"), +} + +# Global dictionary to store intermediate features from hooks +hook_features: Dict[str, torch.Tensor] = {} + + +def get_intermediate_features_hook(name: str): + """Creates a hook function that saves the output of a layer.""" + + def hook(model, input, output): + # Some layers might return tuples (e.g., attention blocks) + # We are usually interested in the first element (hidden states) + if isinstance(output, tuple): + hook_features[name] = output[0] + else: + hook_features[name] = output + + return hook + + +def clear_hook_features(): + """Clears the global feature dictionary.""" + global hook_features + hook_features = {} + + +def normalize_features( + feat1: torch.Tensor, feat2: torch.Tensor, eps: float = 1e-6 +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Normalizes feat1 and feat2 using the statistics of feat2 (predicted features). + Normalization is done per-channel. + """ + # Calculate stats over spatial dimensions (H, W) + dims = tuple(range(2, feat2.ndim)) # Dims to reduce over (usually 2, 3 for H, W) + mean = torch.mean(feat2, dim=dims, keepdim=True) + std = torch.std(feat2, dim=dims, keepdim=True) + eps + + feat1_norm = (feat1 - mean) / std + feat2_norm = (feat2 - mean) / std + return feat1_norm, feat2_norm + + +def get_decoder_layer_names(decoder: nn.Module) -> List[str]: + """Helper to get potential layer names for hooks in the VAE decoder.""" + layer_names = [] + for name, module in decoder.named_modules(): + # Example: Target ResnetBlocks and potentially UpBlocks + if isinstance(module, (diffusers.models.resnet.ResnetBlock2D, diffusers.models.unet_2d_blocks.UpBlock2D)): + # Filter out redundant names if UpBlock contains ResnetBlocks already named + is_child = any( + name.startswith(parent + ".") + for parent in layer_names + if isinstance(decoder.get_submodule(parent), diffusers.models.unet_2d_blocks.UpBlock2D) + ) + if not is_child: + layer_names.append(name) + # A basic default selection if complex logic fails + if not layer_names: + layer_names = [ + name for name, module in decoder.named_modules() if re.match(r"up_blocks\.\d+\.resnets\.\d+$", name) + ] + return layer_names + + +def save_model_card( + repo_id: str, + images: list = None, + validation_prompt: str = None, + base_model: str = None, + dataset_name: str = None, + repo_folder: str = None, + vae_path: str = None, +): + img_str = "" + if images is not None: + for i, image in enumerate(images): + image.save(os.path.join(repo_folder, f"image_{i}.png")) + img_str += f"![img_{i}](./image_{i}.png)\n" + + model_description = f""" +# Text-to-image finetuning - {repo_id} + +This pipeline was finetuned from **{base_model}** on the **{dataset_name}** dataset. Below are some example images generated with the finetuned pipeline using the following prompt: {validation_prompt}: \n +{img_str} + +Special VAE used for training: {vae_path}. +""" + + model_card = load_or_create_model_card( + repo_id_or_path=repo_id, + from_training=True, + license="creativeml-openrail-m", + base_model=base_model, + model_description=model_description, + inference=True, + ) + + tags = [ + "stable-diffusion-xl", + "stable-diffusion-xl-diffusers", + "text-to-image", + "diffusers-training", + "diffusers", + ] + model_card = populate_model_card(model_card, tags=tags) + + model_card.save(os.path.join(repo_folder, "README.md")) + + +def import_model_class_from_model_name_or_path( + pretrained_model_name_or_path: str, revision: str, subfolder: str = "text_encoder" +): + text_encoder_config = PretrainedConfig.from_pretrained( + pretrained_model_name_or_path, subfolder=subfolder, revision=revision + ) + model_class = text_encoder_config.architectures[0] + + if model_class == "CLIPTextModel": + from transformers import CLIPTextModel + + return CLIPTextModel + elif model_class == "CLIPTextModelWithProjection": + from transformers import CLIPTextModelWithProjection + + return CLIPTextModelWithProjection + else: + raise ValueError(f"{model_class} is not supported.") + + +def parse_args(input_args=None): + parser = argparse.ArgumentParser(description="LPL based training script of Stable Diffusion XL.") + parser.add_argument( + "--pretrained_model_name_or_path", + type=str, + default=None, + required=True, + help="Path to pretrained model or model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--pretrained_vae_model_name_or_path", + type=str, + default=None, + help="Path to pretrained VAE model with better numerical stability. More details: https://github.com/huggingface/diffusers/pull/4038.", + ) + parser.add_argument( + "--revision", + type=str, + default=None, + required=False, + help="Revision of pretrained model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--variant", + type=str, + default=None, + help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16", + ) + parser.add_argument( + "--dataset_name", + type=str, + default=None, + help=( + "The name of the Dataset (from the HuggingFace hub) to train on (could be your own, possibly private," + " dataset). It can also be a path pointing to a local copy of a dataset in your filesystem," + " or to a folder containing files that 🤗 Datasets can understand." + ), + ) + parser.add_argument( + "--dataset_config_name", + type=str, + default=None, + help="The config of the Dataset, leave as None if there's only one config.", + ) + parser.add_argument( + "--train_data_dir", + type=str, + default=None, + help=( + "A folder containing the training data. Folder contents must follow the structure described in" + " https://huggingface.co/docs/datasets/image_dataset#imagefolder. In particular, a `metadata.jsonl` file" + " must exist to provide the captions for the images. Ignored if `dataset_name` is specified." + ), + ) + parser.add_argument( + "--image_column", type=str, default="image", help="The column of the dataset containing an image." + ) + parser.add_argument( + "--caption_column", + type=str, + default="text", + help="The column of the dataset containing a caption or a list of captions.", + ) + parser.add_argument( + "--validation_prompt", + type=str, + default=None, + help="A prompt that is used during validation to verify that the model is learning.", + ) + parser.add_argument( + "--num_validation_images", + type=int, + default=4, + help="Number of images that should be generated during validation with `validation_prompt`.", + ) + parser.add_argument( + "--validation_epochs", + type=int, + default=1, + help=( + "Run fine-tuning validation every X epochs. The validation process consists of running the prompt" + " `args.validation_prompt` multiple times: `args.num_validation_images`." + ), + ) + parser.add_argument( + "--max_train_samples", + type=int, + default=None, + help=( + "For debugging purposes or quicker training, truncate the number of training examples to this " + "value if set." + ), + ) + parser.add_argument( + "--proportion_empty_prompts", + type=float, + default=0, + help="Proportion of image prompts to be replaced with empty strings. Defaults to 0 (no prompt replacement).", + ) + parser.add_argument( + "--output_dir", + type=str, + default="sdxl-model-finetuned", + help="The output directory where the model predictions and checkpoints will be written.", + ) + parser.add_argument( + "--cache_dir", + type=str, + default=None, + help="The directory where the downloaded models and datasets will be stored.", + ) + parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") + parser.add_argument( + "--resolution", + type=int, + default=1024, + help=( + "The resolution for input images, all the images in the train/validation dataset will be resized to this" + " resolution" + ), + ) + parser.add_argument( + "--center_crop", + default=False, + action="store_true", + help=( + "Whether to center crop the input images to the resolution. If not set, the images will be randomly" + " cropped. The images will be resized to the resolution first before cropping." + ), + ) + parser.add_argument( + "--random_flip", + action="store_true", + help="whether to randomly flip images horizontally", + ) + parser.add_argument( + "--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader." + ) + parser.add_argument("--num_train_epochs", type=int, default=100) + parser.add_argument( + "--max_train_steps", + type=int, + default=None, + help="Total number of training steps to perform. If provided, overrides num_train_epochs.", + ) + parser.add_argument( + "--checkpointing_steps", + type=int, + default=500, + help=( + "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final" + " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming" + " training using `--resume_from_checkpoint`." + ), + ) + parser.add_argument( + "--checkpoints_total_limit", + type=int, + default=None, + help=("Max number of checkpoints to store."), + ) + parser.add_argument( + "--resume_from_checkpoint", + type=str, + default=None, + help=( + "Whether training should be resumed from a previous checkpoint. Use a path saved by" + ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' + ), + ) + parser.add_argument( + "--gradient_accumulation_steps", + type=int, + default=1, + help="Number of updates steps to accumulate before performing a backward/update pass.", + ) + parser.add_argument( + "--gradient_checkpointing", + action="store_true", + help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.", + ) + parser.add_argument( + "--learning_rate", + type=float, + default=1e-4, + help="Initial learning rate (after the potential warmup period) to use.", + ) + parser.add_argument( + "--scale_lr", + action="store_true", + default=False, + help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.", + ) + parser.add_argument( + "--lr_scheduler", + type=str, + default="constant", + help=( + 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' + ' "constant", "constant_with_warmup"]' + ), + ) + parser.add_argument( + "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler." + ) + parser.add_argument( + "--timestep_bias_strategy", + type=str, + default="none", + choices=["earlier", "later", "range", "none"], + help=( + "The timestep bias strategy, which may help direct the model toward learning low or high frequency details." + " Choices: ['earlier', 'later', 'range', 'none']." + " The default is 'none', which means no bias is applied, and training proceeds normally." + " The value of 'later' will increase the frequency of the model's final training timesteps." + ), + ) + parser.add_argument( + "--timestep_bias_multiplier", + type=float, + default=1.0, + help=( + "The multiplier for the bias. Defaults to 1.0, which means no bias is applied." + " A value of 2.0 will double the weight of the bias, and a value of 0.5 will halve it." + ), + ) + parser.add_argument( + "--timestep_bias_begin", + type=int, + default=0, + help=( + "When using `--timestep_bias_strategy=range`, the beginning (inclusive) timestep to bias." + " Defaults to zero, which equates to having no specific bias." + ), + ) + parser.add_argument( + "--timestep_bias_end", + type=int, + default=1000, + help=( + "When using `--timestep_bias_strategy=range`, the final timestep (inclusive) to bias." + " Defaults to 1000, which is the number of timesteps that Stable Diffusion is trained on." + ), + ) + parser.add_argument( + "--timestep_bias_portion", + type=float, + default=0.25, + help=( + "The portion of timesteps to bias. Defaults to 0.25, which 25% of timesteps will be biased." + " A value of 0.5 will bias one half of the timesteps. The value provided for `--timestep_bias_strategy` determines" + " whether the biased portions are in the earlier or later timesteps." + ), + ) + parser.add_argument( + "--snr_gamma", + type=float, + default=None, + help="SNR weighting gamma to be used if rebalancing the loss. Recommended value is 5.0. " + "More details here: https://arxiv.org/abs/2303.09556.", + ) + parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.") + parser.add_argument( + "--allow_tf32", + action="store_true", + help=( + "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" + " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" + ), + ) + parser.add_argument( + "--dataloader_num_workers", + type=int, + default=0, + help=( + "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process." + ), + ) + parser.add_argument( + "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes." + ) + parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.") + parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.") + parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.") + parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer") + parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") + parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.") + parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.") + parser.add_argument( + "--prediction_type", + type=str, + default=None, + help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediction_type` is chosen.", + ) + parser.add_argument( + "--hub_model_id", + type=str, + default=None, + help="The name of the repository to keep in sync with the local `output_dir`.", + ) + parser.add_argument( + "--logging_dir", + type=str, + default="logs", + help=( + "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to" + " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***." + ), + ) + parser.add_argument( + "--report_to", + type=str, + default="tensorboard", + help=( + 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' + ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' + ), + ) + parser.add_argument( + "--mixed_precision", + type=str, + default=None, + choices=["no", "fp16", "bf16"], + help=( + "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" + " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" + " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." + ), + ) + parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") + parser.add_argument( + "--enable_npu_flash_attention", action="store_true", help="Whether or not to use npu flash attention." + ) + parser.add_argument( + "--enable_xformers_memory_efficient_attention", action="store_true", help="Whether or not to use xformers." + ) + parser.add_argument("--noise_offset", type=float, default=0, help="The scale of noise offset.") + + parser.add_argument( + "--use_lpl", + action="store_true", + help="Whether to use Latent Perceptual Loss (LPL). Increases memory usage.", + ) + parser.add_argument( + "--lpl_weight", + type=float, + default=1.0, + help="Weight for the Latent Perceptual Loss.", + ) + parser.add_argument( + "--lpl_t_threshold", + type=int, + default=200, + help="Apply LPL only for timesteps t < lpl_t_threshold. Corresponds to high SNR.", + ) + parser.add_argument( + "--lpl_loss_type", + type=str, + default="mse", + choices=["mse", "l1"], + help="Type of loss to use for LPL.", + ) + parser.add_argument( + "--lpl_norm_type", + type=str, + default="default", + choices=["default", "shared", "batch"], + help="Type of normalization to use for LPL features.", + ) + parser.add_argument( + "--lpl_pow_law", + action="store_true", + help="Whether to use power law weighting for LPL layers.", + ) + parser.add_argument( + "--lpl_num_blocks", + type=int, + default=4, + help="Number of up blocks to use for LPL feature extraction.", + ) + parser.add_argument( + "--lpl_remove_outliers", + action="store_true", + help="Whether to remove outliers in LPL feature maps.", + ) + parser.add_argument( + "--lpl_scale", + action="store_true", + help="Whether to scale LPL loss by noise level weights.", + ) + parser.add_argument( + "--lpl_start", + type=int, + default=0, + help="Step to start applying LPL loss.", + ) + + if input_args is not None: + args = parser.parse_args(input_args) + else: + args = parser.parse_args() + + env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) + if env_local_rank != -1 and env_local_rank != args.local_rank: + args.local_rank = env_local_rank + + # Sanity checks + if args.dataset_name is None and args.train_data_dir is None: + raise ValueError("Need either a dataset name or a training folder.") + if args.proportion_empty_prompts < 0 or args.proportion_empty_prompts > 1: + raise ValueError("`--proportion_empty_prompts` must be in the range [0, 1].") + + return args + + +# Adapted from pipelines.StableDiffusionXLPipeline.encode_prompt +def encode_prompt(batch, text_encoders, tokenizers, proportion_empty_prompts, caption_column, is_train=True): + prompt_embeds_list = [] + prompt_batch = batch[caption_column] + + captions = [] + for caption in prompt_batch: + if random.random() < proportion_empty_prompts: + captions.append("") + elif isinstance(caption, str): + captions.append(caption) + elif isinstance(caption, (list, np.ndarray)): + # take a random caption if there are multiple + captions.append(random.choice(caption) if is_train else caption[0]) + + with torch.no_grad(): + for tokenizer, text_encoder in zip(tokenizers, text_encoders): + text_inputs = tokenizer( + captions, + padding="max_length", + max_length=tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + prompt_embeds = text_encoder( + text_input_ids.to(text_encoder.device), + output_hidden_states=True, + return_dict=False, + ) + + # We are only ALWAYS interested in the pooled output of the final text encoder + pooled_prompt_embeds = prompt_embeds[0] + prompt_embeds = prompt_embeds[-1][-2] + bs_embed, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.view(bs_embed, seq_len, -1) + prompt_embeds_list.append(prompt_embeds) + + prompt_embeds = torch.concat(prompt_embeds_list, dim=-1) + pooled_prompt_embeds = pooled_prompt_embeds.view(bs_embed, -1) + return {"prompt_embeds": prompt_embeds.cpu(), "pooled_prompt_embeds": pooled_prompt_embeds.cpu()} + + +def compute_vae_encodings(batch, vae): + images = batch.pop("pixel_values") + pixel_values = torch.stack(list(images)) + pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float() + pixel_values = pixel_values.to(vae.device, dtype=vae.dtype) + + with torch.no_grad(): + model_input = vae.encode(pixel_values).latent_dist.sample() + model_input = model_input * vae.config.scaling_factor + + # There might have slightly performance improvement + # by changing model_input.cpu() to accelerator.gather(model_input) + return {"model_input": model_input.cpu()} + + +def generate_timestep_weights(args, num_timesteps): + weights = torch.ones(num_timesteps) + + # Determine the indices to bias + num_to_bias = int(args.timestep_bias_portion * num_timesteps) + + if args.timestep_bias_strategy == "later": + bias_indices = slice(-num_to_bias, None) + elif args.timestep_bias_strategy == "earlier": + bias_indices = slice(0, num_to_bias) + elif args.timestep_bias_strategy == "range": + # Out of the possible 1000 timesteps, we might want to focus on eg. 200-500. + range_begin = args.timestep_bias_begin + range_end = args.timestep_bias_end + if range_begin < 0: + raise ValueError( + "When using the range strategy for timestep bias, you must provide a beginning timestep greater or equal to zero." + ) + if range_end > num_timesteps: + raise ValueError( + "When using the range strategy for timestep bias, you must provide an ending timestep smaller than the number of timesteps." + ) + bias_indices = slice(range_begin, range_end) + else: # 'none' or any other string + return weights + if args.timestep_bias_multiplier <= 0: + return ValueError( + "The parameter --timestep_bias_multiplier is not intended to be used to disable the training of specific timesteps." + " If it was intended to disable timestep bias, use `--timestep_bias_strategy none` instead." + " A timestep bias multiplier less than or equal to 0 is not allowed." + ) + + # Apply the bias + weights[bias_indices] *= args.timestep_bias_multiplier + + # Normalize + weights /= weights.sum() + + return weights + + +def main(args): + if args.report_to == "wandb" and args.hub_token is not None: + raise ValueError( + "You cannot use both --report_to=wandb and --hub_token due to a security risk of exposing your token." + " Please use `huggingface-cli login` to authenticate with the Hub." + ) + + logging_dir = Path(args.output_dir, args.logging_dir) + + accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) + + if torch.backends.mps.is_available() and args.mixed_precision == "bf16": + # due to pytorch#99272, MPS does not yet support bfloat16. + raise ValueError( + "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead." + ) + + accelerator = Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + mixed_precision=args.mixed_precision, + log_with=args.report_to, + project_config=accelerator_project_config, + ) + + # Disable AMP for MPS. + if torch.backends.mps.is_available(): + accelerator.native_amp = False + + if args.report_to == "wandb": + if not is_wandb_available(): + raise ImportError("Make sure to install wandb if you want to use it for logging during training.") + import wandb + + # Make one log on every process with the configuration for debugging. + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + ) + logger.info(accelerator.state, main_process_only=False) + if accelerator.is_local_main_process: + datasets.utils.logging.set_verbosity_warning() + transformers.utils.logging.set_verbosity_warning() + diffusers.utils.logging.set_verbosity_info() + else: + datasets.utils.logging.set_verbosity_error() + transformers.utils.logging.set_verbosity_error() + diffusers.utils.logging.set_verbosity_error() + + # If passed along, set the training seed now. + if args.seed is not None: + set_seed(args.seed) + + # Handle the repository creation + if accelerator.is_main_process: + if args.output_dir is not None: + os.makedirs(args.output_dir, exist_ok=True) + + if args.push_to_hub: + repo_id = create_repo( + repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, token=args.hub_token + ).repo_id + + # Load the tokenizers + tokenizer_one = AutoTokenizer.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="tokenizer", + revision=args.revision, + use_fast=False, + ) + tokenizer_two = AutoTokenizer.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="tokenizer_2", + revision=args.revision, + use_fast=False, + ) + + # import correct text encoder classes + text_encoder_cls_one = import_model_class_from_model_name_or_path( + args.pretrained_model_name_or_path, args.revision + ) + text_encoder_cls_two = import_model_class_from_model_name_or_path( + args.pretrained_model_name_or_path, args.revision, subfolder="text_encoder_2" + ) + + # Load scheduler and models + noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler") + # Check for terminal SNR in combination with SNR Gamma + text_encoder_one = text_encoder_cls_one.from_pretrained( + args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision, variant=args.variant + ) + text_encoder_two = text_encoder_cls_two.from_pretrained( + args.pretrained_model_name_or_path, subfolder="text_encoder_2", revision=args.revision, variant=args.variant + ) + vae_path = ( + args.pretrained_model_name_or_path + if args.pretrained_vae_model_name_or_path is None + else args.pretrained_vae_model_name_or_path + ) + vae = AutoencoderKL.from_pretrained( + vae_path, + subfolder="vae" if args.pretrained_vae_model_name_or_path is None else None, + revision=args.revision, + variant=args.variant, + ) + unet = UNet2DConditionModel.from_pretrained( + args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision, variant=args.variant + ) + + # Freeze vae and text encoders. + vae.requires_grad_(False) + text_encoder_one.requires_grad_(False) + text_encoder_two.requires_grad_(False) + # Set unet as trainable. + unet.train() + + # For mixed precision training we cast all non-trainable weights to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + weight_dtype = torch.float32 + if accelerator.mixed_precision == "fp16": + weight_dtype = torch.float16 + elif accelerator.mixed_precision == "bf16": + weight_dtype = torch.bfloat16 + + # Move unet, vae and text_encoder to device and cast to weight_dtype + # The VAE is in float32 to avoid NaN losses. + vae.to(accelerator.device, dtype=torch.float32) + text_encoder_one.to(accelerator.device, dtype=weight_dtype) + text_encoder_two.to(accelerator.device, dtype=weight_dtype) + + # Create EMA for the unet. + if args.use_ema: + ema_unet = UNet2DConditionModel.from_pretrained( + args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision, variant=args.variant + ) + ema_unet = EMAModel(ema_unet.parameters(), model_cls=UNet2DConditionModel, model_config=ema_unet.config) + if args.enable_npu_flash_attention: + if is_torch_npu_available(): + logger.info("npu flash attention enabled.") + unet.enable_npu_flash_attention() + else: + raise ValueError("npu flash attention requires torch_npu extensions and is supported only on npu devices.") + if args.enable_xformers_memory_efficient_attention: + if is_xformers_available(): + import xformers + + xformers_version = version.parse(xformers.__version__) + if xformers_version == version.parse("0.0.16"): + logger.warning( + "xFormers 0.0.16 cannot be used for training in some GPUs. If you observe problems during training, please update xFormers to at least 0.0.17. See https://huggingface.co/docs/diffusers/main/en/optimization/xformers for more details." + ) + unet.enable_xformers_memory_efficient_attention() + else: + raise ValueError("xformers is not available. Make sure it is installed correctly") + + # `accelerate` 0.16.0 will have better support for customized saving + if version.parse(accelerate.__version__) >= version.parse("0.16.0"): + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + def save_model_hook(models, weights, output_dir): + if accelerator.is_main_process: + if args.use_ema: + ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema")) + + for i, model in enumerate(models): + model.save_pretrained(os.path.join(output_dir, "unet")) + + # make sure to pop weight so that corresponding model is not saved again + if weights: + weights.pop() + + def load_model_hook(models, input_dir): + if args.use_ema: + load_model = EMAModel.from_pretrained(os.path.join(input_dir, "unet_ema"), UNet2DConditionModel) + ema_unet.load_state_dict(load_model.state_dict()) + ema_unet.to(accelerator.device) + del load_model + + for _ in range(len(models)): + # pop models so that they are not loaded again + model = models.pop() + + # load diffusers style into model + load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet") + model.register_to_config(**load_model.config) + + model.load_state_dict(load_model.state_dict()) + del load_model + + accelerator.register_save_state_pre_hook(save_model_hook) + accelerator.register_load_state_pre_hook(load_model_hook) + + if args.gradient_checkpointing: + unet.enable_gradient_checkpointing() + + # Enable TF32 for faster training on Ampere GPUs, + # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices + if args.allow_tf32: + torch.backends.cuda.matmul.allow_tf32 = True + + if args.scale_lr: + args.learning_rate = ( + args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes + ) + + # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs + if args.use_8bit_adam: + try: + import bitsandbytes as bnb + except ImportError: + raise ImportError( + "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`." + ) + + optimizer_class = bnb.optim.AdamW8bit + else: + optimizer_class = torch.optim.AdamW + + # Optimizer creation + params_to_optimize = unet.parameters() + optimizer = optimizer_class( + params_to_optimize, + lr=args.learning_rate, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + + # Get the datasets: you can either provide your own training and evaluation files (see below) + # or specify a Dataset from the hub (the dataset will be downloaded automatically from the datasets Hub). + + # In distributed training, the load_dataset function guarantees that only one local process can concurrently + # download the dataset. + if args.dataset_name is not None: + # Downloading and loading a dataset from the hub. + dataset = load_dataset( + args.dataset_name, args.dataset_config_name, cache_dir=args.cache_dir, data_dir=args.train_data_dir + ) + else: + data_files = {} + if args.train_data_dir is not None: + data_files["train"] = os.path.join(args.train_data_dir, "**") + dataset = load_dataset( + "imagefolder", + data_files=data_files, + cache_dir=args.cache_dir, + ) + # See more about loading custom images at + # https://huggingface.co/docs/datasets/v2.4.0/en/image_load#imagefolder + + # Preprocessing the datasets. + # We need to tokenize inputs and targets. + column_names = dataset["train"].column_names + + # 6. Get the column names for input/target. + dataset_columns = DATASET_NAME_MAPPING.get(args.dataset_name, None) + if args.image_column is None: + image_column = dataset_columns[0] if dataset_columns is not None else column_names[0] + else: + image_column = args.image_column + if image_column not in column_names: + raise ValueError( + f"--image_column' value '{args.image_column}' needs to be one of: {', '.join(column_names)}" + ) + if args.caption_column is None: + caption_column = dataset_columns[1] if dataset_columns is not None else column_names[1] + else: + caption_column = args.caption_column + if caption_column not in column_names: + raise ValueError( + f"--caption_column' value '{args.caption_column}' needs to be one of: {', '.join(column_names)}" + ) + + # Preprocessing the datasets. + train_resize = transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR) + train_crop = transforms.CenterCrop(args.resolution) if args.center_crop else transforms.RandomCrop(args.resolution) + train_flip = transforms.RandomHorizontalFlip(p=1.0) + train_transforms = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]) + + def preprocess_train(examples): + images = [image.convert("RGB") for image in examples[image_column]] + # image aug + original_sizes = [] + all_images = [] + crop_top_lefts = [] + for image in images: + original_sizes.append((image.height, image.width)) + image = train_resize(image) + if args.random_flip and random.random() < 0.5: + # flip + image = train_flip(image) + if args.center_crop: + y1 = max(0, int(round((image.height - args.resolution) / 2.0))) + x1 = max(0, int(round((image.width - args.resolution) / 2.0))) + image = train_crop(image) + else: + y1, x1, h, w = train_crop.get_params(image, (args.resolution, args.resolution)) + image = crop(image, y1, x1, h, w) + crop_top_left = (y1, x1) + crop_top_lefts.append(crop_top_left) + image = train_transforms(image) + all_images.append(image) + + examples["original_sizes"] = original_sizes + examples["crop_top_lefts"] = crop_top_lefts + examples["pixel_values"] = all_images + return examples + + with accelerator.main_process_first(): + if args.max_train_samples is not None: + dataset["train"] = dataset["train"].shuffle(seed=args.seed).select(range(args.max_train_samples)) + # Set the training transforms + train_dataset = dataset["train"].with_transform(preprocess_train) + + # Let's first compute all the embeddings so that we can free up the text encoders + # from memory. We will pre-compute the VAE encodings too. + text_encoders = [text_encoder_one, text_encoder_two] + tokenizers = [tokenizer_one, tokenizer_two] + compute_embeddings_fn = functools.partial( + encode_prompt, + text_encoders=text_encoders, + tokenizers=tokenizers, + proportion_empty_prompts=args.proportion_empty_prompts, + caption_column=args.caption_column, + ) + compute_vae_encodings_fn = functools.partial(compute_vae_encodings, vae=vae) + with accelerator.main_process_first(): + from datasets.fingerprint import Hasher + + # fingerprint used by the cache for the other processes to load the result + # details: https://github.com/huggingface/diffusers/pull/4038#discussion_r1266078401 + new_fingerprint = Hasher.hash(args) + new_fingerprint_for_vae = Hasher.hash((vae_path, args)) + train_dataset_with_embeddings = train_dataset.map( + compute_embeddings_fn, batched=True, new_fingerprint=new_fingerprint + ) + train_dataset_with_vae = train_dataset.map( + compute_vae_encodings_fn, + batched=True, + batch_size=args.train_batch_size, + new_fingerprint=new_fingerprint_for_vae, + ) + precomputed_dataset = concatenate_datasets( + [train_dataset_with_embeddings, train_dataset_with_vae.remove_columns(["image", "text"])], axis=1 + ) + precomputed_dataset = precomputed_dataset.with_transform(preprocess_train) + + del compute_vae_encodings_fn, compute_embeddings_fn, text_encoder_one, text_encoder_two + del text_encoders, tokenizers + if not args.use_lpl: + del vae + gc.collect() + + if is_torch_npu_available(): + torch_npu.npu.empty_cache() + elif torch.cuda.is_available(): + torch.cuda.empty_cache() + + def collate_fn(examples): + model_input = torch.stack([torch.tensor(example["model_input"]) for example in examples]) + original_sizes = [example["original_sizes"] for example in examples] + crop_top_lefts = [example["crop_top_lefts"] for example in examples] + prompt_embeds = torch.stack([torch.tensor(example["prompt_embeds"]) for example in examples]) + pooled_prompt_embeds = torch.stack([torch.tensor(example["pooled_prompt_embeds"]) for example in examples]) + + return { + "model_input": model_input, + "prompt_embeds": prompt_embeds, + "pooled_prompt_embeds": pooled_prompt_embeds, + "original_sizes": original_sizes, + "crop_top_lefts": crop_top_lefts, + } + + # DataLoaders creation: + train_dataloader = torch.utils.data.DataLoader( + precomputed_dataset, + shuffle=True, + collate_fn=collate_fn, + batch_size=args.train_batch_size, + num_workers=args.dataloader_num_workers, + ) + + # Scheduler and math around the number of training steps. + overrode_max_train_steps = False + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if args.max_train_steps is None: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + overrode_max_train_steps = True + + lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=optimizer, + num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps, + num_training_steps=args.max_train_steps * args.gradient_accumulation_steps, + ) + + # Prepare everything with our `accelerator`. + unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( + unet, optimizer, train_dataloader, lr_scheduler + ) + + if args.use_ema: + ema_unet.to(accelerator.device) + + # We need to recalculate our total training steps as the size of the training dataloader may have changed. + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if overrode_max_train_steps: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + # Afterwards we recalculate our number of training epochs + args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + + # We need to initialize the trackers we use, and also store our configuration. + # The trackers initializes automatically on the main process. + if accelerator.is_main_process: + accelerator.init_trackers("text2image-fine-tune-sdxl", config=vars(args)) + + if args.use_lpl: + lpl_fn = LatentPerceptualLoss( + vae=vae, + loss_type=args.lpl_loss_type, + grad_ckpt=args.gradient_checkpointing, + pow_law=args.lpl_pow_law, + norm_type=args.lpl_norm_type, + num_mid_blocks=args.lpl_num_blocks, + feature_type="feature", + remove_outliers=args.lpl_remove_outliers, + ) + lpl_fn.to(accelerator.device) + else: + lpl_fn = None + + # Function for unwrapping if torch.compile() was used in accelerate. + def unwrap_model(model): + model = accelerator.unwrap_model(model) + model = model._orig_mod if is_compiled_module(model) else model + return model + + if torch.backends.mps.is_available() or "playground" in args.pretrained_model_name_or_path: + autocast_ctx = nullcontext() + else: + autocast_ctx = torch.autocast(accelerator.device.type) + + # Train! + total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps + + logger.info("***** Running training *****") + logger.info(f" Num examples = {len(precomputed_dataset)}") + logger.info(f" Num Epochs = {args.num_train_epochs}") + logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") + logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") + logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") + logger.info(f" Total optimization steps = {args.max_train_steps}") + global_step = 0 + first_epoch = 0 + + # Potentially load in the weights and states from a previous save + if args.resume_from_checkpoint: + if args.resume_from_checkpoint != "latest": + path = os.path.basename(args.resume_from_checkpoint) + else: + # Get the most recent checkpoint + dirs = os.listdir(args.output_dir) + dirs = [d for d in dirs if d.startswith("checkpoint")] + dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) + path = dirs[-1] if len(dirs) > 0 else None + + if path is None: + accelerator.print( + f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run." + ) + args.resume_from_checkpoint = None + initial_global_step = 0 + else: + accelerator.print(f"Resuming from checkpoint {path}") + accelerator.load_state(os.path.join(args.output_dir, path)) + global_step = int(path.split("-")[1]) + + initial_global_step = global_step + first_epoch = global_step // num_update_steps_per_epoch + + else: + initial_global_step = 0 + + progress_bar = tqdm( + range(0, args.max_train_steps), + initial=initial_global_step, + desc="Steps", + # Only show the progress bar once on each machine. + disable=not accelerator.is_local_main_process, + ) + + # Get scheduler alphas and sigmas for LPL z0_hat calculation + alphas_cumprod = noise_scheduler.alphas_cumprod.to(accelerator.device) + + for epoch in range(first_epoch, args.num_train_epochs): + train_loss = 0.0 + for step, batch in enumerate(train_dataloader): + with accelerator.accumulate(unet): + # Sample noise that we'll add to the latents + model_input = batch["model_input"].to(accelerator.device) + noise = torch.randn_like(model_input) + if args.noise_offset: + # https://www.crosslabs.org//blog/diffusion-with-offset-noise + noise += args.noise_offset * torch.randn( + (model_input.shape[0], model_input.shape[1], 1, 1), device=model_input.device + ) + + bsz = model_input.shape[0] + if args.timestep_bias_strategy == "none": + # Sample a random timestep for each image without bias. + timesteps = torch.randint( + 0, noise_scheduler.config.num_train_timesteps, (bsz,), device=model_input.device + ) + else: + # Sample a random timestep for each image, potentially biased by the timestep weights. + # Biasing the timestep weights allows us to spend less time training irrelevant timesteps. + weights = generate_timestep_weights(args, noise_scheduler.config.num_train_timesteps).to( + model_input.device + ) + timesteps = torch.multinomial(weights, bsz, replacement=True).long() + + # Add noise to the model input according to the noise magnitude at each timestep + # (this is the forward diffusion process) + noisy_model_input = noise_scheduler.add_noise(model_input, noise, timesteps).to(dtype=weight_dtype) + + # time ids + def compute_time_ids(original_size, crops_coords_top_left): + # Adapted from pipeline.StableDiffusionXLPipeline._get_add_time_ids + target_size = (args.resolution, args.resolution) + add_time_ids = list(original_size + crops_coords_top_left + target_size) + add_time_ids = torch.tensor([add_time_ids], device=accelerator.device, dtype=weight_dtype) + return add_time_ids + + add_time_ids = torch.cat( + [compute_time_ids(s, c) for s, c in zip(batch["original_sizes"], batch["crop_top_lefts"])] + ) + + # Predict the noise residual + unet_added_conditions = {"time_ids": add_time_ids} + prompt_embeds = batch["prompt_embeds"].to(accelerator.device, dtype=weight_dtype) + pooled_prompt_embeds = batch["pooled_prompt_embeds"].to(accelerator.device) + unet_added_conditions.update({"text_embeds": pooled_prompt_embeds}) + model_pred = unet( + noisy_model_input, + timesteps, + prompt_embeds, + added_cond_kwargs=unet_added_conditions, + return_dict=False, + )[0] + + # Get the target for loss depending on the prediction type + if args.prediction_type is not None: + # set prediction_type of scheduler if defined + noise_scheduler.register_to_config(prediction_type=args.prediction_type) + + if noise_scheduler.config.prediction_type == "epsilon": + target = noise + elif noise_scheduler.config.prediction_type == "v_prediction": + target = noise_scheduler.get_velocity(model_input, noise, timesteps) + elif noise_scheduler.config.prediction_type == "sample": + # We set the target to latents here, but the model_pred will return the noise sample prediction. + target = model_input + # We will have to subtract the noise residual from the prediction to get the target sample. + model_pred = model_pred - noise + else: + raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}") + + if args.snr_gamma is None: + loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean") + else: + # Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556. + # Since we predict the noise instead of x_0, the original formulation is slightly changed. + # This is discussed in Section 4.2 of the same paper. + snr = compute_snr(noise_scheduler, timesteps) + mse_loss_weights = torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min( + dim=1 + )[0] + if noise_scheduler.config.prediction_type == "epsilon": + mse_loss_weights = mse_loss_weights / snr + elif noise_scheduler.config.prediction_type == "v_prediction": + mse_loss_weights = mse_loss_weights / (snr + 1) + + loss = F.mse_loss(model_pred.float(), target.float(), reduction="none") + loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights + loss = loss.mean() + + lpl_loss_value = torch.tensor(0.0, device=accelerator.device) + if args.use_lpl and lpl_fn is not None and global_step >= args.lpl_start: + # Apply LPL only below the timestep threshold + lpl_mask = timesteps < args.lpl_t_threshold + if lpl_mask.any(): + # Select samples that meet the threshold + masked_indices = torch.where(lpl_mask)[0] + z0_masked = model_input[masked_indices] + zt_masked = noisy_model_input[masked_indices] + t_masked = timesteps[masked_indices] + model_pred_masked = model_pred[masked_indices] + + # Calculate z0_hat for the masked samples + alpha_t = alphas_cumprod[t_masked].sqrt().to(torch.float32) + sigma_t = (1 - alphas_cumprod[t_masked]).sqrt().to(torch.float32) + alpha_t = alpha_t.view(-1, 1, 1, 1) + sigma_t = sigma_t.view(-1, 1, 1, 1) + + if noise_scheduler.config.prediction_type == "epsilon": + z0_hat_masked = (zt_masked.float() - sigma_t * model_pred_masked.float()) / alpha_t + elif noise_scheduler.config.prediction_type == "v_prediction": + z0_hat_masked = alpha_t * zt_masked.float() - sigma_t * model_pred_masked.float() + else: # sample prediction + z0_hat_masked = model_pred_masked.float() + + with accelerator.autocast(): + lpl_loss_value = lpl_fn.get_loss(z0_hat_masked, z0_masked) + + if args.lpl_scale: + if args.snr_gamma is not None: + # Use SNR-based weights if available + snr = compute_snr(noise_scheduler, t_masked) + snr_weights = torch.stack( + [snr, args.snr_gamma * torch.ones_like(t_masked)], dim=1 + ).min(dim=1)[0] + if noise_scheduler.config.prediction_type == "epsilon": + snr_weights = snr_weights / snr + elif noise_scheduler.config.prediction_type == "v_prediction": + snr_weights = snr_weights / (snr + 1) + lpl_loss_value = (lpl_loss_value * snr_weights).mean() + else: + # If no SNR weighting, just use mean + lpl_loss_value = lpl_loss_value.mean() + else: + lpl_loss_value = lpl_loss_value.mean() + + # Combine losses + total_loss = loss + args.lpl_weight * lpl_loss_value + + # Gather the losses across all processes for logging + avg_loss = accelerator.gather(total_loss.repeat(args.train_batch_size)).mean() + train_loss += avg_loss.item() / args.gradient_accumulation_steps + + # Backpropagate + accelerator.backward(total_loss) + if accelerator.sync_gradients: + params_to_clip = unet.parameters() + accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm) + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + # Checks if the accelerator has performed an optimization step behind the scenes + if accelerator.sync_gradients: + if args.use_ema: + ema_unet.step(unet.parameters()) + progress_bar.update(1) + global_step += 1 + + # Enhanced logging for LPL metrics + log_data = { + "train_loss": train_loss, + "diffusion_loss": loss.item(), + "learning_rate": lr_scheduler.get_last_lr()[0], + } + + if args.use_lpl and lpl_fn is not None and global_step >= args.lpl_start: + if lpl_mask.any(): + # LPL application statistics + log_data.update( + { + "lpl/loss": lpl_loss_value.item(), + "lpl/num_samples": lpl_mask.sum().item(), + "lpl/application_ratio": lpl_mask.float().mean().item(), + "lpl/weight": args.lpl_weight, + "lpl/weighted_loss": (args.lpl_weight * lpl_loss_value).item(), + } + ) + + # SNR statistics for LPL-applied samples + if args.snr_gamma is not None: + snr_values = snr[masked_indices] + log_data.update( + { + "lpl/snr_mean": snr_values.mean().item(), + "lpl/snr_std": snr_values.std().item(), + "lpl/snr_min": snr_values.min().item(), + "lpl/snr_max": snr_values.max().item(), + } + ) + + # Feature statistics if available + if hasattr(lpl_fn, "last_feature_stats"): + for layer_idx, stats in enumerate(lpl_fn.last_feature_stats): + log_data.update( + { + f"lpl/features/layer_{layer_idx}/mean": stats["mean"], + f"lpl/features/layer_{layer_idx}/std": stats["std"], + f"lpl/features/layer_{layer_idx}/outlier_ratio": stats.get( + "outlier_ratio", 0.0 + ), + } + ) + + # Memory usage if available + if torch.cuda.is_available(): + log_data.update( + { + "lpl/memory/allocated": torch.cuda.memory_allocated() / 1024**2, # MB + "lpl/memory/reserved": torch.cuda.memory_reserved() / 1024**2, # MB + } + ) + + # Log to accelerator + accelerator.log(log_data, step=global_step) + + # Update progress bar with more metrics + progress_bar_logs = { + "loss": loss.detach().item(), + "lr": lr_scheduler.get_last_lr()[0], + } + if args.use_lpl and lpl_loss_value.item() > 0: + progress_bar_logs.update( + { + "lpl": lpl_loss_value.item(), + "lpl_ratio": lpl_mask.float().mean().item() if lpl_mask.any() else 0.0, + } + ) + progress_bar.set_postfix(**progress_bar_logs) + + # DeepSpeed requires saving weights on every device; saving weights only on the main process would cause issues. + if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process: + if global_step % args.checkpointing_steps == 0: + # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` + if args.checkpoints_total_limit is not None: + checkpoints = os.listdir(args.output_dir) + checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] + checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) + + # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints + if len(checkpoints) >= args.checkpoints_total_limit: + num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 + removing_checkpoints = checkpoints[0:num_to_remove] + + logger.info( + f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" + ) + logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") + + for removing_checkpoint in removing_checkpoints: + removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) + shutil.rmtree(removing_checkpoint) + + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + accelerator.save_state(save_path) + logger.info(f"Saved state to {save_path}") + + logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} + progress_bar.set_postfix(**logs) + + if global_step >= args.max_train_steps: + break + + if accelerator.is_main_process: + if args.validation_prompt is not None and epoch % args.validation_epochs == 0: + logger.info( + f"Running validation... \n Generating {args.num_validation_images} images with prompt:" + f" {args.validation_prompt}." + ) + if args.use_ema: + # Store the UNet parameters temporarily and load the EMA parameters to perform inference. + ema_unet.store(unet.parameters()) + ema_unet.copy_to(unet.parameters()) + + # create pipeline + vae = AutoencoderKL.from_pretrained( + vae_path, + subfolder="vae" if args.pretrained_vae_model_name_or_path is None else None, + revision=args.revision, + variant=args.variant, + ) + pipeline = StableDiffusionXLPipeline.from_pretrained( + args.pretrained_model_name_or_path, + vae=vae, + unet=accelerator.unwrap_model(unet), + revision=args.revision, + variant=args.variant, + torch_dtype=weight_dtype, + ) + if args.prediction_type is not None: + scheduler_args = {"prediction_type": args.prediction_type} + pipeline.scheduler = pipeline.scheduler.from_config(pipeline.scheduler.config, **scheduler_args) + + pipeline = pipeline.to(accelerator.device) + pipeline.set_progress_bar_config(disable=True) + + # run inference + generator = ( + torch.Generator(device=accelerator.device).manual_seed(args.seed) + if args.seed is not None + else None + ) + pipeline_args = {"prompt": args.validation_prompt} + + with autocast_ctx: + images = [ + pipeline(**pipeline_args, generator=generator, num_inference_steps=25).images[0] + for _ in range(args.num_validation_images) + ] + + for tracker in accelerator.trackers: + if tracker.name == "tensorboard": + np_images = np.stack([np.asarray(img) for img in images]) + tracker.writer.add_images("validation", np_images, epoch, dataformats="NHWC") + if tracker.name == "wandb": + tracker.log( + { + "validation": [ + wandb.Image(image, caption=f"{i}: {args.validation_prompt}") + for i, image in enumerate(images) + ] + } + ) + + del pipeline + if is_torch_npu_available(): + torch_npu.npu.empty_cache() + elif torch.cuda.is_available(): + torch.cuda.empty_cache() + + if args.use_ema: + # Switch back to the original UNet parameters. + ema_unet.restore(unet.parameters()) + + accelerator.wait_for_everyone() + if accelerator.is_main_process: + unet = unwrap_model(unet) + if args.use_ema: + ema_unet.copy_to(unet.parameters()) + + # Serialize pipeline. + vae = AutoencoderKL.from_pretrained( + vae_path, + subfolder="vae" if args.pretrained_vae_model_name_or_path is None else None, + revision=args.revision, + variant=args.variant, + torch_dtype=weight_dtype, + ) + pipeline = StableDiffusionXLPipeline.from_pretrained( + args.pretrained_model_name_or_path, + unet=unet, + vae=vae, + revision=args.revision, + variant=args.variant, + torch_dtype=weight_dtype, + ) + if args.prediction_type is not None: + scheduler_args = {"prediction_type": args.prediction_type} + pipeline.scheduler = pipeline.scheduler.from_config(pipeline.scheduler.config, **scheduler_args) + pipeline.save_pretrained(args.output_dir) + + # run inference + images = [] + if args.validation_prompt and args.num_validation_images > 0: + pipeline = pipeline.to(accelerator.device) + generator = ( + torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed is not None else None + ) + + with autocast_ctx: + images = [ + pipeline(args.validation_prompt, num_inference_steps=25, generator=generator).images[0] + for _ in range(args.num_validation_images) + ] + + for tracker in accelerator.trackers: + if tracker.name == "tensorboard": + np_images = np.stack([np.asarray(img) for img in images]) + tracker.writer.add_images("test", np_images, epoch, dataformats="NHWC") + if tracker.name == "wandb": + tracker.log( + { + "test": [ + wandb.Image(image, caption=f"{i}: {args.validation_prompt}") + for i, image in enumerate(images) + ] + } + ) + + if args.push_to_hub: + save_model_card( + repo_id=repo_id, + images=images, + validation_prompt=args.validation_prompt, + base_model=args.pretrained_model_name_or_path, + dataset_name=args.dataset_name, + repo_folder=args.output_dir, + vae_path=args.pretrained_vae_model_name_or_path, + ) + upload_folder( + repo_id=repo_id, + folder_path=args.output_dir, + commit_message="End of training", + ignore_patterns=["step_*", "epoch_*"], + ) + + accelerator.end_training() + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/examples/research_projects/onnxruntime/text_to_image/README.md b/examples/research_projects/onnxruntime/text_to_image/README.md index f398f081663a..1d688471ba74 100644 --- a/examples/research_projects/onnxruntime/text_to_image/README.md +++ b/examples/research_projects/onnxruntime/text_to_image/README.md @@ -4,7 +4,7 @@ The `train_text_to_image.py` script shows how to fine-tune stable diffusion mode ___Note___: -___This script is experimental. The script fine-tunes the whole model and often times the model overfits and runs into issues like catastrophic forgetting. It's recommended to try different hyperparamters to get the best result on your dataset.___ +___This script is experimental. The script fine-tunes the whole model and often times the model overfits and runs into issues like catastrophic forgetting. It's recommended to try different hyperparameters to get the best result on your dataset.___ ## Running locally with PyTorch diff --git a/examples/research_projects/pixart/controlnet_pixart_alpha.py b/examples/research_projects/pixart/controlnet_pixart_alpha.py index 8f2eb974398d..faf118798ecc 100644 --- a/examples/research_projects/pixart/controlnet_pixart_alpha.py +++ b/examples/research_projects/pixart/controlnet_pixart_alpha.py @@ -26,7 +26,7 @@ def __init__( norm_type: str = "ada_norm_single", norm_elementwise_affine: bool = False, norm_eps: float = 1e-6, - attention_type: Optional[str] = "default", + attention_type: str | None = "default", ): super().__init__() diff --git a/examples/research_projects/pixart/pipeline_pixart_alpha_controlnet.py b/examples/research_projects/pixart/pipeline_pixart_alpha_controlnet.py index 89228983d4d8..27273594f474 100644 --- a/examples/research_projects/pixart/pipeline_pixart_alpha_controlnet.py +++ b/examples/research_projects/pixart/pipeline_pixart_alpha_controlnet.py @@ -808,7 +808,7 @@ def __call__( # rc todo: controlnet_conditioning_scale: Union[float, List[float]] = 1.0, # rc todo: control_guidance_start = 0.0, # rc todo: control_guidance_end = 1.0, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: int = 1, diff --git a/examples/research_projects/promptdiffusion/convert_original_promptdiffusion_to_diffusers.py b/examples/research_projects/promptdiffusion/convert_original_promptdiffusion_to_diffusers.py index c9efcffa5bb8..31105b6a5890 100644 --- a/examples/research_projects/promptdiffusion/convert_original_promptdiffusion_to_diffusers.py +++ b/examples/research_projects/promptdiffusion/convert_original_promptdiffusion_to_diffusers.py @@ -1110,7 +1110,7 @@ def stable_unclip_image_encoder(original_config, local_files_only=False): def stable_unclip_image_noising_components( - original_config, clip_stats_path: Optional[str] = None, device: Optional[str] = None + original_config, clip_stats_path: str | None = None, device: str | None = None ): """ Returns the noising components for the img2img and txt2img unclip pipelines. @@ -1267,9 +1267,9 @@ def download_from_original_stable_diffusion_ckpt( upcast_attention: Optional[bool] = None, device: str = None, from_safetensors: bool = False, - stable_unclip: Optional[str] = None, - stable_unclip_prior: Optional[str] = None, - clip_stats_path: Optional[str] = None, + stable_unclip: str | None = None, + stable_unclip_prior: str | None = None, + clip_stats_path: str | None = None, controlnet: Optional[bool] = None, adapter: Optional[bool] = None, load_safety_checker: bool = True, diff --git a/examples/research_projects/promptdiffusion/pipeline_prompt_diffusion.py b/examples/research_projects/promptdiffusion/pipeline_prompt_diffusion.py index 233df1276563..8b23570aea77 100644 --- a/examples/research_projects/promptdiffusion/pipeline_prompt_diffusion.py +++ b/examples/research_projects/promptdiffusion/pipeline_prompt_diffusion.py @@ -923,7 +923,7 @@ def __call__( prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, ip_adapter_image: Optional[PipelineImageInput] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, controlnet_conditioning_scale: Union[float, List[float]] = 1.0, diff --git a/examples/research_projects/promptdiffusion/promptdiffusioncontrolnet.py b/examples/research_projects/promptdiffusion/promptdiffusioncontrolnet.py index 7853695f0566..aff0302ac202 100644 --- a/examples/research_projects/promptdiffusion/promptdiffusioncontrolnet.py +++ b/examples/research_projects/promptdiffusion/promptdiffusioncontrolnet.py @@ -111,7 +111,7 @@ def __init__( "CrossAttnDownBlock2D", "DownBlock2D", ), - mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn", + mid_block_type: str | None = "UNetMidBlock2DCrossAttn", only_cross_attention: Union[bool, Tuple[bool]] = False, block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280), layers_per_block: int = 2, @@ -123,12 +123,12 @@ def __init__( cross_attention_dim: int = 1280, transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1, encoder_hid_dim: Optional[int] = None, - encoder_hid_dim_type: Optional[str] = None, + encoder_hid_dim_type: str | None = None, attention_head_dim: Union[int, Tuple[int, ...]] = 8, num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None, use_linear_projection: bool = False, - class_embed_type: Optional[str] = None, - addition_embed_type: Optional[str] = None, + class_embed_type: str | None = None, + addition_embed_type: str | None = None, addition_time_embed_dim: Optional[int] = None, num_class_embeds: Optional[int] = None, upcast_attention: bool = False, diff --git a/examples/research_projects/rdm/pipeline_rdm.py b/examples/research_projects/rdm/pipeline_rdm.py index 9b696874c5d1..08a941438ac4 100644 --- a/examples/research_projects/rdm/pipeline_rdm.py +++ b/examples/research_projects/rdm/pipeline_rdm.py @@ -162,10 +162,10 @@ def __call__( guidance_scale: float = 7.5, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, latents: Optional[torch.Tensor] = None, prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, callback_steps: Optional[int] = 1, diff --git a/examples/research_projects/sdxl_flax/README.md b/examples/research_projects/sdxl_flax/README.md index dfbe90e63bde..e2c02d7bb617 100644 --- a/examples/research_projects/sdxl_flax/README.md +++ b/examples/research_projects/sdxl_flax/README.md @@ -231,7 +231,7 @@ images = generate(prompt, neg_prompt) print(f"First inference in {time.time() - start}") ``` -From this point forward, any calls to generate should result in a faster inference +From this point forward, Any calls to generate should result in a faster inference time and it won't change. ```python diff --git a/examples/research_projects/sdxl_flax/sdxl_single.py b/examples/research_projects/sdxl_flax/sdxl_single.py index 5b9b862d99b5..c3cbf6ca24f0 100644 --- a/examples/research_projects/sdxl_flax/sdxl_single.py +++ b/examples/research_projects/sdxl_flax/sdxl_single.py @@ -18,7 +18,7 @@ NUM_DEVICES = jax.device_count() # 1. Let's start by downloading the model and loading it into our pipeline class -# Adhering to JAX's functional approach, the model's parameters are returned seperatetely and +# Adhering to JAX's functional approach, the model's parameters are returned separately and # will have to be passed to the pipeline during inference pipeline, params = FlaxStableDiffusionXLPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", revision="refs/pr/95", split_head_dim=True diff --git a/examples/research_projects/sdxl_flax/sdxl_single_aot.py b/examples/research_projects/sdxl_flax/sdxl_single_aot.py index 08bd13902aa9..a0f6b243b5df 100644 --- a/examples/research_projects/sdxl_flax/sdxl_single_aot.py +++ b/examples/research_projects/sdxl_flax/sdxl_single_aot.py @@ -131,7 +131,7 @@ def generate(prompt, negative_prompt, seed=default_seed, guidance_scale=default_ images = generate(prompt, neg_prompt) print(f"First inference in {time.time() - start}") -# 9. From this point forward, any calls to generate should result in a faster inference +# 9. From this point forward, Any calls to generate should result in a faster inference # time and it won't change. start = time.time() prompt = "photo of a rhino dressed suit and tie sitting at a table in a bar with a bar stools, award winning photography, Elke vogelsang" diff --git a/examples/research_projects/vae/vae_roundtrip.py b/examples/research_projects/vae/vae_roundtrip.py index 922cb42615e9..bfb32e484381 100644 --- a/examples/research_projects/vae/vae_roundtrip.py +++ b/examples/research_projects/vae/vae_roundtrip.py @@ -41,10 +41,10 @@ def load_vae_model( *, device: torch.device, model_name_or_path: str, - revision: Optional[str], - variant: Optional[str], + revision: str | None, + variant: str | None, # NOTE: use subfolder="vae" if the pointed model is for stable diffusion as a whole instead of just the VAE - subfolder: Optional[str], + subfolder: str | None, use_tiny_nn: bool, ) -> SupportedAutoencoder: if use_tiny_nn: @@ -156,9 +156,9 @@ def main_kwargs( device: torch.device, input_image_path: str, pretrained_model_name_or_path: str, - revision: Optional[str], - variant: Optional[str], - subfolder: Optional[str], + revision: str | None, + variant: str | None, + subfolder: str | None, use_tiny_nn: bool, ) -> None: vae = load_vae_model( diff --git a/examples/server-async/utils/requestscopedpipeline.py b/examples/server-async/utils/requestscopedpipeline.py index 57d1e2567169..d8d2ecd77e2c 100644 --- a/examples/server-async/utils/requestscopedpipeline.py +++ b/examples/server-async/utils/requestscopedpipeline.py @@ -7,16 +7,12 @@ from diffusers.utils import logging from .scheduler import BaseAsyncScheduler, async_retrieve_timesteps +from .wrappers import ThreadSafeImageProcessorWrapper, ThreadSafeTokenizerWrapper, ThreadSafeVAEWrapper logger = logging.get_logger(__name__) -def safe_tokenize(tokenizer, *args, lock, **kwargs): - with lock: - return tokenizer(*args, **kwargs) - - class RequestScopedPipeline: DEFAULT_MUTABLE_ATTRS = [ "_all_hooks", @@ -38,24 +34,41 @@ def __init__( wrap_scheduler: bool = True, ): self._base = pipeline + self.unet = getattr(pipeline, "unet", None) self.vae = getattr(pipeline, "vae", None) self.text_encoder = getattr(pipeline, "text_encoder", None) self.components = getattr(pipeline, "components", None) + self.transformer = getattr(pipeline, "transformer", None) + if wrap_scheduler and hasattr(pipeline, "scheduler") and pipeline.scheduler is not None: if not isinstance(pipeline.scheduler, BaseAsyncScheduler): pipeline.scheduler = BaseAsyncScheduler(pipeline.scheduler) self._mutable_attrs = list(mutable_attrs) if mutable_attrs is not None else list(self.DEFAULT_MUTABLE_ATTRS) + self._tokenizer_lock = tokenizer_lock if tokenizer_lock is not None else threading.Lock() + self._vae_lock = threading.Lock() + self._image_lock = threading.Lock() + self._auto_detect_mutables = bool(auto_detect_mutables) self._tensor_numel_threshold = int(tensor_numel_threshold) - self._auto_detected_attrs: List[str] = [] - def _make_local_scheduler(self, num_inference_steps: int, device: Optional[str] = None, **clone_kwargs): + def _detect_kernel_pipeline(self, pipeline) -> bool: + kernel_indicators = [ + "text_encoding_cache", + "memory_manager", + "enable_optimizations", + "_create_request_context", + "get_optimization_stats", + ] + + return any(hasattr(pipeline, attr) for attr in kernel_indicators) + + def _make_local_scheduler(self, num_inference_steps: int, device: str | None = None, **clone_kwargs): base_sched = getattr(self._base, "scheduler", None) if base_sched is None: return None @@ -70,11 +83,21 @@ def _make_local_scheduler(self, num_inference_steps: int, device: Optional[str] num_inference_steps=num_inference_steps, device=device, **clone_kwargs ) except Exception as e: - logger.debug(f"clone_for_request failed: {e}; falling back to deepcopy()") + logger.debug(f"clone_for_request failed: {e}; trying shallow copy fallback") try: - return copy.deepcopy(wrapped_scheduler) - except Exception as e: - logger.warning(f"Deepcopy of scheduler failed: {e}. Returning original scheduler (*risky*).") + if hasattr(wrapped_scheduler, "scheduler"): + try: + copied_scheduler = copy.copy(wrapped_scheduler.scheduler) + return BaseAsyncScheduler(copied_scheduler) + except Exception: + return wrapped_scheduler + else: + copied_scheduler = copy.copy(wrapped_scheduler) + return BaseAsyncScheduler(copied_scheduler) + except Exception as e2: + logger.warning( + f"Shallow copy of scheduler also failed: {e2}. Using original scheduler (*thread-unsafe but functional*)." + ) return wrapped_scheduler def _autodetect_mutables(self, max_attrs: int = 40): @@ -86,6 +109,7 @@ def _autodetect_mutables(self, max_attrs: int = 40): candidates: List[str] = [] seen = set() + for name in dir(self._base): if name.startswith("__"): continue @@ -93,6 +117,7 @@ def _autodetect_mutables(self, max_attrs: int = 40): continue if name in ("to", "save_pretrained", "from_pretrained"): continue + try: val = getattr(self._base, name) except Exception: @@ -100,11 +125,9 @@ def _autodetect_mutables(self, max_attrs: int = 40): import types - # skip callables and modules if callable(val) or isinstance(val, (types.ModuleType, types.FunctionType, types.MethodType)): continue - # containers -> candidate if isinstance(val, (dict, list, set, tuple, bytearray)): candidates.append(name) seen.add(name) @@ -205,7 +228,10 @@ def _is_tokenizer_component(self, component) -> bool: return has_tokenizer_methods and (has_tokenizer_in_name or has_tokenizer_attrs) - def generate(self, *args, num_inference_steps: int = 50, device: Optional[str] = None, **kwargs): + def _should_wrap_tokenizers(self) -> bool: + return True + + def generate(self, *args, num_inference_steps: int = 50, device: str | None = None, **kwargs): local_scheduler = self._make_local_scheduler(num_inference_steps=num_inference_steps, device=device) try: @@ -214,6 +240,25 @@ def generate(self, *args, num_inference_steps: int = 50, device: Optional[str] = logger.warning(f"copy.copy(self._base) failed: {e}. Falling back to deepcopy (may increase memory).") local_pipe = copy.deepcopy(self._base) + try: + if ( + hasattr(local_pipe, "vae") + and local_pipe.vae is not None + and not isinstance(local_pipe.vae, ThreadSafeVAEWrapper) + ): + local_pipe.vae = ThreadSafeVAEWrapper(local_pipe.vae, self._vae_lock) + + if ( + hasattr(local_pipe, "image_processor") + and local_pipe.image_processor is not None + and not isinstance(local_pipe.image_processor, ThreadSafeImageProcessorWrapper) + ): + local_pipe.image_processor = ThreadSafeImageProcessorWrapper( + local_pipe.image_processor, self._image_lock + ) + except Exception as e: + logger.debug(f"Could not wrap vae/image_processor: {e}") + if local_scheduler is not None: try: timesteps, num_steps, configured_scheduler = async_retrieve_timesteps( @@ -231,47 +276,42 @@ def generate(self, *args, num_inference_steps: int = 50, device: Optional[str] = self._clone_mutable_attrs(self._base, local_pipe) - # 4) wrap tokenizers on the local pipe with the lock wrapper - tokenizer_wrappers = {} # name -> original_tokenizer - try: - # a) wrap direct tokenizer attributes (tokenizer, tokenizer_2, ...) - for name in dir(local_pipe): - if "tokenizer" in name and not name.startswith("_"): - tok = getattr(local_pipe, name, None) - if tok is not None and self._is_tokenizer_component(tok): - tokenizer_wrappers[name] = tok - setattr( - local_pipe, - name, - lambda *args, tok=tok, **kwargs: safe_tokenize( - tok, *args, lock=self._tokenizer_lock, **kwargs - ), - ) - - # b) wrap tokenizers in components dict - if hasattr(local_pipe, "components") and isinstance(local_pipe.components, dict): - for key, val in local_pipe.components.items(): - if val is None: - continue - - if self._is_tokenizer_component(val): - tokenizer_wrappers[f"components[{key}]"] = val - local_pipe.components[key] = lambda *args, tokenizer=val, **kwargs: safe_tokenize( - tokenizer, *args, lock=self._tokenizer_lock, **kwargs - ) + original_tokenizers = {} - except Exception as e: - logger.debug(f"Tokenizer wrapping step encountered an error: {e}") + if self._should_wrap_tokenizers(): + try: + for name in dir(local_pipe): + if "tokenizer" in name and not name.startswith("_"): + tok = getattr(local_pipe, name, None) + if tok is not None and self._is_tokenizer_component(tok): + if not isinstance(tok, ThreadSafeTokenizerWrapper): + original_tokenizers[name] = tok + wrapped_tokenizer = ThreadSafeTokenizerWrapper(tok, self._tokenizer_lock) + setattr(local_pipe, name, wrapped_tokenizer) + + if hasattr(local_pipe, "components") and isinstance(local_pipe.components, dict): + for key, val in local_pipe.components.items(): + if val is None: + continue + + if self._is_tokenizer_component(val): + if not isinstance(val, ThreadSafeTokenizerWrapper): + original_tokenizers[f"components[{key}]"] = val + wrapped_tokenizer = ThreadSafeTokenizerWrapper(val, self._tokenizer_lock) + local_pipe.components[key] = wrapped_tokenizer + + except Exception as e: + logger.debug(f"Tokenizer wrapping step encountered an error: {e}") result = None cm = getattr(local_pipe, "model_cpu_offload_context", None) + try: if callable(cm): try: with cm(): result = local_pipe(*args, num_inference_steps=num_inference_steps, **kwargs) except TypeError: - # cm might be a context manager instance rather than callable try: with cm: result = local_pipe(*args, num_inference_steps=num_inference_steps, **kwargs) @@ -279,18 +319,18 @@ def generate(self, *args, num_inference_steps: int = 50, device: Optional[str] = logger.debug(f"model_cpu_offload_context usage failed: {e}. Proceeding without it.") result = local_pipe(*args, num_inference_steps=num_inference_steps, **kwargs) else: - # no offload context available — call directly result = local_pipe(*args, num_inference_steps=num_inference_steps, **kwargs) return result finally: try: - for name, tok in tokenizer_wrappers.items(): + for name, tok in original_tokenizers.items(): if name.startswith("components["): key = name[len("components[") : -1] - local_pipe.components[key] = tok + if hasattr(local_pipe, "components") and isinstance(local_pipe.components, dict): + local_pipe.components[key] = tok else: setattr(local_pipe, name, tok) except Exception as e: - logger.debug(f"Error restoring wrapped tokenizers: {e}") + logger.debug(f"Error restoring original tokenizers: {e}") diff --git a/examples/server-async/utils/wrappers.py b/examples/server-async/utils/wrappers.py new file mode 100644 index 000000000000..1e8474eabf3f --- /dev/null +++ b/examples/server-async/utils/wrappers.py @@ -0,0 +1,86 @@ +class ThreadSafeTokenizerWrapper: + def __init__(self, tokenizer, lock): + self._tokenizer = tokenizer + self._lock = lock + + self._thread_safe_methods = { + "__call__", + "encode", + "decode", + "tokenize", + "encode_plus", + "batch_encode_plus", + "batch_decode", + } + + def __getattr__(self, name): + attr = getattr(self._tokenizer, name) + + if name in self._thread_safe_methods and callable(attr): + + def wrapped_method(*args, **kwargs): + with self._lock: + return attr(*args, **kwargs) + + return wrapped_method + + return attr + + def __call__(self, *args, **kwargs): + with self._lock: + return self._tokenizer(*args, **kwargs) + + def __setattr__(self, name, value): + if name.startswith("_"): + super().__setattr__(name, value) + else: + setattr(self._tokenizer, name, value) + + def __dir__(self): + return dir(self._tokenizer) + + +class ThreadSafeVAEWrapper: + def __init__(self, vae, lock): + self._vae = vae + self._lock = lock + + def __getattr__(self, name): + attr = getattr(self._vae, name) + if name in {"decode", "encode", "forward"} and callable(attr): + + def wrapped(*args, **kwargs): + with self._lock: + return attr(*args, **kwargs) + + return wrapped + return attr + + def __setattr__(self, name, value): + if name.startswith("_"): + super().__setattr__(name, value) + else: + setattr(self._vae, name, value) + + +class ThreadSafeImageProcessorWrapper: + def __init__(self, proc, lock): + self._proc = proc + self._lock = lock + + def __getattr__(self, name): + attr = getattr(self._proc, name) + if name in {"postprocess", "preprocess"} and callable(attr): + + def wrapped(*args, **kwargs): + with self._lock: + return attr(*args, **kwargs) + + return wrapped + return attr + + def __setattr__(self, name, value): + if name.startswith("_"): + super().__setattr__(name, value) + else: + setattr(self._proc, name, value) diff --git a/examples/t2i_adapter/train_t2i_adapter_sdxl.py b/examples/t2i_adapter/train_t2i_adapter_sdxl.py index 32128ebbd4df..eb393418c5d7 100644 --- a/examples/t2i_adapter/train_t2i_adapter_sdxl.py +++ b/examples/t2i_adapter/train_t2i_adapter_sdxl.py @@ -61,7 +61,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/text_to_image/train_text_to_image.py b/examples/text_to_image/train_text_to_image.py index 90dd06d33c5e..7b76594a8dd0 100644 --- a/examples/text_to_image/train_text_to_image.py +++ b/examples/text_to_image/train_text_to_image.py @@ -57,7 +57,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__, log_level="INFO") diff --git a/examples/text_to_image/train_text_to_image_flax.py b/examples/text_to_image/train_text_to_image_flax.py index e474445d9afe..4fe710089981 100644 --- a/examples/text_to_image/train_text_to_image_flax.py +++ b/examples/text_to_image/train_text_to_image_flax.py @@ -49,7 +49,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = logging.getLogger(__name__) diff --git a/examples/text_to_image/train_text_to_image_lora.py b/examples/text_to_image/train_text_to_image_lora.py index 310a50ac4e9a..55c2c42d74c0 100644 --- a/examples/text_to_image/train_text_to_image_lora.py +++ b/examples/text_to_image/train_text_to_image_lora.py @@ -61,7 +61,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__, log_level="INFO") diff --git a/examples/text_to_image/train_text_to_image_lora_sdxl.py b/examples/text_to_image/train_text_to_image_lora_sdxl.py index 88f5c3cede6e..e211ad95ff43 100644 --- a/examples/text_to_image/train_text_to_image_lora_sdxl.py +++ b/examples/text_to_image/train_text_to_image_lora_sdxl.py @@ -68,7 +68,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) if is_torch_npu_available(): diff --git a/examples/text_to_image/train_text_to_image_sdxl.py b/examples/text_to_image/train_text_to_image_sdxl.py index 4eafa8f28a19..95749d4dcde4 100644 --- a/examples/text_to_image/train_text_to_image_sdxl.py +++ b/examples/text_to_image/train_text_to_image_sdxl.py @@ -55,7 +55,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) if is_torch_npu_available(): diff --git a/examples/textual_inversion/textual_inversion.py b/examples/textual_inversion/textual_inversion.py index 0d8c25349fca..1aaa701d8ceb 100644 --- a/examples/textual_inversion/textual_inversion.py +++ b/examples/textual_inversion/textual_inversion.py @@ -82,7 +82,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/textual_inversion/textual_inversion_flax.py b/examples/textual_inversion/textual_inversion_flax.py index 7fb394a1bd15..66a5da1fcd8f 100644 --- a/examples/textual_inversion/textual_inversion_flax.py +++ b/examples/textual_inversion/textual_inversion_flax.py @@ -56,7 +56,7 @@ # ------------------------------------------------------------------------------ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = logging.getLogger(__name__) diff --git a/examples/textual_inversion/textual_inversion_sdxl.py b/examples/textual_inversion/textual_inversion_sdxl.py index 3f482341ca4a..3e9151034eaa 100644 --- a/examples/textual_inversion/textual_inversion_sdxl.py +++ b/examples/textual_inversion/textual_inversion_sdxl.py @@ -77,7 +77,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__) diff --git a/examples/unconditional_image_generation/train_unconditional.py b/examples/unconditional_image_generation/train_unconditional.py index ed7d2db43700..649fc8c2facd 100644 --- a/examples/unconditional_image_generation/train_unconditional.py +++ b/examples/unconditional_image_generation/train_unconditional.py @@ -29,7 +29,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__, log_level="INFO") diff --git a/examples/vqgan/train_vqgan.py b/examples/vqgan/train_vqgan.py index d9ad2774e897..4684c9ce61c6 100644 --- a/examples/vqgan/train_vqgan.py +++ b/examples/vqgan/train_vqgan.py @@ -50,7 +50,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.37.0.dev0") +check_min_version("0.38.0.dev0") logger = get_logger(__name__, log_level="INFO") diff --git a/scripts/convert_animatediff_sparsectrl_to_diffusers.py b/scripts/convert_animatediff_sparsectrl_to_diffusers.py index f246dceb97f8..b96b42d60f44 100644 --- a/scripts/convert_animatediff_sparsectrl_to_diffusers.py +++ b/scripts/convert_animatediff_sparsectrl_to_diffusers.py @@ -18,7 +18,7 @@ } -def convert(original_state_dict: Dict[str, nn.Module]) -> Dict[str, nn.Module]: +def convert(original_state_dict: Dict[str, nn.Module]) -> dict[str, nn.Module]: converted_state_dict = {} for key in list(original_state_dict.keys()): diff --git a/scripts/convert_asymmetric_vqgan_to_diffusers.py b/scripts/convert_asymmetric_vqgan_to_diffusers.py index ffb735e18224..a9f83a46c2d0 100644 --- a/scripts/convert_asymmetric_vqgan_to_diffusers.py +++ b/scripts/convert_asymmetric_vqgan_to_diffusers.py @@ -61,7 +61,7 @@ } -def convert_asymmetric_autoencoder_kl_state_dict(original_state_dict: Dict[str, Any]) -> Dict[str, Any]: +def convert_asymmetric_autoencoder_kl_state_dict(original_state_dict: Dict[str, Any]) -> dict[str, Any]: converted_state_dict = {} for k, v in original_state_dict.items(): if k.startswith("encoder."): diff --git a/scripts/convert_cogvideox_to_diffusers.py b/scripts/convert_cogvideox_to_diffusers.py index 7eeed240c4de..5220aa0bde9d 100644 --- a/scripts/convert_cogvideox_to_diffusers.py +++ b/scripts/convert_cogvideox_to_diffusers.py @@ -120,7 +120,7 @@ def replace_up_keys_inplace(key: str, state_dict: Dict[str, Any]): TOKENIZER_MAX_LENGTH = 226 -def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]: +def get_state_dict(saved_dict: Dict[str, Any]) -> dict[str, Any]: state_dict = saved_dict if "model" in saved_dict.keys(): state_dict = state_dict["model"] @@ -131,7 +131,7 @@ def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]: return state_dict -def update_state_dict_inplace(state_dict: Dict[str, Any], old_key: str, new_key: str) -> Dict[str, Any]: +def update_state_dict_inplace(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]: state_dict[new_key] = state_dict.pop(old_key) diff --git a/scripts/convert_cosmos_to_diffusers.py b/scripts/convert_cosmos_to_diffusers.py index bc6014068e87..b6f321b5b2fb 100644 --- a/scripts/convert_cosmos_to_diffusers.py +++ b/scripts/convert_cosmos_to_diffusers.py @@ -78,12 +78,85 @@ --save_pipeline ``` +# Cosmos 2.5 Transfer + +Download checkpoint +```bash +hf download nvidia/Cosmos-Transfer2.5-2B +``` + +Convert checkpoint +```bash +# depth +transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Transfer2.5-2B/snapshots/eb5325b77d358944da58a690157dd2b8071bbf85/general/depth/626e6618-bfcd-4d9a-a077-1409e2ce353f_ema_bf16.pt + +python scripts/convert_cosmos_to_diffusers.py \ + --transformer_type Cosmos-2.5-Transfer-General-2B \ + --transformer_ckpt_path $transformer_ckpt_path \ + --vae_type wan2.1 \ + --output_path converted/transfer/2b/general/depth/pipeline \ + --save_pipeline + +python scripts/convert_cosmos_to_diffusers.py \ + --transformer_type Cosmos-2.5-Transfer-General-2B \ + --transformer_ckpt_path $transformer_ckpt_path \ + --vae_type wan2.1 \ + --output_path converted/transfer/2b/general/depth/models + +# edge +transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Transfer2.5-2B/snapshots/eb5325b77d358944da58a690157dd2b8071bbf85/general/edge/61f5694b-0ad5-4ecd-8ad7-c8545627d125_ema_bf16.pt + +python scripts/convert_cosmos_to_diffusers.py \ + --transformer_type Cosmos-2.5-Transfer-General-2B \ + --transformer_ckpt_path $transformer_ckpt_path \ + --vae_type wan2.1 \ + --output_path converted/transfer/2b/general/edge/pipeline \ + --save_pipeline + +python scripts/convert_cosmos_to_diffusers.py \ + --transformer_type Cosmos-2.5-Transfer-General-2B \ + --transformer_ckpt_path $transformer_ckpt_path \ + --vae_type wan2.1 \ + --output_path converted/transfer/2b/general/edge/models + +# blur +transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Transfer2.5-2B/snapshots/eb5325b77d358944da58a690157dd2b8071bbf85/general/blur/ba2f44f2-c726-4fe7-949f-597069d9b91c_ema_bf16.pt + +python scripts/convert_cosmos_to_diffusers.py \ + --transformer_type Cosmos-2.5-Transfer-General-2B \ + --transformer_ckpt_path $transformer_ckpt_path \ + --vae_type wan2.1 \ + --output_path converted/transfer/2b/general/blur/pipeline \ + --save_pipeline + +python scripts/convert_cosmos_to_diffusers.py \ + --transformer_type Cosmos-2.5-Transfer-General-2B \ + --transformer_ckpt_path $transformer_ckpt_path \ + --vae_type wan2.1 \ + --output_path converted/transfer/2b/general/blur/models + +# seg +transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Transfer2.5-2B/snapshots/eb5325b77d358944da58a690157dd2b8071bbf85/general/seg/5136ef49-6d8d-42e8-8abf-7dac722a304a_ema_bf16.pt + +python scripts/convert_cosmos_to_diffusers.py \ + --transformer_type Cosmos-2.5-Transfer-General-2B \ + --transformer_ckpt_path $transformer_ckpt_path \ + --vae_type wan2.1 \ + --output_path converted/transfer/2b/general/seg/pipeline \ + --save_pipeline + +python scripts/convert_cosmos_to_diffusers.py \ + --transformer_type Cosmos-2.5-Transfer-General-2B \ + --transformer_ckpt_path $transformer_ckpt_path \ + --vae_type wan2.1 \ + --output_path converted/transfer/2b/general/seg/models +``` """ import argparse import pathlib import sys -from typing import Any, Dict +from typing import Any, Dict, Optional import torch from accelerate import init_empty_weights @@ -95,6 +168,7 @@ AutoencoderKLWan, Cosmos2TextToImagePipeline, Cosmos2VideoToWorldPipeline, + CosmosControlNetModel, CosmosTextToWorldPipeline, CosmosTransformer3DModel, CosmosVideoToWorldPipeline, @@ -103,13 +177,14 @@ UniPCMultistepScheduler, ) from diffusers.pipelines.cosmos.pipeline_cosmos2_5_predict import Cosmos2_5_PredictBasePipeline +from diffusers.pipelines.cosmos.pipeline_cosmos2_5_transfer import Cosmos2_5_TransferPipeline def remove_keys_(key: str, state_dict: Dict[str, Any]): state_dict.pop(key) -def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> Dict[str, Any]: +def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]: state_dict[new_key] = state_dict.pop(old_key) @@ -356,8 +431,62 @@ def rename_transformer_blocks_(key: str, state_dict: Dict[str, Any]): "crossattn_proj_in_channels": 100352, "encoder_hidden_states_channels": 1024, }, + "Cosmos-2.5-Transfer-General-2B": { + "in_channels": 16 + 1, + "out_channels": 16, + "num_attention_heads": 16, + "attention_head_dim": 128, + "num_layers": 28, + "mlp_ratio": 4.0, + "text_embed_dim": 1024, + "adaln_lora_dim": 256, + "max_size": (128, 240, 240), + "patch_size": (1, 2, 2), + "rope_scale": (1.0, 3.0, 3.0), + "concat_padding_mask": True, + "extra_pos_embed_type": None, + "use_crossattn_projection": True, + "crossattn_proj_in_channels": 100352, + "encoder_hidden_states_channels": 1024, + "controlnet_block_every_n": 7, + "img_context_dim_in": 1152, + "img_context_dim_out": 2048, + "img_context_num_tokens": 256, + }, +} + +CONTROLNET_CONFIGS = { + "Cosmos-2.5-Transfer-General-2B": { + "n_controlnet_blocks": 4, + "model_channels": 2048, + "in_channels": 130, + "latent_channels": 18, # (16 latent + 1 condition_mask) + 1 padding_mask = 18 + "num_attention_heads": 16, + "attention_head_dim": 128, + "mlp_ratio": 4.0, + "text_embed_dim": 1024, + "adaln_lora_dim": 256, + "patch_size": (1, 2, 2), + "max_size": (128, 240, 240), + "rope_scale": (1.0, 3.0, 3.0), + "extra_pos_embed_type": None, + "img_context_dim_in": 1152, + "img_context_dim_out": 2048, + "use_crossattn_projection": True, + "crossattn_proj_in_channels": 100352, + "encoder_hidden_states_channels": 1024, + }, +} + +CONTROLNET_KEYS_RENAME_DICT = { + **TRANSFORMER_KEYS_RENAME_DICT_COSMOS_2_0, + "blocks": "blocks", + "control_embedder.proj.1": "patch_embed.proj", } + +CONTROLNET_SPECIAL_KEYS_REMAP = {**TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_2_0} + VAE_KEYS_RENAME_DICT = { "down.0": "down_blocks.0", "down.1": "down_blocks.1", @@ -436,7 +565,7 @@ def rename_transformer_blocks_(key: str, state_dict: Dict[str, Any]): } -def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]: +def get_state_dict(saved_dict: Dict[str, Any]) -> dict[str, Any]: state_dict = saved_dict if "model" in saved_dict.keys(): state_dict = state_dict["model"] @@ -447,9 +576,12 @@ def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]: return state_dict -def convert_transformer(transformer_type: str, ckpt_path: str, weights_only: bool = True): +def convert_transformer( + transformer_type: str, + state_dict: Optional[Dict[str, Any]] = None, + weights_only: bool = True, +): PREFIX_KEY = "net." - original_state_dict = get_state_dict(torch.load(ckpt_path, map_location="cpu", weights_only=weights_only)) if "Cosmos-1.0" in transformer_type: TRANSFORMER_KEYS_RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT_COSMOS_1_0 @@ -467,23 +599,29 @@ def convert_transformer(transformer_type: str, ckpt_path: str, weights_only: boo config = TRANSFORMER_CONFIGS[transformer_type] transformer = CosmosTransformer3DModel(**config) - for key in list(original_state_dict.keys()): + old2new = {} + new2old = {} + for key in list(state_dict.keys()): new_key = key[:] if new_key.startswith(PREFIX_KEY): new_key = new_key.removeprefix(PREFIX_KEY) for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items(): new_key = new_key.replace(replace_key, rename_key) print(key, "->", new_key, flush=True) - update_state_dict_(original_state_dict, key, new_key) + assert new_key not in new2old, f"new key {new_key} already mapped" + assert key not in old2new, f"old key {key} already mapped" + old2new[key] = new_key + new2old[new_key] = key + update_state_dict_(state_dict, key, new_key) - for key in list(original_state_dict.keys()): + for key in list(state_dict.keys()): for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items(): if special_key not in key: continue - handler_fn_inplace(key, original_state_dict) + handler_fn_inplace(key, state_dict) expected_keys = set(transformer.state_dict().keys()) - mapped_keys = set(original_state_dict.keys()) + mapped_keys = set(state_dict.keys()) missing_keys = expected_keys - mapped_keys unexpected_keys = mapped_keys - expected_keys if missing_keys: @@ -497,10 +635,86 @@ def convert_transformer(transformer_type: str, ckpt_path: str, weights_only: boo print(k) sys.exit(2) - transformer.load_state_dict(original_state_dict, strict=True, assign=True) + transformer.load_state_dict(state_dict, strict=True, assign=True) return transformer +def convert_controlnet( + transformer_type: str, + control_state_dict: Dict[str, Any], + base_state_dict: Dict[str, Any], + weights_only: bool = True, +): + """ + Convert controlnet weights. + + Args: + transformer_type: The type of transformer/controlnet + control_state_dict: State dict containing controlnet-specific weights + base_state_dict: State dict containing base transformer weights (for shared modules) + weights_only: Whether to use weights_only loading + """ + if transformer_type not in CONTROLNET_CONFIGS: + raise AssertionError(f"{transformer_type} does not define a ControlNet config") + + PREFIX_KEY = "net." + + # Process control-specific keys + for key in list(control_state_dict.keys()): + new_key = key[:] + if new_key.startswith(PREFIX_KEY): + new_key = new_key.removeprefix(PREFIX_KEY) + for replace_key, rename_key in CONTROLNET_KEYS_RENAME_DICT.items(): + new_key = new_key.replace(replace_key, rename_key) + update_state_dict_(control_state_dict, key, new_key) + + for key in list(control_state_dict.keys()): + for special_key, handler_fn_inplace in CONTROLNET_SPECIAL_KEYS_REMAP.items(): + if special_key not in key: + continue + handler_fn_inplace(key, control_state_dict) + + # Copy shared weights from base transformer to controlnet + # These are the duplicated modules: patch_embed_base, time_embed, learnable_pos_embed, img_context_proj, crossattn_proj + shared_module_mappings = { + # transformer key prefix -> controlnet key prefix + "patch_embed.": "patch_embed_base.", + "time_embed.": "time_embed.", + "learnable_pos_embed.": "learnable_pos_embed.", + "img_context_proj.": "img_context_proj.", + "crossattn_proj.": "crossattn_proj.", + } + + for key in list(base_state_dict.keys()): + for transformer_prefix, controlnet_prefix in shared_module_mappings.items(): + if key.startswith(transformer_prefix): + controlnet_key = controlnet_prefix + key[len(transformer_prefix) :] + control_state_dict[controlnet_key] = base_state_dict[key].clone() + print(f"Copied shared weight: {key} -> {controlnet_key}", flush=True) + break + + cfg = CONTROLNET_CONFIGS[transformer_type] + controlnet = CosmosControlNetModel(**cfg) + + expected_keys = set(controlnet.state_dict().keys()) + mapped_keys = set(control_state_dict.keys()) + missing_keys = expected_keys - mapped_keys + unexpected_keys = mapped_keys - expected_keys + if missing_keys: + print(f"WARNING: missing controlnet keys ({len(missing_keys)}):", file=sys.stderr, flush=True) + for k in sorted(missing_keys): + print(k, file=sys.stderr) + sys.exit(3) + if unexpected_keys: + print(f"WARNING: unexpected controlnet keys ({len(unexpected_keys)}):", file=sys.stderr, flush=True) + for k in sorted(unexpected_keys): + print(k, file=sys.stderr) + sys.exit(4) + + controlnet.load_state_dict(control_state_dict, strict=True, assign=True) + return controlnet + + def convert_vae(vae_type: str): model_name = VAE_CONFIGS[vae_type]["name"] snapshot_directory = snapshot_download(model_name, repo_type="model") @@ -586,7 +800,7 @@ def save_pipeline_cosmos_2_0(args, transformer, vae): pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") -def save_pipeline_cosmos2_5(args, transformer, vae): +def save_pipeline_cosmos2_5_predict(args, transformer, vae): text_encoder_path = args.text_encoder_path or "nvidia/Cosmos-Reason1-7B" tokenizer_path = args.tokenizer_path or "Qwen/Qwen2.5-VL-7B-Instruct" @@ -614,6 +828,35 @@ def save_pipeline_cosmos2_5(args, transformer, vae): pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") +def save_pipeline_cosmos2_5_transfer(args, transformer, controlnet, vae): + text_encoder_path = args.text_encoder_path or "nvidia/Cosmos-Reason1-7B" + tokenizer_path = args.tokenizer_path or "Qwen/Qwen2.5-VL-7B-Instruct" + + text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( + text_encoder_path, torch_dtype="auto", device_map="cpu" + ) + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) + + scheduler = UniPCMultistepScheduler( + use_karras_sigmas=True, + use_flow_sigmas=True, + prediction_type="flow_prediction", + sigma_max=200.0, + sigma_min=0.01, + ) + + pipe = Cosmos2_5_TransferPipeline( + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + controlnet=controlnet, + vae=vae, + scheduler=scheduler, + safety_checker=lambda *args, **kwargs: None, + ) + pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + + def get_args(): parser = argparse.ArgumentParser() parser.add_argument("--transformer_type", type=str, default=None, choices=list(TRANSFORMER_CONFIGS.keys())) @@ -642,18 +885,61 @@ def get_args(): args = get_args() transformer = None + controlnet = None dtype = DTYPE_MAPPING[args.dtype] if args.save_pipeline: assert args.transformer_ckpt_path is not None assert args.vae_type is not None + raw_state_dict = None if args.transformer_ckpt_path is not None: weights_only = "Cosmos-1.0" in args.transformer_type - transformer = convert_transformer(args.transformer_type, args.transformer_ckpt_path, weights_only) - transformer = transformer.to(dtype=dtype) - if not args.save_pipeline: - transformer.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + raw_state_dict = get_state_dict( + torch.load(args.transformer_ckpt_path, map_location="cpu", weights_only=weights_only) + ) + + if raw_state_dict is not None: + if "Transfer" in args.transformer_type: + base_state_dict = {} + control_state_dict = {} + for k, v in raw_state_dict.items(): + plain_key = k.removeprefix("net.") if k.startswith("net.") else k + if "control" in plain_key.lower(): + control_state_dict[k] = v + else: + base_state_dict[k] = v + assert len(base_state_dict.keys() & control_state_dict.keys()) == 0 + + # Convert transformer first to get the processed base state dict + transformer = convert_transformer( + args.transformer_type, state_dict=base_state_dict, weights_only=weights_only + ) + transformer = transformer.to(dtype=dtype) + + # Get converted transformer state dict to copy shared weights to controlnet + converted_base_state_dict = transformer.state_dict() + + # Convert controlnet with both control-specific and shared weights from transformer + controlnet = convert_controlnet( + args.transformer_type, control_state_dict, converted_base_state_dict, weights_only=weights_only + ) + controlnet = controlnet.to(dtype=dtype) + + if not args.save_pipeline: + transformer.save_pretrained( + pathlib.Path(args.output_path) / "transformer", safe_serialization=True, max_shard_size="5GB" + ) + controlnet.save_pretrained( + pathlib.Path(args.output_path) / "controlnet", safe_serialization=True, max_shard_size="5GB" + ) + else: + transformer = convert_transformer( + args.transformer_type, state_dict=raw_state_dict, weights_only=weights_only + ) + transformer = transformer.to(dtype=dtype) + if not args.save_pipeline: + transformer.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") if args.vae_type is not None: if "Cosmos-1.0" in args.transformer_type: @@ -667,6 +953,8 @@ def get_args(): if not args.save_pipeline: vae.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + else: + vae = None if args.save_pipeline: if "Cosmos-1.0" in args.transformer_type: @@ -678,6 +966,11 @@ def get_args(): assert args.tokenizer_path is not None save_pipeline_cosmos_2_0(args, transformer, vae) elif "Cosmos-2.5" in args.transformer_type: - save_pipeline_cosmos2_5(args, transformer, vae) + if "Predict" in args.transformer_type: + save_pipeline_cosmos2_5_predict(args, transformer, vae) + elif "Transfer" in args.transformer_type: + save_pipeline_cosmos2_5_transfer(args, transformer, None, vae) + else: + raise AssertionError(f"{args.transformer_type} not supported") else: raise AssertionError(f"{args.transformer_type} not supported") diff --git a/scripts/convert_dcae_to_diffusers.py b/scripts/convert_dcae_to_diffusers.py index 15f79a8154e6..1a94c4b06640 100644 --- a/scripts/convert_dcae_to_diffusers.py +++ b/scripts/convert_dcae_to_diffusers.py @@ -79,7 +79,7 @@ def remap_proj_conv_(key: str, state_dict: Dict[str, Any]): } -def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]: +def get_state_dict(saved_dict: Dict[str, Any]) -> dict[str, Any]: state_dict = saved_dict if "model" in saved_dict.keys(): state_dict = state_dict["model"] @@ -90,7 +90,7 @@ def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]: return state_dict -def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> Dict[str, Any]: +def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]: state_dict[new_key] = state_dict.pop(old_key) diff --git a/scripts/convert_flux2_to_diffusers.py b/scripts/convert_flux2_to_diffusers.py index 2973913fa215..a8fa6f87eee1 100644 --- a/scripts/convert_flux2_to_diffusers.py +++ b/scripts/convert_flux2_to_diffusers.py @@ -44,7 +44,7 @@ parser = argparse.ArgumentParser() parser.add_argument("--original_state_dict_repo_id", default=None, type=str) parser.add_argument("--vae_filename", default="flux2-vae.sft", type=str) -parser.add_argument("--dit_filename", default="flux-dev-dummy.sft", type=str) +parser.add_argument("--dit_filename", default="flux2-dev.safetensors", type=str) parser.add_argument("--vae", action="store_true") parser.add_argument("--dit", action="store_true") parser.add_argument("--vae_dtype", type=str, default="fp32") @@ -385,9 +385,9 @@ def update_state_dict(state_dict: Dict[str, Any], old_key: str, new_key: str) -> def get_flux2_transformer_config(model_type: str) -> Tuple[Dict[str, Any], ...]: - if model_type == "test" or model_type == "dummy-flux2": + if model_type == "flux2-dev": config = { - "model_id": "diffusers-internal-dev/dummy-flux2", + "model_id": "black-forest-labs/FLUX.2-dev", "diffusers_config": { "patch_size": 1, "in_channels": 128, @@ -405,6 +405,53 @@ def get_flux2_transformer_config(model_type: str) -> Tuple[Dict[str, Any], ...]: } rename_dict = FLUX2_TRANSFORMER_KEYS_RENAME_DICT special_keys_remap = TRANSFORMER_SPECIAL_KEYS_REMAP + elif model_type == "klein-4b": + config = { + "model_id": "diffusers-internal-dev/dummy0115", + "diffusers_config": { + "patch_size": 1, + "in_channels": 128, + "num_layers": 5, + "num_single_layers": 20, + "attention_head_dim": 128, + "num_attention_heads": 24, + "joint_attention_dim": 7680, + "timestep_guidance_channels": 256, + "mlp_ratio": 3.0, + "axes_dims_rope": (32, 32, 32, 32), + "rope_theta": 2000, + "eps": 1e-6, + "guidance_embeds": False, + }, + } + rename_dict = FLUX2_TRANSFORMER_KEYS_RENAME_DICT + special_keys_remap = TRANSFORMER_SPECIAL_KEYS_REMAP + + elif model_type == "klein-9b": + config = { + "model_id": "diffusers-internal-dev/dummy0115", + "diffusers_config": { + "patch_size": 1, + "in_channels": 128, + "num_layers": 8, + "num_single_layers": 24, + "attention_head_dim": 128, + "num_attention_heads": 32, + "joint_attention_dim": 12288, + "timestep_guidance_channels": 256, + "mlp_ratio": 3.0, + "axes_dims_rope": (32, 32, 32, 32), + "rope_theta": 2000, + "eps": 1e-6, + "guidance_embeds": False, + }, + } + rename_dict = FLUX2_TRANSFORMER_KEYS_RENAME_DICT + special_keys_remap = TRANSFORMER_SPECIAL_KEYS_REMAP + + else: + raise ValueError(f"Unknown model_type: {model_type}. Choose from: flux2-dev, klein-4b, klein-9b") + return config, rename_dict, special_keys_remap @@ -447,7 +494,14 @@ def main(args): if args.dit: original_dit_ckpt = load_original_checkpoint(args, filename=args.dit_filename) - transformer = convert_flux2_transformer_to_diffusers(original_dit_ckpt, "test") + + if "klein-4b" in args.dit_filename: + model_type = "klein-4b" + elif "klein-9b" in args.dit_filename: + model_type = "klein-9b" + else: + model_type = "flux2-dev" + transformer = convert_flux2_transformer_to_diffusers(original_dit_ckpt, model_type) if not args.full_pipe: dit_dtype = torch.bfloat16 if args.dit_dtype == "bf16" else torch.float32 transformer.to(dit_dtype).save_pretrained(f"{args.output_path}/transformer") @@ -465,8 +519,15 @@ def main(args): "black-forest-labs/FLUX.1-dev", subfolder="scheduler" ) + if_distilled = "base" not in args.dit_filename + pipe = Flux2Pipeline( - vae=vae, transformer=transformer, text_encoder=text_encoder, tokenizer=tokenizer, scheduler=scheduler + vae=vae, + transformer=transformer, + text_encoder=text_encoder, + tokenizer=tokenizer, + scheduler=scheduler, + if_distilled=if_distilled, ) pipe.save_pretrained(args.output_path) diff --git a/scripts/convert_hunyuan_video_to_diffusers.py b/scripts/convert_hunyuan_video_to_diffusers.py index c84809d7f68a..fb604c98b65b 100644 --- a/scripts/convert_hunyuan_video_to_diffusers.py +++ b/scripts/convert_hunyuan_video_to_diffusers.py @@ -203,11 +203,11 @@ def remap_single_transformer_blocks_(key, state_dict): } -def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> Dict[str, Any]: +def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]: state_dict[new_key] = state_dict.pop(old_key) -def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]: +def get_state_dict(saved_dict: Dict[str, Any]) -> dict[str, Any]: state_dict = saved_dict if "model" in saved_dict.keys(): state_dict = state_dict["model"] diff --git a/scripts/convert_ltx2_to_diffusers.py b/scripts/convert_ltx2_to_diffusers.py new file mode 100644 index 000000000000..9dee954af6d0 --- /dev/null +++ b/scripts/convert_ltx2_to_diffusers.py @@ -0,0 +1,902 @@ +import argparse +import os +from contextlib import nullcontext +from typing import Any + +import safetensors.torch +import torch +from accelerate import init_empty_weights +from huggingface_hub import hf_hub_download +from transformers import AutoTokenizer, Gemma3ForConditionalGeneration + +from diffusers import ( + AutoencoderKLLTX2Audio, + AutoencoderKLLTX2Video, + FlowMatchEulerDiscreteScheduler, + LTX2LatentUpsamplePipeline, + LTX2Pipeline, + LTX2VideoTransformer3DModel, +) +from diffusers.pipelines.ltx2 import LTX2LatentUpsamplerModel, LTX2TextConnectors, LTX2Vocoder +from diffusers.utils.import_utils import is_accelerate_available + + +CTX = init_empty_weights if is_accelerate_available() else nullcontext + + +LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT = { + # Input Patchify Projections + "patchify_proj": "proj_in", + "audio_patchify_proj": "audio_proj_in", + # Modulation Parameters + # Handle adaln_single --> time_embed, audioln_single --> audio_time_embed separately as the original keys are + # substrings of the other modulation parameters below + "av_ca_video_scale_shift_adaln_single": "av_cross_attn_video_scale_shift", + "av_ca_a2v_gate_adaln_single": "av_cross_attn_video_a2v_gate", + "av_ca_audio_scale_shift_adaln_single": "av_cross_attn_audio_scale_shift", + "av_ca_v2a_gate_adaln_single": "av_cross_attn_audio_v2a_gate", + # Transformer Blocks + # Per-Block Cross Attention Modulatin Parameters + "scale_shift_table_a2v_ca_video": "video_a2v_cross_attn_scale_shift_table", + "scale_shift_table_a2v_ca_audio": "audio_a2v_cross_attn_scale_shift_table", + # Attention QK Norms + "q_norm": "norm_q", + "k_norm": "norm_k", +} + +LTX_2_0_VIDEO_VAE_RENAME_DICT = { + # Encoder + "down_blocks.0": "down_blocks.0", + "down_blocks.1": "down_blocks.0.downsamplers.0", + "down_blocks.2": "down_blocks.1", + "down_blocks.3": "down_blocks.1.downsamplers.0", + "down_blocks.4": "down_blocks.2", + "down_blocks.5": "down_blocks.2.downsamplers.0", + "down_blocks.6": "down_blocks.3", + "down_blocks.7": "down_blocks.3.downsamplers.0", + "down_blocks.8": "mid_block", + # Decoder + "up_blocks.0": "mid_block", + "up_blocks.1": "up_blocks.0.upsamplers.0", + "up_blocks.2": "up_blocks.0", + "up_blocks.3": "up_blocks.1.upsamplers.0", + "up_blocks.4": "up_blocks.1", + "up_blocks.5": "up_blocks.2.upsamplers.0", + "up_blocks.6": "up_blocks.2", + "last_time_embedder": "time_embedder", + "last_scale_shift_table": "scale_shift_table", + # Common + # For all 3D ResNets + "res_blocks": "resnets", + "per_channel_statistics.mean-of-means": "latents_mean", + "per_channel_statistics.std-of-means": "latents_std", +} + +LTX_2_0_AUDIO_VAE_RENAME_DICT = { + "per_channel_statistics.mean-of-means": "latents_mean", + "per_channel_statistics.std-of-means": "latents_std", +} + +LTX_2_0_VOCODER_RENAME_DICT = { + "ups": "upsamplers", + "resblocks": "resnets", + "conv_pre": "conv_in", + "conv_post": "conv_out", +} + +LTX_2_0_TEXT_ENCODER_RENAME_DICT = { + "video_embeddings_connector": "video_connector", + "audio_embeddings_connector": "audio_connector", + "transformer_1d_blocks": "transformer_blocks", + # Attention QK Norms + "q_norm": "norm_q", + "k_norm": "norm_k", +} + + +def update_state_dict_inplace(state_dict: dict[str, Any], old_key: str, new_key: str) -> None: + state_dict[new_key] = state_dict.pop(old_key) + + +def remove_keys_inplace(key: str, state_dict: dict[str, Any]) -> None: + state_dict.pop(key) + + +def convert_ltx2_transformer_adaln_single(key: str, state_dict: dict[str, Any]) -> None: + # Skip if not a weight, bias + if ".weight" not in key and ".bias" not in key: + return + + if key.startswith("adaln_single."): + new_key = key.replace("adaln_single.", "time_embed.") + param = state_dict.pop(key) + state_dict[new_key] = param + + if key.startswith("audio_adaln_single."): + new_key = key.replace("audio_adaln_single.", "audio_time_embed.") + param = state_dict.pop(key) + state_dict[new_key] = param + + return + + +def convert_ltx2_audio_vae_per_channel_statistics(key: str, state_dict: dict[str, Any]) -> None: + if key.startswith("per_channel_statistics"): + new_key = ".".join(["decoder", key]) + param = state_dict.pop(key) + state_dict[new_key] = param + + return + + +LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP = { + "video_embeddings_connector": remove_keys_inplace, + "audio_embeddings_connector": remove_keys_inplace, + "adaln_single": convert_ltx2_transformer_adaln_single, +} + +LTX_2_0_CONNECTORS_KEYS_RENAME_DICT = { + "connectors.": "", + "video_embeddings_connector": "video_connector", + "audio_embeddings_connector": "audio_connector", + "transformer_1d_blocks": "transformer_blocks", + "text_embedding_projection.aggregate_embed": "text_proj_in", + # Attention QK Norms + "q_norm": "norm_q", + "k_norm": "norm_k", +} + +LTX_2_0_VAE_SPECIAL_KEYS_REMAP = { + "per_channel_statistics.channel": remove_keys_inplace, + "per_channel_statistics.mean-of-stds": remove_keys_inplace, +} + +LTX_2_0_AUDIO_VAE_SPECIAL_KEYS_REMAP = {} + +LTX_2_0_VOCODER_SPECIAL_KEYS_REMAP = {} + + +def split_transformer_and_connector_state_dict(state_dict: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]: + connector_prefixes = ( + "video_embeddings_connector", + "audio_embeddings_connector", + "transformer_1d_blocks", + "text_embedding_projection.aggregate_embed", + "connectors.", + "video_connector", + "audio_connector", + "text_proj_in", + ) + + transformer_state_dict, connector_state_dict = {}, {} + for key, value in state_dict.items(): + if key.startswith(connector_prefixes): + connector_state_dict[key] = value + else: + transformer_state_dict[key] = value + + return transformer_state_dict, connector_state_dict + + +def get_ltx2_transformer_config(version: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + if version == "test": + # Produces a transformer of the same size as used in test_models_transformer_ltx2.py + config = { + "model_id": "diffusers-internal-dev/dummy-ltx2", + "diffusers_config": { + "in_channels": 4, + "out_channels": 4, + "patch_size": 1, + "patch_size_t": 1, + "num_attention_heads": 2, + "attention_head_dim": 8, + "cross_attention_dim": 16, + "vae_scale_factors": (8, 32, 32), + "pos_embed_max_pos": 20, + "base_height": 2048, + "base_width": 2048, + "audio_in_channels": 4, + "audio_out_channels": 4, + "audio_patch_size": 1, + "audio_patch_size_t": 1, + "audio_num_attention_heads": 2, + "audio_attention_head_dim": 4, + "audio_cross_attention_dim": 8, + "audio_scale_factor": 4, + "audio_pos_embed_max_pos": 20, + "audio_sampling_rate": 16000, + "audio_hop_length": 160, + "num_layers": 2, + "activation_fn": "gelu-approximate", + "qk_norm": "rms_norm_across_heads", + "norm_elementwise_affine": False, + "norm_eps": 1e-6, + "caption_channels": 16, + "attention_bias": True, + "attention_out_bias": True, + "rope_theta": 10000.0, + "rope_double_precision": False, + "causal_offset": 1, + "timestep_scale_multiplier": 1000, + "cross_attn_timestep_scale_multiplier": 1, + }, + } + rename_dict = LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT + special_keys_remap = LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP + elif version == "2.0": + config = { + "model_id": "diffusers-internal-dev/new-ltx-model", + "diffusers_config": { + "in_channels": 128, + "out_channels": 128, + "patch_size": 1, + "patch_size_t": 1, + "num_attention_heads": 32, + "attention_head_dim": 128, + "cross_attention_dim": 4096, + "vae_scale_factors": (8, 32, 32), + "pos_embed_max_pos": 20, + "base_height": 2048, + "base_width": 2048, + "audio_in_channels": 128, + "audio_out_channels": 128, + "audio_patch_size": 1, + "audio_patch_size_t": 1, + "audio_num_attention_heads": 32, + "audio_attention_head_dim": 64, + "audio_cross_attention_dim": 2048, + "audio_scale_factor": 4, + "audio_pos_embed_max_pos": 20, + "audio_sampling_rate": 16000, + "audio_hop_length": 160, + "num_layers": 48, + "activation_fn": "gelu-approximate", + "qk_norm": "rms_norm_across_heads", + "norm_elementwise_affine": False, + "norm_eps": 1e-6, + "caption_channels": 3840, + "attention_bias": True, + "attention_out_bias": True, + "rope_theta": 10000.0, + "rope_double_precision": True, + "causal_offset": 1, + "timestep_scale_multiplier": 1000, + "cross_attn_timestep_scale_multiplier": 1000, + "rope_type": "split", + }, + } + rename_dict = LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT + special_keys_remap = LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP + return config, rename_dict, special_keys_remap + + +def get_ltx2_connectors_config(version: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + if version == "test": + config = { + "model_id": "diffusers-internal-dev/dummy-ltx2", + "diffusers_config": { + "caption_channels": 16, + "text_proj_in_factor": 3, + "video_connector_num_attention_heads": 4, + "video_connector_attention_head_dim": 8, + "video_connector_num_layers": 1, + "video_connector_num_learnable_registers": None, + "audio_connector_num_attention_heads": 4, + "audio_connector_attention_head_dim": 8, + "audio_connector_num_layers": 1, + "audio_connector_num_learnable_registers": None, + "connector_rope_base_seq_len": 32, + "rope_theta": 10000.0, + "rope_double_precision": False, + "causal_temporal_positioning": False, + }, + } + elif version == "2.0": + config = { + "model_id": "diffusers-internal-dev/new-ltx-model", + "diffusers_config": { + "caption_channels": 3840, + "text_proj_in_factor": 49, + "video_connector_num_attention_heads": 30, + "video_connector_attention_head_dim": 128, + "video_connector_num_layers": 2, + "video_connector_num_learnable_registers": 128, + "audio_connector_num_attention_heads": 30, + "audio_connector_attention_head_dim": 128, + "audio_connector_num_layers": 2, + "audio_connector_num_learnable_registers": 128, + "connector_rope_base_seq_len": 4096, + "rope_theta": 10000.0, + "rope_double_precision": True, + "causal_temporal_positioning": False, + "rope_type": "split", + }, + } + + rename_dict = LTX_2_0_CONNECTORS_KEYS_RENAME_DICT + special_keys_remap = {} + + return config, rename_dict, special_keys_remap + + +def convert_ltx2_transformer(original_state_dict: dict[str, Any], version: str) -> dict[str, Any]: + config, rename_dict, special_keys_remap = get_ltx2_transformer_config(version) + diffusers_config = config["diffusers_config"] + + transformer_state_dict, _ = split_transformer_and_connector_state_dict(original_state_dict) + + with init_empty_weights(): + transformer = LTX2VideoTransformer3DModel.from_config(diffusers_config) + + # Handle official code --> diffusers key remapping via the remap dict + for key in list(transformer_state_dict.keys()): + new_key = key[:] + for replace_key, rename_key in rename_dict.items(): + new_key = new_key.replace(replace_key, rename_key) + update_state_dict_inplace(transformer_state_dict, key, new_key) + + # Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in + # special_keys_remap + for key in list(transformer_state_dict.keys()): + for special_key, handler_fn_inplace in special_keys_remap.items(): + if special_key not in key: + continue + handler_fn_inplace(key, transformer_state_dict) + + transformer.load_state_dict(transformer_state_dict, strict=True, assign=True) + return transformer + + +def convert_ltx2_connectors(original_state_dict: dict[str, Any], version: str) -> LTX2TextConnectors: + config, rename_dict, special_keys_remap = get_ltx2_connectors_config(version) + diffusers_config = config["diffusers_config"] + + _, connector_state_dict = split_transformer_and_connector_state_dict(original_state_dict) + if len(connector_state_dict) == 0: + raise ValueError("No connector weights found in the provided state dict.") + + with init_empty_weights(): + connectors = LTX2TextConnectors.from_config(diffusers_config) + + for key in list(connector_state_dict.keys()): + new_key = key[:] + for replace_key, rename_key in rename_dict.items(): + new_key = new_key.replace(replace_key, rename_key) + update_state_dict_inplace(connector_state_dict, key, new_key) + + for key in list(connector_state_dict.keys()): + for special_key, handler_fn_inplace in special_keys_remap.items(): + if special_key not in key: + continue + handler_fn_inplace(key, connector_state_dict) + + connectors.load_state_dict(connector_state_dict, strict=True, assign=True) + return connectors + + +def get_ltx2_video_vae_config( + version: str, timestep_conditioning: bool = False +) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + if version == "test": + config = { + "model_id": "diffusers-internal-dev/dummy-ltx2", + "diffusers_config": { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 128, + "block_out_channels": (256, 512, 1024, 2048), + "down_block_types": ( + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + ), + "decoder_block_out_channels": (256, 512, 1024), + "layers_per_block": (4, 6, 6, 2, 2), + "decoder_layers_per_block": (5, 5, 5, 5), + "spatio_temporal_scaling": (True, True, True, True), + "decoder_spatio_temporal_scaling": (True, True, True), + "decoder_inject_noise": (False, False, False, False), + "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"), + "upsample_residual": (True, True, True), + "upsample_factor": (2, 2, 2), + "timestep_conditioning": timestep_conditioning, + "patch_size": 4, + "patch_size_t": 1, + "resnet_norm_eps": 1e-6, + "encoder_causal": True, + "decoder_causal": False, + "encoder_spatial_padding_mode": "zeros", + "decoder_spatial_padding_mode": "reflect", + "spatial_compression_ratio": 32, + "temporal_compression_ratio": 8, + }, + } + rename_dict = LTX_2_0_VIDEO_VAE_RENAME_DICT + special_keys_remap = LTX_2_0_VAE_SPECIAL_KEYS_REMAP + elif version == "2.0": + config = { + "model_id": "diffusers-internal-dev/dummy-ltx2", + "diffusers_config": { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 128, + "block_out_channels": (256, 512, 1024, 2048), + "down_block_types": ( + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + ), + "decoder_block_out_channels": (256, 512, 1024), + "layers_per_block": (4, 6, 6, 2, 2), + "decoder_layers_per_block": (5, 5, 5, 5), + "spatio_temporal_scaling": (True, True, True, True), + "decoder_spatio_temporal_scaling": (True, True, True), + "decoder_inject_noise": (False, False, False, False), + "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"), + "upsample_residual": (True, True, True), + "upsample_factor": (2, 2, 2), + "timestep_conditioning": timestep_conditioning, + "patch_size": 4, + "patch_size_t": 1, + "resnet_norm_eps": 1e-6, + "encoder_causal": True, + "decoder_causal": False, + "encoder_spatial_padding_mode": "zeros", + "decoder_spatial_padding_mode": "reflect", + "spatial_compression_ratio": 32, + "temporal_compression_ratio": 8, + }, + } + rename_dict = LTX_2_0_VIDEO_VAE_RENAME_DICT + special_keys_remap = LTX_2_0_VAE_SPECIAL_KEYS_REMAP + return config, rename_dict, special_keys_remap + + +def convert_ltx2_video_vae( + original_state_dict: dict[str, Any], version: str, timestep_conditioning: bool +) -> dict[str, Any]: + config, rename_dict, special_keys_remap = get_ltx2_video_vae_config(version, timestep_conditioning) + diffusers_config = config["diffusers_config"] + + with init_empty_weights(): + vae = AutoencoderKLLTX2Video.from_config(diffusers_config) + + # Handle official code --> diffusers key remapping via the remap dict + for key in list(original_state_dict.keys()): + new_key = key[:] + for replace_key, rename_key in rename_dict.items(): + new_key = new_key.replace(replace_key, rename_key) + update_state_dict_inplace(original_state_dict, key, new_key) + + # Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in + # special_keys_remap + for key in list(original_state_dict.keys()): + for special_key, handler_fn_inplace in special_keys_remap.items(): + if special_key not in key: + continue + handler_fn_inplace(key, original_state_dict) + + vae.load_state_dict(original_state_dict, strict=True, assign=True) + return vae + + +def get_ltx2_audio_vae_config(version: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + if version == "2.0": + config = { + "model_id": "diffusers-internal-dev/new-ltx-model", + "diffusers_config": { + "base_channels": 128, + "output_channels": 2, + "ch_mult": (1, 2, 4), + "num_res_blocks": 2, + "attn_resolutions": None, + "in_channels": 2, + "resolution": 256, + "latent_channels": 8, + "norm_type": "pixel", + "causality_axis": "height", + "dropout": 0.0, + "mid_block_add_attention": False, + "sample_rate": 16000, + "mel_hop_length": 160, + "is_causal": True, + "mel_bins": 64, + "double_z": True, + }, + } + rename_dict = LTX_2_0_AUDIO_VAE_RENAME_DICT + special_keys_remap = LTX_2_0_AUDIO_VAE_SPECIAL_KEYS_REMAP + return config, rename_dict, special_keys_remap + + +def convert_ltx2_audio_vae(original_state_dict: dict[str, Any], version: str) -> dict[str, Any]: + config, rename_dict, special_keys_remap = get_ltx2_audio_vae_config(version) + diffusers_config = config["diffusers_config"] + + with init_empty_weights(): + vae = AutoencoderKLLTX2Audio.from_config(diffusers_config) + + # Handle official code --> diffusers key remapping via the remap dict + for key in list(original_state_dict.keys()): + new_key = key[:] + for replace_key, rename_key in rename_dict.items(): + new_key = new_key.replace(replace_key, rename_key) + update_state_dict_inplace(original_state_dict, key, new_key) + + # Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in + # special_keys_remap + for key in list(original_state_dict.keys()): + for special_key, handler_fn_inplace in special_keys_remap.items(): + if special_key not in key: + continue + handler_fn_inplace(key, original_state_dict) + + vae.load_state_dict(original_state_dict, strict=True, assign=True) + return vae + + +def get_ltx2_vocoder_config(version: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + if version == "2.0": + config = { + "model_id": "diffusers-internal-dev/new-ltx-model", + "diffusers_config": { + "in_channels": 128, + "hidden_channels": 1024, + "out_channels": 2, + "upsample_kernel_sizes": [16, 15, 8, 4, 4], + "upsample_factors": [6, 5, 2, 2, 2], + "resnet_kernel_sizes": [3, 7, 11], + "resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + "leaky_relu_negative_slope": 0.1, + "output_sampling_rate": 24000, + }, + } + rename_dict = LTX_2_0_VOCODER_RENAME_DICT + special_keys_remap = LTX_2_0_VOCODER_SPECIAL_KEYS_REMAP + return config, rename_dict, special_keys_remap + + +def convert_ltx2_vocoder(original_state_dict: dict[str, Any], version: str) -> dict[str, Any]: + config, rename_dict, special_keys_remap = get_ltx2_vocoder_config(version) + diffusers_config = config["diffusers_config"] + + with init_empty_weights(): + vocoder = LTX2Vocoder.from_config(diffusers_config) + + # Handle official code --> diffusers key remapping via the remap dict + for key in list(original_state_dict.keys()): + new_key = key[:] + for replace_key, rename_key in rename_dict.items(): + new_key = new_key.replace(replace_key, rename_key) + update_state_dict_inplace(original_state_dict, key, new_key) + + # Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in + # special_keys_remap + for key in list(original_state_dict.keys()): + for special_key, handler_fn_inplace in special_keys_remap.items(): + if special_key not in key: + continue + handler_fn_inplace(key, original_state_dict) + + vocoder.load_state_dict(original_state_dict, strict=True, assign=True) + return vocoder + + +def get_ltx2_spatial_latent_upsampler_config(version: str): + if version == "2.0": + config = { + "in_channels": 128, + "mid_channels": 1024, + "num_blocks_per_stage": 4, + "dims": 3, + "spatial_upsample": True, + "temporal_upsample": False, + "rational_spatial_scale": 2.0, + } + else: + raise ValueError(f"Unsupported version: {version}") + return config + + +def convert_ltx2_spatial_latent_upsampler( + original_state_dict: dict[str, Any], config: dict[str, Any], dtype: torch.dtype +): + with init_empty_weights(): + latent_upsampler = LTX2LatentUpsamplerModel(**config) + + latent_upsampler.load_state_dict(original_state_dict, strict=True, assign=True) + latent_upsampler.to(dtype) + return latent_upsampler + + +def load_original_checkpoint(args, filename: str | None) -> dict[str, Any]: + if args.original_state_dict_repo_id is not None: + ckpt_path = hf_hub_download(repo_id=args.original_state_dict_repo_id, filename=filename) + elif args.checkpoint_path is not None: + ckpt_path = args.checkpoint_path + else: + raise ValueError("Please provide either `original_state_dict_repo_id` or a local `checkpoint_path`") + + original_state_dict = safetensors.torch.load_file(ckpt_path) + return original_state_dict + + +def load_hub_or_local_checkpoint(repo_id: str | None = None, filename: str | None = None) -> dict[str, Any]: + if repo_id is None and filename is None: + raise ValueError("Please supply at least one of `repo_id` or `filename`") + + if repo_id is not None: + if filename is None: + raise ValueError("If repo_id is specified, filename must also be specified.") + ckpt_path = hf_hub_download(repo_id=repo_id, filename=filename) + else: + ckpt_path = filename + + _, ext = os.path.splitext(ckpt_path) + if ext in [".safetensors", ".sft"]: + state_dict = safetensors.torch.load_file(ckpt_path) + else: + state_dict = torch.load(ckpt_path, map_location="cpu") + + return state_dict + + +def get_model_state_dict_from_combined_ckpt(combined_ckpt: dict[str, Any], prefix: str) -> dict[str, Any]: + # Ensure that the key prefix ends with a dot (.) + if not prefix.endswith("."): + prefix = prefix + "." + + model_state_dict = {} + for param_name, param in combined_ckpt.items(): + if param_name.startswith(prefix): + model_state_dict[param_name.replace(prefix, "")] = param + + if prefix == "model.diffusion_model.": + # Some checkpoints store the text connector projection outside the diffusion model prefix. + connector_key = "text_embedding_projection.aggregate_embed.weight" + if connector_key in combined_ckpt and connector_key not in model_state_dict: + model_state_dict[connector_key] = combined_ckpt[connector_key] + + return model_state_dict + + +def get_args(): + parser = argparse.ArgumentParser() + + def none_or_str(value: str): + if isinstance(value, str) and value.lower() == "none": + return None + return value + + parser.add_argument( + "--original_state_dict_repo_id", + default="Lightricks/LTX-2", + type=none_or_str, + help="HF Hub repo id with LTX 2.0 checkpoint", + ) + parser.add_argument( + "--checkpoint_path", + default=None, + type=str, + help="Local checkpoint path for LTX 2.0. Will be used if `original_state_dict_repo_id` is not specified.", + ) + parser.add_argument( + "--version", + type=str, + default="2.0", + choices=["test", "2.0"], + help="Version of the LTX 2.0 model", + ) + + parser.add_argument( + "--combined_filename", + default="ltx-2-19b-dev.safetensors", + type=none_or_str, + help="Filename for combined checkpoint with all LTX 2.0 models (VAE, DiT, etc.)", + ) + parser.add_argument("--vae_prefix", default="vae.", type=str) + parser.add_argument("--audio_vae_prefix", default="audio_vae.", type=str) + parser.add_argument("--dit_prefix", default="model.diffusion_model.", type=str) + parser.add_argument("--vocoder_prefix", default="vocoder.", type=str) + + parser.add_argument("--vae_filename", default=None, type=str, help="VAE filename; overrides combined ckpt if set") + parser.add_argument( + "--audio_vae_filename", default=None, type=str, help="Audio VAE filename; overrides combined ckpt if set" + ) + parser.add_argument("--dit_filename", default=None, type=str, help="DiT filename; overrides combined ckpt if set") + parser.add_argument( + "--vocoder_filename", default=None, type=str, help="Vocoder filename; overrides combined ckpt if set" + ) + parser.add_argument( + "--text_encoder_model_id", + default="google/gemma-3-12b-it-qat-q4_0-unquantized", + type=none_or_str, + help="HF Hub id for the LTX 2.0 base text encoder model", + ) + parser.add_argument( + "--tokenizer_id", + default="google/gemma-3-12b-it-qat-q4_0-unquantized", + type=none_or_str, + help="HF Hub id for the LTX 2.0 text tokenizer", + ) + parser.add_argument( + "--latent_upsampler_filename", + default="ltx-2-spatial-upscaler-x2-1.0.safetensors", + type=none_or_str, + help="Latent upsampler filename", + ) + + parser.add_argument( + "--timestep_conditioning", action="store_true", help="Whether to add timestep condition to the video VAE model" + ) + parser.add_argument("--vae", action="store_true", help="Whether to convert the video VAE model") + parser.add_argument("--audio_vae", action="store_true", help="Whether to convert the audio VAE model") + parser.add_argument("--dit", action="store_true", help="Whether to convert the DiT model") + parser.add_argument("--connectors", action="store_true", help="Whether to convert the connector model") + parser.add_argument("--vocoder", action="store_true", help="Whether to convert the vocoder model") + parser.add_argument("--text_encoder", action="store_true", help="Whether to conver the text encoder") + parser.add_argument("--latent_upsampler", action="store_true", help="Whether to convert the latent upsampler") + parser.add_argument( + "--full_pipeline", + action="store_true", + help="Whether to save the pipeline. This will attempt to convert all models (e.g. vae, dit, etc.)", + ) + parser.add_argument( + "--upsample_pipeline", + action="store_true", + help="Whether to save a latent upsampling pipeline", + ) + + parser.add_argument("--vae_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"]) + parser.add_argument("--audio_vae_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"]) + parser.add_argument("--dit_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"]) + parser.add_argument("--vocoder_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"]) + parser.add_argument("--text_encoder_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"]) + + parser.add_argument("--output_path", type=str, required=True, help="Path where converted model should be saved") + + return parser.parse_args() + + +DTYPE_MAPPING = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + +VARIANT_MAPPING = { + "fp32": None, + "fp16": "fp16", + "bf16": "bf16", +} + + +def main(args): + vae_dtype = DTYPE_MAPPING[args.vae_dtype] + audio_vae_dtype = DTYPE_MAPPING[args.audio_vae_dtype] + dit_dtype = DTYPE_MAPPING[args.dit_dtype] + vocoder_dtype = DTYPE_MAPPING[args.vocoder_dtype] + text_encoder_dtype = DTYPE_MAPPING[args.text_encoder_dtype] + + combined_ckpt = None + load_combined_models = any( + [ + args.vae, + args.audio_vae, + args.dit, + args.vocoder, + args.text_encoder, + args.full_pipeline, + args.upsample_pipeline, + ] + ) + if args.combined_filename is not None and load_combined_models: + combined_ckpt = load_original_checkpoint(args, filename=args.combined_filename) + + if args.vae or args.full_pipeline or args.upsample_pipeline: + if args.vae_filename is not None: + original_vae_ckpt = load_hub_or_local_checkpoint(filename=args.vae_filename) + elif combined_ckpt is not None: + original_vae_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.vae_prefix) + vae = convert_ltx2_video_vae( + original_vae_ckpt, version=args.version, timestep_conditioning=args.timestep_conditioning + ) + if not args.full_pipeline and not args.upsample_pipeline: + vae.to(vae_dtype).save_pretrained(os.path.join(args.output_path, "vae")) + + if args.audio_vae or args.full_pipeline: + if args.audio_vae_filename is not None: + original_audio_vae_ckpt = load_hub_or_local_checkpoint(filename=args.audio_vae_filename) + elif combined_ckpt is not None: + original_audio_vae_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.audio_vae_prefix) + audio_vae = convert_ltx2_audio_vae(original_audio_vae_ckpt, version=args.version) + if not args.full_pipeline: + audio_vae.to(audio_vae_dtype).save_pretrained(os.path.join(args.output_path, "audio_vae")) + + if args.dit or args.full_pipeline: + if args.dit_filename is not None: + original_dit_ckpt = load_hub_or_local_checkpoint(filename=args.dit_filename) + elif combined_ckpt is not None: + original_dit_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.dit_prefix) + transformer = convert_ltx2_transformer(original_dit_ckpt, version=args.version) + if not args.full_pipeline: + transformer.to(dit_dtype).save_pretrained(os.path.join(args.output_path, "transformer")) + + if args.connectors or args.full_pipeline: + if args.dit_filename is not None: + original_connectors_ckpt = load_hub_or_local_checkpoint(filename=args.dit_filename) + elif combined_ckpt is not None: + original_connectors_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.dit_prefix) + connectors = convert_ltx2_connectors(original_connectors_ckpt, version=args.version) + if not args.full_pipeline: + connectors.to(dit_dtype).save_pretrained(os.path.join(args.output_path, "connectors")) + + if args.vocoder or args.full_pipeline: + if args.vocoder_filename is not None: + original_vocoder_ckpt = load_hub_or_local_checkpoint(filename=args.vocoder_filename) + elif combined_ckpt is not None: + original_vocoder_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.vocoder_prefix) + vocoder = convert_ltx2_vocoder(original_vocoder_ckpt, version=args.version) + if not args.full_pipeline: + vocoder.to(vocoder_dtype).save_pretrained(os.path.join(args.output_path, "vocoder")) + + if args.text_encoder or args.full_pipeline: + # text_encoder = AutoModel.from_pretrained(args.text_encoder_model_id) + text_encoder = Gemma3ForConditionalGeneration.from_pretrained(args.text_encoder_model_id) + if not args.full_pipeline: + text_encoder.to(text_encoder_dtype).save_pretrained(os.path.join(args.output_path, "text_encoder")) + + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_id) + if not args.full_pipeline: + tokenizer.save_pretrained(os.path.join(args.output_path, "tokenizer")) + + if args.latent_upsampler or args.full_pipeline or args.upsample_pipeline: + original_latent_upsampler_ckpt = load_hub_or_local_checkpoint( + repo_id=args.original_state_dict_repo_id, filename=args.latent_upsampler_filename + ) + latent_upsampler_config = get_ltx2_spatial_latent_upsampler_config(args.version) + latent_upsampler = convert_ltx2_spatial_latent_upsampler( + original_latent_upsampler_ckpt, + latent_upsampler_config, + dtype=vae_dtype, + ) + if not args.full_pipeline and not args.upsample_pipeline: + latent_upsampler.save_pretrained(os.path.join(args.output_path, "latent_upsampler")) + + if args.full_pipeline: + scheduler = FlowMatchEulerDiscreteScheduler( + use_dynamic_shifting=True, + base_shift=0.95, + max_shift=2.05, + base_image_seq_len=1024, + max_image_seq_len=4096, + shift_terminal=0.1, + ) + + pipe = LTX2Pipeline( + scheduler=scheduler, + vae=vae, + audio_vae=audio_vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + connectors=connectors, + transformer=transformer, + vocoder=vocoder, + ) + + pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + + if args.upsample_pipeline: + pipe = LTX2LatentUpsamplePipeline(vae=vae, latent_upsampler=latent_upsampler) + + # Put latent upsampling pipeline in its own subdirectory so it doesn't mess with the full pipeline + pipe.save_pretrained( + os.path.join(args.output_path, "upsample_pipeline"), safe_serialization=True, max_shard_size="5GB" + ) + + +if __name__ == "__main__": + args = get_args() + main(args) diff --git a/scripts/convert_ltx_to_diffusers.py b/scripts/convert_ltx_to_diffusers.py index 19e5602039e5..b4ae57f9c8f0 100644 --- a/scripts/convert_ltx_to_diffusers.py +++ b/scripts/convert_ltx_to_diffusers.py @@ -116,7 +116,7 @@ def remove_keys_(key: str, state_dict: Dict[str, Any]): } -def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]: +def get_state_dict(saved_dict: Dict[str, Any]) -> dict[str, Any]: state_dict = saved_dict if "model" in saved_dict.keys(): state_dict = state_dict["model"] @@ -127,7 +127,7 @@ def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]: return state_dict -def update_state_dict_inplace(state_dict: Dict[str, Any], old_key: str, new_key: str) -> Dict[str, Any]: +def update_state_dict_inplace(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]: state_dict[new_key] = state_dict.pop(old_key) @@ -192,7 +192,7 @@ def convert_spatial_latent_upsampler(ckpt_path: str, config, dtype: torch.dtype) return latent_upsampler -def get_transformer_config(version: str) -> Dict[str, Any]: +def get_transformer_config(version: str) -> dict[str, Any]: if version == "0.9.7": config = { "in_channels": 128, @@ -232,7 +232,7 @@ def get_transformer_config(version: str) -> Dict[str, Any]: return config -def get_vae_config(version: str) -> Dict[str, Any]: +def get_vae_config(version: str) -> dict[str, Any]: if version in ["0.9.0"]: config = { "in_channels": 3, @@ -359,7 +359,7 @@ def get_vae_config(version: str) -> Dict[str, Any]: return config -def get_spatial_latent_upsampler_config(version: str) -> Dict[str, Any]: +def get_spatial_latent_upsampler_config(version: str) -> dict[str, Any]: if version == "0.9.7": config = { "in_channels": 128, diff --git a/scripts/convert_prx_to_diffusers.py b/scripts/convert_prx_to_diffusers.py index d9bde2f34d56..00bb3f6fe99e 100644 --- a/scripts/convert_prx_to_diffusers.py +++ b/scripts/convert_prx_to_diffusers.py @@ -86,7 +86,7 @@ def create_parameter_mapping(depth: int) -> dict: return mapping -def convert_checkpoint_parameters(old_state_dict: Dict[str, torch.Tensor], depth: int) -> Dict[str, torch.Tensor]: +def convert_checkpoint_parameters(old_state_dict: Dict[str, torch.Tensor], depth: int) -> dict[str, torch.Tensor]: """Convert old checkpoint parameters to new diffusers format.""" print("Converting checkpoint parameters...") diff --git a/scripts/convert_rae_to_diffusers.py b/scripts/convert_rae_to_diffusers.py new file mode 100644 index 000000000000..0502e49ef30c --- /dev/null +++ b/scripts/convert_rae_to_diffusers.py @@ -0,0 +1,406 @@ +import argparse +from pathlib import Path +from typing import Any + +import torch +from huggingface_hub import HfApi, hf_hub_download + +from diffusers import AutoencoderRAE + + +DECODER_CONFIGS = { + "ViTB": { + "decoder_hidden_size": 768, + "decoder_intermediate_size": 3072, + "decoder_num_attention_heads": 12, + "decoder_num_hidden_layers": 12, + }, + "ViTL": { + "decoder_hidden_size": 1024, + "decoder_intermediate_size": 4096, + "decoder_num_attention_heads": 16, + "decoder_num_hidden_layers": 24, + }, + "ViTXL": { + "decoder_hidden_size": 1152, + "decoder_intermediate_size": 4096, + "decoder_num_attention_heads": 16, + "decoder_num_hidden_layers": 28, + }, +} + +ENCODER_DEFAULT_NAME_OR_PATH = { + "dinov2": "facebook/dinov2-with-registers-base", + "siglip2": "google/siglip2-base-patch16-256", + "mae": "facebook/vit-mae-base", +} + +ENCODER_HIDDEN_SIZE = { + "dinov2": 768, + "siglip2": 768, + "mae": 768, +} + +ENCODER_PATCH_SIZE = { + "dinov2": 14, + "siglip2": 16, + "mae": 16, +} + +DEFAULT_DECODER_SUBDIR = { + "dinov2": "decoders/dinov2/wReg_base", + "mae": "decoders/mae/base_p16", + "siglip2": "decoders/siglip2/base_p16_i256", +} + +DEFAULT_STATS_SUBDIR = { + "dinov2": "stats/dinov2/wReg_base", + "mae": "stats/mae/base_p16", + "siglip2": "stats/siglip2/base_p16_i256", +} + +DECODER_FILE_CANDIDATES = ("dinov2_decoder.pt", "model.pt") +STATS_FILE_CANDIDATES = ("stat.pt",) + + +def dataset_case_candidates(name: str) -> tuple[str, ...]: + return (name, name.lower(), name.upper(), name.title(), "imagenet1k", "ImageNet1k") + + +class RepoAccessor: + def __init__(self, repo_or_path: str, cache_dir: str | None = None): + self.repo_or_path = repo_or_path + self.cache_dir = cache_dir + self.local_root: Path | None = None + self.repo_id: str | None = None + self.repo_files: set[str] | None = None + + root = Path(repo_or_path) + if root.exists() and root.is_dir(): + self.local_root = root + else: + self.repo_id = repo_or_path + self.repo_files = set(HfApi().list_repo_files(repo_or_path)) + + def exists(self, relative_path: str) -> bool: + relative_path = relative_path.replace("\\", "/") + if self.local_root is not None: + return (self.local_root / relative_path).is_file() + return relative_path in self.repo_files + + def fetch(self, relative_path: str) -> Path: + relative_path = relative_path.replace("\\", "/") + if self.local_root is not None: + return self.local_root / relative_path + downloaded = hf_hub_download(repo_id=self.repo_id, filename=relative_path, cache_dir=self.cache_dir) + return Path(downloaded) + + +def unwrap_state_dict(maybe_wrapped: dict[str, Any]) -> dict[str, Any]: + state_dict = maybe_wrapped + for k in ("model", "module", "state_dict"): + if isinstance(state_dict, dict) and k in state_dict and isinstance(state_dict[k], dict): + state_dict = state_dict[k] + + out = dict(state_dict) + if len(out) > 0 and all(key.startswith("module.") for key in out): + out = {key[len("module.") :]: value for key, value in out.items()} + if len(out) > 0 and all(key.startswith("decoder.") for key in out): + out = {key[len("decoder.") :]: value for key, value in out.items()} + return out + + +def remap_decoder_attention_keys_for_diffusers(state_dict: dict[str, Any]) -> dict[str, Any]: + """ + Map official RAE decoder attention key layout to diffusers Attention layout used by AutoencoderRAE decoder. + + Example mappings: + - `...attention.attention.query.*` -> `...attention.to_q.*` + - `...attention.attention.key.*` -> `...attention.to_k.*` + - `...attention.attention.value.*` -> `...attention.to_v.*` + - `...attention.output.dense.*` -> `...attention.to_out.0.*` + """ + remapped: dict[str, Any] = {} + for key, value in state_dict.items(): + new_key = key + new_key = new_key.replace(".attention.attention.query.", ".attention.to_q.") + new_key = new_key.replace(".attention.attention.key.", ".attention.to_k.") + new_key = new_key.replace(".attention.attention.value.", ".attention.to_v.") + new_key = new_key.replace(".attention.output.dense.", ".attention.to_out.0.") + remapped[new_key] = value + return remapped + + +def resolve_decoder_file( + accessor: RepoAccessor, encoder_type: str, variant: str, decoder_checkpoint: str | None +) -> str: + if decoder_checkpoint is not None: + if accessor.exists(decoder_checkpoint): + return decoder_checkpoint + raise FileNotFoundError(f"Decoder checkpoint not found: {decoder_checkpoint}") + + base = f"{DEFAULT_DECODER_SUBDIR[encoder_type]}/{variant}" + for name in DECODER_FILE_CANDIDATES: + candidate = f"{base}/{name}" + if accessor.exists(candidate): + return candidate + + raise FileNotFoundError( + f"Could not find decoder checkpoint under `{base}`. Tried: {list(DECODER_FILE_CANDIDATES)}" + ) + + +def resolve_stats_file( + accessor: RepoAccessor, + encoder_type: str, + dataset_name: str, + stats_checkpoint: str | None, +) -> str | None: + if stats_checkpoint is not None: + if accessor.exists(stats_checkpoint): + return stats_checkpoint + raise FileNotFoundError(f"Stats checkpoint not found: {stats_checkpoint}") + + base = DEFAULT_STATS_SUBDIR[encoder_type] + for dataset in dataset_case_candidates(dataset_name): + for name in STATS_FILE_CANDIDATES: + candidate = f"{base}/{dataset}/{name}" + if accessor.exists(candidate): + return candidate + + return None + + +def extract_latent_stats(stats_obj: Any) -> tuple[Any | None, Any | None]: + if not isinstance(stats_obj, dict): + return None, None + + if "latents_mean" in stats_obj or "latents_std" in stats_obj: + return stats_obj.get("latents_mean", None), stats_obj.get("latents_std", None) + + mean = stats_obj.get("mean", None) + var = stats_obj.get("var", None) + if mean is None and var is None: + return None, None + + latents_std = None + if var is not None: + if isinstance(var, torch.Tensor): + latents_std = torch.sqrt(var + 1e-5) + else: + latents_std = torch.sqrt(torch.tensor(var) + 1e-5) + return mean, latents_std + + +def _strip_final_layernorm_affine(state_dict: dict[str, Any], prefix: str = "") -> dict[str, Any]: + """Remove final layernorm weight/bias from encoder state dict. + + RAE uses non-affine layernorm (weight=1, bias=0 is the default identity). + Stripping these keys means the model keeps its default init values, which + is functionally equivalent to setting elementwise_affine=False. + """ + keys_to_strip = {f"{prefix}weight", f"{prefix}bias"} + return {k: v for k, v in state_dict.items() if k not in keys_to_strip} + + +def _load_hf_encoder_state_dict(encoder_type: str, encoder_name_or_path: str) -> dict[str, Any]: + """Download the HF encoder and extract the state dict for the inner model.""" + if encoder_type == "dinov2": + from transformers import Dinov2WithRegistersModel + + hf_model = Dinov2WithRegistersModel.from_pretrained(encoder_name_or_path) + sd = hf_model.state_dict() + return _strip_final_layernorm_affine(sd, prefix="layernorm.") + elif encoder_type == "siglip2": + from transformers import SiglipModel + + # SiglipModel.vision_model is a SiglipVisionTransformer. + # Our Siglip2Encoder wraps it inside SiglipVisionModel which nests it + # under .vision_model, so we add the prefix to match the diffusers key layout. + hf_model = SiglipModel.from_pretrained(encoder_name_or_path).vision_model + sd = {f"vision_model.{k}": v for k, v in hf_model.state_dict().items()} + return _strip_final_layernorm_affine(sd, prefix="vision_model.post_layernorm.") + elif encoder_type == "mae": + from transformers import ViTMAEForPreTraining + + hf_model = ViTMAEForPreTraining.from_pretrained(encoder_name_or_path).vit + sd = hf_model.state_dict() + return _strip_final_layernorm_affine(sd, prefix="layernorm.") + else: + raise ValueError(f"Unknown encoder_type: {encoder_type}") + + +def convert(args: argparse.Namespace) -> None: + accessor = RepoAccessor(args.repo_or_path, cache_dir=args.cache_dir) + encoder_name_or_path = args.encoder_name_or_path or ENCODER_DEFAULT_NAME_OR_PATH[args.encoder_type] + + decoder_relpath = resolve_decoder_file(accessor, args.encoder_type, args.variant, args.decoder_checkpoint) + stats_relpath = resolve_stats_file(accessor, args.encoder_type, args.dataset_name, args.stats_checkpoint) + + print(f"Using decoder checkpoint: {decoder_relpath}") + if stats_relpath is not None: + print(f"Using stats checkpoint: {stats_relpath}") + else: + print("No stats checkpoint found; conversion will proceed without latent stats.") + + if args.dry_run: + return + + decoder_path = accessor.fetch(decoder_relpath) + decoder_obj = torch.load(decoder_path, map_location="cpu") + decoder_state_dict = unwrap_state_dict(decoder_obj) + decoder_state_dict = remap_decoder_attention_keys_for_diffusers(decoder_state_dict) + + latents_mean, latents_std = None, None + if stats_relpath is not None: + stats_path = accessor.fetch(stats_relpath) + stats_obj = torch.load(stats_path, map_location="cpu") + latents_mean, latents_std = extract_latent_stats(stats_obj) + + decoder_cfg = DECODER_CONFIGS[args.decoder_config_name] + + # Read encoder normalization stats from the HF image processor (only place that downloads encoder info) + from transformers import AutoConfig, AutoImageProcessor + + proc = AutoImageProcessor.from_pretrained(encoder_name_or_path) + encoder_norm_mean = list(proc.image_mean) + encoder_norm_std = list(proc.image_std) + + # Read encoder hidden size and patch size from HF config + encoder_hidden_size = ENCODER_HIDDEN_SIZE[args.encoder_type] + encoder_patch_size = ENCODER_PATCH_SIZE[args.encoder_type] + try: + hf_config = AutoConfig.from_pretrained(encoder_name_or_path) + # For models like SigLIP that nest vision config + if hasattr(hf_config, "vision_config"): + hf_config = hf_config.vision_config + encoder_hidden_size = hf_config.hidden_size + encoder_patch_size = hf_config.patch_size + except Exception: + pass + + # Load the actual encoder weights from HF to include in the saved model + encoder_state_dict = _load_hf_encoder_state_dict(args.encoder_type, encoder_name_or_path) + + # Build model on meta device to avoid double init overhead + with torch.device("meta"): + model = AutoencoderRAE( + encoder_type=args.encoder_type, + encoder_hidden_size=encoder_hidden_size, + encoder_patch_size=encoder_patch_size, + encoder_input_size=args.encoder_input_size, + patch_size=args.patch_size, + image_size=args.image_size, + num_channels=args.num_channels, + encoder_norm_mean=encoder_norm_mean, + encoder_norm_std=encoder_norm_std, + decoder_hidden_size=decoder_cfg["decoder_hidden_size"], + decoder_num_hidden_layers=decoder_cfg["decoder_num_hidden_layers"], + decoder_num_attention_heads=decoder_cfg["decoder_num_attention_heads"], + decoder_intermediate_size=decoder_cfg["decoder_intermediate_size"], + latents_mean=latents_mean, + latents_std=latents_std, + scaling_factor=args.scaling_factor, + ) + + # Assemble full state dict and load with assign=True + full_state_dict = {} + + # Encoder weights (prefixed with "encoder.") + for k, v in encoder_state_dict.items(): + full_state_dict[f"encoder.{k}"] = v + + # Decoder weights (prefixed with "decoder.") + for k, v in decoder_state_dict.items(): + full_state_dict[f"decoder.{k}"] = v + + # Buffers from config + full_state_dict["encoder_mean"] = torch.tensor(encoder_norm_mean, dtype=torch.float32).view(1, 3, 1, 1) + full_state_dict["encoder_std"] = torch.tensor(encoder_norm_std, dtype=torch.float32).view(1, 3, 1, 1) + if latents_mean is not None: + latents_mean_t = latents_mean if isinstance(latents_mean, torch.Tensor) else torch.tensor(latents_mean) + full_state_dict["_latents_mean"] = latents_mean_t + else: + full_state_dict["_latents_mean"] = torch.zeros(1) + if latents_std is not None: + latents_std_t = latents_std if isinstance(latents_std, torch.Tensor) else torch.tensor(latents_std) + full_state_dict["_latents_std"] = latents_std_t + else: + full_state_dict["_latents_std"] = torch.ones(1) + + model.load_state_dict(full_state_dict, strict=False, assign=True) + + # Verify no critical keys are missing + model_keys = {name for name, _ in model.named_parameters()} + model_keys |= {name for name, _ in model.named_buffers()} + loaded_keys = set(full_state_dict.keys()) + missing = model_keys - loaded_keys + # decoder_pos_embed is initialized in-model. trainable_cls_token is only + # allowed to be missing if it was absent in the source decoder checkpoint. + allowed_missing = {"decoder.decoder_pos_embed"} + if "trainable_cls_token" not in decoder_state_dict: + allowed_missing.add("decoder.trainable_cls_token") + if missing - allowed_missing: + print(f"Warning: missing keys after conversion: {sorted(missing - allowed_missing)}") + + output_path = Path(args.output_path) + output_path.mkdir(parents=True, exist_ok=True) + model.save_pretrained(output_path) + + if args.verify_load: + print("Verifying converted checkpoint with AutoencoderRAE.from_pretrained(low_cpu_mem_usage=False)...") + loaded_model = AutoencoderRAE.from_pretrained(output_path, low_cpu_mem_usage=False) + if not isinstance(loaded_model, AutoencoderRAE): + raise RuntimeError("Verification failed: loaded object is not AutoencoderRAE.") + print("Verification passed.") + + print(f"Saved converted AutoencoderRAE to: {output_path}") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Convert RAE decoder checkpoints to diffusers AutoencoderRAE format") + parser.add_argument( + "--repo_or_path", type=str, required=True, help="Hub repo id (e.g. nyu-visionx/RAE-collections) or local path" + ) + parser.add_argument("--output_path", type=str, required=True, help="Directory to save converted model") + + parser.add_argument("--encoder_type", type=str, choices=["dinov2", "mae", "siglip2"], required=True) + parser.add_argument( + "--encoder_name_or_path", type=str, default=None, help="Optional encoder HF model id or local path override" + ) + + parser.add_argument("--variant", type=str, default="ViTXL_n08", help="Decoder variant folder name") + parser.add_argument("--dataset_name", type=str, default="imagenet1k", help="Stats dataset folder name") + + parser.add_argument( + "--decoder_checkpoint", type=str, default=None, help="Relative path to decoder checkpoint inside repo/path" + ) + parser.add_argument( + "--stats_checkpoint", type=str, default=None, help="Relative path to stats checkpoint inside repo/path" + ) + + parser.add_argument("--decoder_config_name", type=str, choices=list(DECODER_CONFIGS.keys()), default="ViTXL") + parser.add_argument("--encoder_input_size", type=int, default=224) + parser.add_argument("--patch_size", type=int, default=16) + parser.add_argument("--image_size", type=int, default=None) + parser.add_argument("--num_channels", type=int, default=3) + parser.add_argument("--scaling_factor", type=float, default=1.0) + + parser.add_argument("--cache_dir", type=str, default=None) + parser.add_argument("--dry_run", action="store_true", help="Only resolve and print selected files") + parser.add_argument( + "--verify_load", + action="store_true", + help="After conversion, load back with AutoencoderRAE.from_pretrained(low_cpu_mem_usage=False).", + ) + + return parser.parse_args() + + +def main() -> None: + args = parse_args() + convert(args) + + +if __name__ == "__main__": + main() diff --git a/scripts/convert_skyreelsv2_to_diffusers.py b/scripts/convert_skyreelsv2_to_diffusers.py index 3bc3c435685b..29d5a7666a6e 100644 --- a/scripts/convert_skyreelsv2_to_diffusers.py +++ b/scripts/convert_skyreelsv2_to_diffusers.py @@ -66,7 +66,7 @@ TRANSFORMER_SPECIAL_KEYS_REMAP = {} -def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> Dict[str, Any]: +def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]: state_dict[new_key] = state_dict.pop(old_key) @@ -81,7 +81,7 @@ def load_sharded_safetensors(dir: pathlib.Path): return state_dict -def get_transformer_config(model_type: str) -> Dict[str, Any]: +def get_transformer_config(model_type: str) -> dict[str, Any]: if model_type == "SkyReels-V2-DF-1.3B-540P": config = { "model_id": "Skywork/SkyReels-V2-DF-1.3B-540P", diff --git a/scripts/convert_wan_to_diffusers.py b/scripts/convert_wan_to_diffusers.py index 06f87409262a..eca1af61ae69 100644 --- a/scripts/convert_wan_to_diffusers.py +++ b/scripts/convert_wan_to_diffusers.py @@ -313,7 +313,7 @@ def convert_animate_face_adapter_weights(key: str, state_dict: Dict[str, Any]) - } -def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> Dict[str, Any]: +def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]: state_dict[new_key] = state_dict.pop(old_key) diff --git a/setup.py b/setup.py index c47124554479..d42da57920a0 100644 --- a/setup.py +++ b/setup.py @@ -101,6 +101,7 @@ "datasets", "filelock", "flax>=0.4.1", + "ftfy", "hf-doc-builder>=0.3.0", "httpx<1.0.0", "huggingface-hub>=0.34.0,<2.0", @@ -111,7 +112,6 @@ "jax>=0.4.1", "jaxlib>=0.4.1", "Jinja2", - "k-diffusion==0.0.12", "torchsde", "note_seq", "librosa", @@ -122,7 +122,7 @@ "pytest", "pytest-timeout", "pytest-xdist", - "python>=3.8.0", + "python>=3.10.0", "ruff==0.9.10", "safetensors>=0.3.1", "sentencepiece>=0.1.91,!=0.1.92", @@ -222,13 +222,14 @@ def run(self): extras["training"] = deps_list("accelerate", "datasets", "protobuf", "tensorboard", "Jinja2", "peft", "timm") extras["test"] = deps_list( "compel", + "ftfy", "GitPython", "datasets", "Jinja2", "invisible-watermark", - "k-diffusion", "librosa", "parameterized", + "protobuf", "pytest", "pytest-timeout", "pytest-xdist", @@ -237,6 +238,7 @@ def run(self): "sentencepiece", "scipy", "tiktoken", + "torchsde", "torchvision", "transformers", "phonemizer", @@ -274,7 +276,7 @@ def run(self): setup( name="diffusers", - version="0.36.0.dev0", # expected format is one of x.y.z.dev0, or x.y.z.rc1 or x.y.z (no to dashes, yes to dots) + version="0.38.0.dev0", # expected format is one of x.y.z.dev0, or x.y.z.rc1 or x.y.z (no to dashes, yes to dots) description="State-of-the-art diffusion in PyTorch and JAX.", long_description=open("README.md", "r", encoding="utf-8").read(), long_description_content_type="text/markdown", @@ -287,7 +289,7 @@ def run(self): packages=find_packages("src"), package_data={"diffusers": ["py.typed"]}, include_package_data=True, - python_requires=">=3.8.0", + python_requires=">=3.10.0", install_requires=list(install_requires), extras_require=extras, entry_points={"console_scripts": ["diffusers-cli=diffusers.commands.diffusers_cli:main"]}, diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index aa11a741af38..cbd2d237f49d 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -1,4 +1,4 @@ -__version__ = "0.37.0.dev0" +__version__ = "0.38.0.dev0" from typing import TYPE_CHECKING @@ -10,7 +10,6 @@ is_bitsandbytes_available, is_flax_available, is_gguf_available, - is_k_diffusion_available, is_librosa_available, is_note_seq_available, is_nvidia_modelopt_available, @@ -23,6 +22,7 @@ is_torchao_available, is_torchsde_available, is_transformers_available, + is_transformers_version, ) @@ -49,8 +49,6 @@ "is_flax_available", "is_inflect_available", "is_invisible_watermark_available", - "is_k_diffusion_available", - "is_k_diffusion_version", "is_librosa_available", "is_note_seq_available", "is_onnx_available", @@ -167,12 +165,14 @@ "FirstBlockCacheConfig", "HookRegistry", "LayerSkipConfig", + "MagCacheConfig", "PyramidAttentionBroadcastConfig", "SmoothedEnergyGuidanceConfig", "TaylorSeerCacheConfig", "apply_faster_cache", "apply_first_block_cache", "apply_layer_skip", + "apply_mag_cache", "apply_pyramid_attention_broadcast", "apply_taylorseer_cache", ] @@ -193,6 +193,8 @@ "AutoencoderKLHunyuanImageRefiner", "AutoencoderKLHunyuanVideo", "AutoencoderKLHunyuanVideo15", + "AutoencoderKLLTX2Audio", + "AutoencoderKLLTX2Video", "AutoencoderKLLTXVideo", "AutoencoderKLMagvit", "AutoencoderKLMochi", @@ -200,12 +202,14 @@ "AutoencoderKLTemporalDecoder", "AutoencoderKLWan", "AutoencoderOobleck", + "AutoencoderRAE", "AutoencoderTiny", "AutoModel", "BriaFiboTransformer2DModel", "BriaTransformer2DModel", "CacheMixin", "ChromaTransformer2DModel", + "ChromaRadianceTransformer2DModel", "ChronoEditTransformer3DModel", "CogVideoXTransformer3DModel", "CogView3PlusTransformer2DModel", @@ -216,6 +220,7 @@ "ControlNetModel", "ControlNetUnionModel", "ControlNetXSAdapter", + "CosmosControlNetModel", "CosmosTransformer3DModel", "DiTTransformer2DModel", "EasyAnimateTransformer3DModel", @@ -223,6 +228,8 @@ "FluxControlNetModel", "FluxMultiControlNetModel", "FluxTransformer2DModel", + "GlmImageTransformer2DModel", + "HeliosTransformer3DModel", "HiDreamImageTransformer2DModel", "HunyuanDiT2DControlNetModel", "HunyuanDiT2DModel", @@ -236,6 +243,7 @@ "Kandinsky5Transformer3DModel", "LatteTransformer3DModel", "LongCatImageTransformer2DModel", + "LTX2VideoTransformer3DModel", "LTXVideoTransformer3DModel", "Lumina2Transformer2DModel", "LuminaNextDiT2DModel", @@ -286,10 +294,17 @@ ) _import_structure["modular_pipelines"].extend( [ + "AutoPipelineBlocks", "ComponentsManager", "ComponentSpec", + "ConditionalPipelineBlocks", + "ConfigSpec", + "InputParam", + "LoopSequentialPipelineBlocks", "ModularPipeline", "ModularPipelineBlocks", + "OutputParam", + "SequentialPipelineBlocks", ] ) _import_structure["optimization"] = [ @@ -347,12 +362,15 @@ "FlowMatchEulerDiscreteScheduler", "FlowMatchHeunDiscreteScheduler", "FlowMatchLCMScheduler", + "HeliosDMDScheduler", + "HeliosScheduler", "HeunDiscreteScheduler", "IPNDMScheduler", "KarrasVeScheduler", "KDPM2AncestralDiscreteScheduler", "KDPM2DiscreteScheduler", "LCMScheduler", + "LTXEulerAncestralRFScheduler", "PNDMScheduler", "RePaintScheduler", "SASolverScheduler", @@ -407,6 +425,10 @@ _import_structure["modular_pipelines"].extend( [ "Flux2AutoBlocks", + "Flux2KleinAutoBlocks", + "Flux2KleinBaseAutoBlocks", + "Flux2KleinBaseModularPipeline", + "Flux2KleinModularPipeline", "Flux2ModularPipeline", "FluxAutoBlocks", "FluxKontextAutoBlocks", @@ -417,11 +439,18 @@ "QwenImageEditModularPipeline", "QwenImageEditPlusAutoBlocks", "QwenImageEditPlusModularPipeline", + "QwenImageLayeredAutoBlocks", + "QwenImageLayeredModularPipeline", "QwenImageModularPipeline", "StableDiffusionXLAutoBlocks", "StableDiffusionXLModularPipeline", - "Wan22AutoBlocks", - "WanAutoBlocks", + "Wan22Blocks", + "Wan22Image2VideoBlocks", + "Wan22Image2VideoModularPipeline", + "Wan22ModularPipeline", + "WanBlocks", + "WanImage2VideoAutoBlocks", + "WanImage2VideoModularPipeline", "WanModularPipeline", "ZImageAutoBlocks", "ZImageModularPipeline", @@ -449,9 +478,12 @@ "AuraFlowPipeline", "BlipDiffusionControlNetPipeline", "BlipDiffusionPipeline", + "BriaFiboEditPipeline", "BriaFiboPipeline", "BriaPipeline", "ChromaImg2ImgPipeline", + "ChromaRadiancePipeline", + "ChromaInpaintPipeline", "ChromaPipeline", "ChronoEditPipeline", "CLIPImageProjection", @@ -464,6 +496,7 @@ "CogView4Pipeline", "ConsisIDPipeline", "Cosmos2_5_PredictBasePipeline", + "Cosmos2_5_TransferPipeline", "Cosmos2TextToImagePipeline", "Cosmos2VideoToWorldPipeline", "CosmosTextToWorldPipeline", @@ -472,6 +505,7 @@ "EasyAnimateControlPipeline", "EasyAnimateInpaintPipeline", "EasyAnimatePipeline", + "Flux2KleinPipeline", "Flux2Pipeline", "FluxControlImg2ImgPipeline", "FluxControlInpaintPipeline", @@ -486,6 +520,9 @@ "FluxKontextPipeline", "FluxPipeline", "FluxPriorReduxPipeline", + "GlmImagePipeline", + "HeliosPipeline", + "HeliosPyramidPipeline", "HiDreamImagePipeline", "HunyuanDiTControlNetPipeline", "HunyuanDiTPAGPipeline", @@ -537,7 +574,12 @@ "LEditsPPPipelineStableDiffusionXL", "LongCatImageEditPipeline", "LongCatImagePipeline", + "LTX2ConditionPipeline", + "LTX2ImageToVideoPipeline", + "LTX2LatentUpsamplePipeline", + "LTX2Pipeline", "LTXConditionPipeline", + "LTXI2VLongMultiPromptPipeline", "LTXImageToVideoPipeline", "LTXLatentUpsamplePipeline", "LTXPipeline", @@ -675,6 +717,7 @@ "ZImageControlNetInpaintPipeline", "ZImageControlNetPipeline", "ZImageImg2ImgPipeline", + "ZImageInpaintPipeline", "ZImageOmniPipeline", "ZImagePipeline", ] @@ -694,19 +737,6 @@ else: _import_structure["pipelines"].extend(["ConsisIDPipeline"]) -try: - if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from .utils import dummy_torch_and_transformers_and_k_diffusion_objects # noqa F403 - - _import_structure["utils.dummy_torch_and_transformers_and_k_diffusion_objects"] = [ - name for name in dir(dummy_torch_and_transformers_and_k_diffusion_objects) if not name.startswith("_") - ] - -else: - _import_structure["pipelines"].extend(["StableDiffusionKDiffusionPipeline", "StableDiffusionXLKDiffusionPipeline"]) - try: if not (is_torch_available() and is_transformers_available() and is_sentencepiece_available()): raise OptionalDependencyNotAvailable() @@ -913,12 +943,14 @@ FirstBlockCacheConfig, HookRegistry, LayerSkipConfig, + MagCacheConfig, PyramidAttentionBroadcastConfig, SmoothedEnergyGuidanceConfig, TaylorSeerCacheConfig, apply_faster_cache, apply_first_block_cache, apply_layer_skip, + apply_mag_cache, apply_pyramid_attention_broadcast, apply_taylorseer_cache, ) @@ -937,6 +969,8 @@ AutoencoderKLHunyuanImageRefiner, AutoencoderKLHunyuanVideo, AutoencoderKLHunyuanVideo15, + AutoencoderKLLTX2Audio, + AutoencoderKLLTX2Video, AutoencoderKLLTXVideo, AutoencoderKLMagvit, AutoencoderKLMochi, @@ -944,11 +978,13 @@ AutoencoderKLTemporalDecoder, AutoencoderKLWan, AutoencoderOobleck, + AutoencoderRAE, AutoencoderTiny, AutoModel, BriaFiboTransformer2DModel, BriaTransformer2DModel, CacheMixin, + ChromaRadianceTransformer2DModel, ChromaTransformer2DModel, ChronoEditTransformer3DModel, CogVideoXTransformer3DModel, @@ -960,6 +996,7 @@ ControlNetModel, ControlNetUnionModel, ControlNetXSAdapter, + CosmosControlNetModel, CosmosTransformer3DModel, DiTTransformer2DModel, EasyAnimateTransformer3DModel, @@ -967,6 +1004,8 @@ FluxControlNetModel, FluxMultiControlNetModel, FluxTransformer2DModel, + GlmImageTransformer2DModel, + HeliosTransformer3DModel, HiDreamImageTransformer2DModel, HunyuanDiT2DControlNetModel, HunyuanDiT2DModel, @@ -980,6 +1019,7 @@ Kandinsky5Transformer3DModel, LatteTransformer3DModel, LongCatImageTransformer2DModel, + LTX2VideoTransformer3DModel, LTXVideoTransformer3DModel, Lumina2Transformer2DModel, LuminaNextDiT2DModel, @@ -1026,7 +1066,19 @@ ZImageTransformer2DModel, attention_backend, ) - from .modular_pipelines import ComponentsManager, ComponentSpec, ModularPipeline, ModularPipelineBlocks + from .modular_pipelines import ( + AutoPipelineBlocks, + ComponentsManager, + ComponentSpec, + ConditionalPipelineBlocks, + ConfigSpec, + InputParam, + LoopSequentialPipelineBlocks, + ModularPipeline, + ModularPipelineBlocks, + OutputParam, + SequentialPipelineBlocks, + ) from .optimization import ( get_constant_schedule, get_constant_schedule_with_warmup, @@ -1082,12 +1134,15 @@ FlowMatchEulerDiscreteScheduler, FlowMatchHeunDiscreteScheduler, FlowMatchLCMScheduler, + HeliosDMDScheduler, + HeliosScheduler, HeunDiscreteScheduler, IPNDMScheduler, KarrasVeScheduler, KDPM2AncestralDiscreteScheduler, KDPM2DiscreteScheduler, LCMScheduler, + LTXEulerAncestralRFScheduler, PNDMScheduler, RePaintScheduler, SASolverScheduler, @@ -1125,6 +1180,10 @@ else: from .modular_pipelines import ( Flux2AutoBlocks, + Flux2KleinAutoBlocks, + Flux2KleinBaseAutoBlocks, + Flux2KleinBaseModularPipeline, + Flux2KleinModularPipeline, Flux2ModularPipeline, FluxAutoBlocks, FluxKontextAutoBlocks, @@ -1135,11 +1194,18 @@ QwenImageEditModularPipeline, QwenImageEditPlusAutoBlocks, QwenImageEditPlusModularPipeline, + QwenImageLayeredAutoBlocks, + QwenImageLayeredModularPipeline, QwenImageModularPipeline, StableDiffusionXLAutoBlocks, StableDiffusionXLModularPipeline, - Wan22AutoBlocks, - WanAutoBlocks, + Wan22Blocks, + Wan22Image2VideoBlocks, + Wan22Image2VideoModularPipeline, + Wan22ModularPipeline, + WanBlocks, + WanImage2VideoAutoBlocks, + WanImage2VideoModularPipeline, WanModularPipeline, ZImageAutoBlocks, ZImageModularPipeline, @@ -1163,10 +1229,13 @@ AudioLDM2UNet2DConditionModel, AudioLDMPipeline, AuraFlowPipeline, + BriaFiboEditPipeline, BriaFiboPipeline, BriaPipeline, ChromaImg2ImgPipeline, + ChromaInpaintPipeline, ChromaPipeline, + ChromaRadiancePipeline, ChronoEditPipeline, CLIPImageProjection, CogVideoXFunControlPipeline, @@ -1178,6 +1247,7 @@ CogView4Pipeline, ConsisIDPipeline, Cosmos2_5_PredictBasePipeline, + Cosmos2_5_TransferPipeline, Cosmos2TextToImagePipeline, Cosmos2VideoToWorldPipeline, CosmosTextToWorldPipeline, @@ -1186,6 +1256,7 @@ EasyAnimateControlPipeline, EasyAnimateInpaintPipeline, EasyAnimatePipeline, + Flux2KleinPipeline, Flux2Pipeline, FluxControlImg2ImgPipeline, FluxControlInpaintPipeline, @@ -1200,6 +1271,9 @@ FluxKontextPipeline, FluxPipeline, FluxPriorReduxPipeline, + GlmImagePipeline, + HeliosPipeline, + HeliosPyramidPipeline, HiDreamImagePipeline, HunyuanDiTControlNetPipeline, HunyuanDiTPAGPipeline, @@ -1251,7 +1325,12 @@ LEditsPPPipelineStableDiffusionXL, LongCatImageEditPipeline, LongCatImagePipeline, + LTX2ConditionPipeline, + LTX2ImageToVideoPipeline, + LTX2LatentUpsamplePipeline, + LTX2Pipeline, LTXConditionPipeline, + LTXI2VLongMultiPromptPipeline, LTXImageToVideoPipeline, LTXLatentUpsamplePipeline, LTXPipeline, @@ -1387,18 +1466,11 @@ ZImageControlNetInpaintPipeline, ZImageControlNetPipeline, ZImageImg2ImgPipeline, + ZImageInpaintPipeline, ZImageOmniPipeline, ZImagePipeline, ) - try: - if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import * # noqa F403 - else: - from .pipelines import StableDiffusionKDiffusionPipeline, StableDiffusionXLKDiffusionPipeline - try: if not (is_torch_available() and is_transformers_available() and is_sentencepiece_available()): raise OptionalDependencyNotAvailable() diff --git a/src/diffusers/callbacks.py b/src/diffusers/callbacks.py index 2a08f091d9f3..087a6b7fee56 100644 --- a/src/diffusers/callbacks.py +++ b/src/diffusers/callbacks.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, List +from typing import Any from .configuration_utils import ConfigMixin, register_to_config from .utils import CONFIG_NAME @@ -33,13 +33,13 @@ def __init__(self, cutoff_step_ratio=1.0, cutoff_step_index=None): raise ValueError("cutoff_step_ratio must be a float between 0.0 and 1.0.") @property - def tensor_inputs(self) -> List[str]: + def tensor_inputs(self) -> list[str]: raise NotImplementedError(f"You need to set the attribute `tensor_inputs` for {self.__class__}") - def callback_fn(self, pipeline, step_index, timesteps, callback_kwargs) -> Dict[str, Any]: + def callback_fn(self, pipeline, step_index, timesteps, callback_kwargs) -> dict[str, Any]: raise NotImplementedError(f"You need to implement the method `callback_fn` for {self.__class__}") - def __call__(self, pipeline, step_index, timestep, callback_kwargs) -> Dict[str, Any]: + def __call__(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]: return self.callback_fn(pipeline, step_index, timestep, callback_kwargs) @@ -49,14 +49,14 @@ class MultiPipelineCallbacks: provides a unified interface for calling all of them. """ - def __init__(self, callbacks: List[PipelineCallback]): + def __init__(self, callbacks: list[PipelineCallback]): self.callbacks = callbacks @property - def tensor_inputs(self) -> List[str]: + def tensor_inputs(self) -> list[str]: return [input for callback in self.callbacks for input in callback.tensor_inputs] - def __call__(self, pipeline, step_index, timestep, callback_kwargs) -> Dict[str, Any]: + def __call__(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]: """ Calls all the callbacks in order with the given arguments and returns the final callback_kwargs. """ @@ -76,7 +76,7 @@ class SDCFGCutoffCallback(PipelineCallback): tensor_inputs = ["prompt_embeds"] - def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> Dict[str, Any]: + def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]: cutoff_step_ratio = self.config.cutoff_step_ratio cutoff_step_index = self.config.cutoff_step_index @@ -109,7 +109,7 @@ class SDXLCFGCutoffCallback(PipelineCallback): "add_time_ids", ] - def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> Dict[str, Any]: + def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]: cutoff_step_ratio = self.config.cutoff_step_ratio cutoff_step_index = self.config.cutoff_step_index @@ -152,7 +152,7 @@ class SDXLControlnetCFGCutoffCallback(PipelineCallback): "image", ] - def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> Dict[str, Any]: + def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]: cutoff_step_ratio = self.config.cutoff_step_ratio cutoff_step_index = self.config.cutoff_step_index @@ -195,7 +195,7 @@ class IPAdapterScaleCutoffCallback(PipelineCallback): tensor_inputs = [] - def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> Dict[str, Any]: + def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]: cutoff_step_ratio = self.config.cutoff_step_ratio cutoff_step_index = self.config.cutoff_step_index @@ -219,7 +219,7 @@ class SD3CFGCutoffCallback(PipelineCallback): tensor_inputs = ["prompt_embeds", "pooled_prompt_embeds"] - def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> Dict[str, Any]: + def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]: cutoff_step_ratio = self.config.cutoff_step_ratio cutoff_step_index = self.config.cutoff_step_index diff --git a/src/diffusers/commands/custom_blocks.py b/src/diffusers/commands/custom_blocks.py index 43d9ea88577a..953240c5a2c3 100644 --- a/src/diffusers/commands/custom_blocks.py +++ b/src/diffusers/commands/custom_blocks.py @@ -89,8 +89,6 @@ def run(self): # automap = self._create_automap(parent_class=parent_class, child_class=child_class) # with open(CONFIG, "w") as f: # json.dump(automap, f) - with open("requirements.txt", "w") as f: - f.write("") def _choose_block(self, candidates, chosen=None): for cls, base in candidates: diff --git a/src/diffusers/commands/fp16_safetensors.py b/src/diffusers/commands/fp16_safetensors.py index 41739261e553..382d6c39bd19 100644 --- a/src/diffusers/commands/fp16_safetensors.py +++ b/src/diffusers/commands/fp16_safetensors.py @@ -35,8 +35,8 @@ def conversion_command_factory(args: Namespace): if args.use_auth_token: warnings.warn( - "The `--use_auth_token` flag is deprecated and will be removed in a future version. Authentication is now" - " handled automatically if user is logged in." + "The `--use_auth_token` flag is deprecated and will be removed in a future version." + "Authentication is now handled automatically if the user is logged in." ) return FP16SafetensorsCommand(args.ckpt_id, args.fp16, args.use_safetensors) @@ -92,8 +92,8 @@ def run(self): pipeline_class = getattr(import_module("diffusers"), pipeline_class_name) self.logger.info(f"Pipeline class imported: {pipeline_class_name}.") - # Load the appropriate pipeline. We could have use `DiffusionPipeline` - # here, but just to avoid any rough edge cases. + # Load the appropriate pipeline. We could have used `DiffusionPipeline` + # here, but just to avoid potential edge cases. pipeline = pipeline_class.from_pretrained( self.ckpt_id, torch_dtype=torch.float16 if self.fp16 else torch.float32 ) diff --git a/src/diffusers/configuration_utils.py b/src/diffusers/configuration_utils.py index 1c4ee33acbfd..7a95ce20aaff 100644 --- a/src/diffusers/configuration_utils.py +++ b/src/diffusers/configuration_utils.py @@ -24,7 +24,7 @@ import re from collections import OrderedDict from pathlib import Path -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any import numpy as np from huggingface_hub import DDUFEntry, create_repo, hf_hub_download @@ -94,10 +94,10 @@ class ConfigMixin: Class attributes: - **config_name** (`str`) -- A filename under which the config should stored when calling [`~ConfigMixin.save_config`] (should be overridden by parent class). - - **ignore_for_config** (`List[str]`) -- A list of attributes that should not be saved in the config (should be + - **ignore_for_config** (`list[str]`) -- A list of attributes that should not be saved in the config (should be overridden by subclass). - **has_compatibles** (`bool`) -- Whether the class has compatible classes (should be overridden by subclass). - - **_deprecated_kwargs** (`List[str]`) -- Keyword arguments that are deprecated. Note that the `init` function + - **_deprecated_kwargs** (`list[str]`) -- Keyword arguments that are deprecated. Note that the `init` function should only have a `kwargs` argument if at least one argument is deprecated (should be overridden by subclass). """ @@ -107,6 +107,38 @@ class ConfigMixin: has_compatibles = False _deprecated_kwargs = [] + _auto_class = None + + @classmethod + def register_for_auto_class(cls, auto_class="AutoModel"): + """ + Register this class with the given auto class so that it can be loaded with `AutoModel.from_pretrained(..., + trust_remote_code=True)`. + + When the config is saved, the resulting `config.json` will include an `auto_map` entry mapping the auto class + to this class's module and class name. + + Args: + auto_class (`str` or type, *optional*, defaults to `"AutoModel"`): + The auto class to register this class with. Can be a string (e.g. `"AutoModel"`) or the class itself. + Currently only `"AutoModel"` is supported. + + Example: + + ```python + from diffusers import ModelMixin, ConfigMixin + + + class MyCustomModel(ModelMixin, ConfigMixin): ... + + + MyCustomModel.register_for_auto_class("AutoModel") + ``` + """ + if auto_class != "AutoModel": + raise ValueError(f"Only 'AutoModel' is supported, got '{auto_class}'.") + + cls._auto_class = auto_class def register_to_config(self, **kwargs): if self.config_name is None: @@ -143,7 +175,7 @@ def __getattr__(self, name: str) -> Any: raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") - def save_config(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs): + def save_config(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs): """ Save a configuration object to the directory specified in `save_directory` so that it can be reloaded using the [`~ConfigMixin.from_config`] class method. @@ -155,7 +187,7 @@ def save_config(self, save_directory: Union[str, os.PathLike], push_to_hub: bool Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the repository you want to push to with `repo_id` (will default to the name of `save_directory` in your namespace). - kwargs (`Dict[str, Any]`, *optional*): + kwargs (`dict[str, Any]`, *optional*): Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. """ if os.path.isfile(save_directory): @@ -189,13 +221,13 @@ def save_config(self, save_directory: Union[str, os.PathLike], push_to_hub: bool @classmethod def from_config( - cls, config: Union[FrozenDict, Dict[str, Any]] = None, return_unused_kwargs=False, **kwargs - ) -> Union[Self, Tuple[Self, Dict[str, Any]]]: + cls, config: FrozenDict | dict[str, Any] = None, return_unused_kwargs=False, **kwargs + ) -> Self | tuple[Self, dict[str, Any]]: r""" Instantiate a Python class from a config dictionary. Parameters: - config (`Dict[str, Any]`): + config (`dict[str, Any]`): A config dictionary from which the Python class is instantiated. Make sure to only load configuration files of compatible classes. return_unused_kwargs (`bool`, *optional*, defaults to `False`): @@ -292,11 +324,11 @@ def get_config_dict(cls, *args, **kwargs): @validate_hf_hub_args def load_config( cls, - pretrained_model_name_or_path: Union[str, os.PathLike], + pretrained_model_name_or_path: str | os.PathLike, return_unused_kwargs=False, return_commit_hash=False, **kwargs, - ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + ) -> tuple[dict[str, Any], dict[str, Any]]: r""" Load a model or scheduler configuration. @@ -309,13 +341,13 @@ def load_config( - A path to a *directory* (for example `./my_model_directory`) containing model weights saved with [`~ConfigMixin.save_config`]. - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. output_loading_info(`bool`, *optional*, defaults to `False`): @@ -352,7 +384,7 @@ def load_config( _ = kwargs.pop("mirror", None) subfolder = kwargs.pop("subfolder", None) user_agent = kwargs.pop("user_agent", {}) - dduf_entries: Optional[Dict[str, DDUFEntry]] = kwargs.pop("dduf_entries", None) + dduf_entries: dict[str, DDUFEntry] | None = kwargs.pop("dduf_entries", None) user_agent = {**user_agent, "file_type": "config"} user_agent = http_user_agent(user_agent) @@ -563,9 +595,7 @@ def extract_init_dict(cls, config_dict, **kwargs): return init_dict, unused_kwargs, hidden_config_dict @classmethod - def _dict_from_json_file( - cls, json_file: Union[str, os.PathLike], dduf_entries: Optional[Dict[str, DDUFEntry]] = None - ): + def _dict_from_json_file(cls, json_file: str | os.PathLike, dduf_entries: dict[str, DDUFEntry] | None = None): if dduf_entries: text = dduf_entries[json_file].read_text() else: @@ -577,12 +607,12 @@ def __repr__(self): return f"{self.__class__.__name__} {self.to_json_string()}" @property - def config(self) -> Dict[str, Any]: + def config(self) -> dict[str, Any]: """ Returns the config of the class as a frozen dictionary Returns: - `Dict[str, Any]`: Config of the class. + `dict[str, Any]`: Config of the class. """ return self._internal_dict @@ -623,9 +653,15 @@ def to_json_saveable(value): # pop the `_pre_quantization_dtype` as torch.dtypes are not serializable. _ = config_dict.pop("_pre_quantization_dtype", None) + if getattr(self, "_auto_class", None) is not None: + module = self.__class__.__module__.split(".")[-1] + auto_map = config_dict.get("auto_map", {}) + auto_map[self._auto_class] = f"{module}.{self.__class__.__name__}" + config_dict["auto_map"] = auto_map + return json.dumps(config_dict, indent=2, sort_keys=True) + "\n" - def to_json_file(self, json_file_path: Union[str, os.PathLike]): + def to_json_file(self, json_file_path: str | os.PathLike): """ Save the configuration instance's parameters to a JSON file. @@ -637,7 +673,7 @@ def to_json_file(self, json_file_path: Union[str, os.PathLike]): writer.write(self.to_json_string()) @classmethod - def _get_config_file_from_dduf(cls, pretrained_model_name_or_path: str, dduf_entries: Dict[str, DDUFEntry]): + def _get_config_file_from_dduf(cls, pretrained_model_name_or_path: str, dduf_entries: dict[str, DDUFEntry]): # paths inside a DDUF file must always be "/" config_file = ( cls.config_name @@ -756,7 +792,7 @@ class LegacyConfigMixin(ConfigMixin): """ @classmethod - def from_config(cls, config: Union[FrozenDict, Dict[str, Any]] = None, return_unused_kwargs=False, **kwargs): + def from_config(cls, config: FrozenDict | dict[str, Any] = None, return_unused_kwargs=False, **kwargs): # To prevent dependency import problem. from .models.model_loading_utils import _fetch_remapped_cls_from_config diff --git a/src/diffusers/dependency_versions_table.py b/src/diffusers/dependency_versions_table.py index 6e5ac630ab08..5b3f84213a2e 100644 --- a/src/diffusers/dependency_versions_table.py +++ b/src/diffusers/dependency_versions_table.py @@ -8,6 +8,7 @@ "datasets": "datasets", "filelock": "filelock", "flax": "flax>=0.4.1", + "ftfy": "ftfy", "hf-doc-builder": "hf-doc-builder>=0.3.0", "httpx": "httpx<1.0.0", "huggingface-hub": "huggingface-hub>=0.34.0,<2.0", @@ -18,7 +19,6 @@ "jax": "jax>=0.4.1", "jaxlib": "jaxlib>=0.4.1", "Jinja2": "Jinja2", - "k-diffusion": "k-diffusion==0.0.12", "torchsde": "torchsde", "note_seq": "note_seq", "librosa": "librosa", @@ -29,7 +29,7 @@ "pytest": "pytest", "pytest-timeout": "pytest-timeout", "pytest-xdist": "pytest-xdist", - "python": "python>=3.8.0", + "python": "python>=3.10.0", "ruff": "ruff==0.9.10", "safetensors": "safetensors>=0.3.1", "sentencepiece": "sentencepiece>=0.1.91,!=0.1.92", diff --git a/src/diffusers/guiders/__init__.py b/src/diffusers/guiders/__init__.py index 58ad0c211b64..b6653817dc95 100644 --- a/src/diffusers/guiders/__init__.py +++ b/src/diffusers/guiders/__init__.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Union from ..utils import is_torch_available, logging diff --git a/src/diffusers/guiders/adaptive_projected_guidance.py b/src/diffusers/guiders/adaptive_projected_guidance.py index 8ec30d02d758..b210cb3e67aa 100644 --- a/src/diffusers/guiders/adaptive_projected_guidance.py +++ b/src/diffusers/guiders/adaptive_projected_guidance.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import math -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING import torch @@ -58,7 +60,7 @@ class AdaptiveProjectedGuidance(BaseGuidance): def __init__( self, guidance_scale: float = 7.5, - adaptive_projected_guidance_momentum: Optional[float] = None, + adaptive_projected_guidance_momentum: float | None = None, adaptive_projected_guidance_rescale: float = 15.0, eta: float = 1.0, guidance_rescale: float = 0.0, @@ -77,7 +79,7 @@ def __init__( self.use_original_formulation = use_original_formulation self.momentum_buffer = None - def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: + def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]: if self._step == 0: if self.adaptive_projected_guidance_momentum is not None: self.momentum_buffer = MomentumBuffer(self.adaptive_projected_guidance_momentum) @@ -89,8 +91,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> return data_batches def prepare_inputs_from_block_state( - self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]] - ) -> List["BlockState"]: + self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]] + ) -> list["BlockState"]: if self._step == 0: if self.adaptive_projected_guidance_momentum is not None: self.momentum_buffer = MomentumBuffer(self.adaptive_projected_guidance_momentum) @@ -101,7 +103,7 @@ def prepare_inputs_from_block_state( data_batches.append(data_batch) return data_batches - def forward(self, pred_cond: torch.Tensor, pred_uncond: Optional[torch.Tensor] = None) -> GuiderOutput: + def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput: pred = None if not self._is_apg_enabled(): @@ -204,7 +206,7 @@ def normalized_guidance( pred_cond: torch.Tensor, pred_uncond: torch.Tensor, guidance_scale: float, - momentum_buffer: Optional[MomentumBuffer] = None, + momentum_buffer: MomentumBuffer | None = None, eta: float = 1.0, norm_threshold: float = 0.0, use_original_formulation: bool = False, diff --git a/src/diffusers/guiders/adaptive_projected_guidance_mix.py b/src/diffusers/guiders/adaptive_projected_guidance_mix.py index bdc97bcf6269..559e30d2aabe 100644 --- a/src/diffusers/guiders/adaptive_projected_guidance_mix.py +++ b/src/diffusers/guiders/adaptive_projected_guidance_mix.py @@ -13,7 +13,7 @@ # limitations under the License. import math -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING import torch @@ -88,7 +88,7 @@ def __init__( self.use_original_formulation = use_original_formulation self.momentum_buffer = None - def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: + def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]: if self._step == 0: if self.adaptive_projected_guidance_momentum is not None: self.momentum_buffer = MomentumBuffer(self.adaptive_projected_guidance_momentum) @@ -100,8 +100,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> return data_batches def prepare_inputs_from_block_state( - self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]] - ) -> List["BlockState"]: + self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]] + ) -> list["BlockState"]: if self._step == 0: if self.adaptive_projected_guidance_momentum is not None: self.momentum_buffer = MomentumBuffer(self.adaptive_projected_guidance_momentum) @@ -112,7 +112,7 @@ def prepare_inputs_from_block_state( data_batches.append(data_batch) return data_batches - def forward(self, pred_cond: torch.Tensor, pred_uncond: Optional[torch.Tensor] = None) -> GuiderOutput: + def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput: pred = None # no guidance @@ -254,7 +254,7 @@ def __repr__(self) -> str: def update_momentum_buffer( pred_cond: torch.Tensor, pred_uncond: torch.Tensor, - momentum_buffer: Optional[MomentumBuffer] = None, + momentum_buffer: MomentumBuffer | None = None, ): diff = pred_cond - pred_uncond if momentum_buffer is not None: @@ -265,7 +265,7 @@ def normalized_guidance( pred_cond: torch.Tensor, pred_uncond: torch.Tensor, guidance_scale: float, - momentum_buffer: Optional[MomentumBuffer] = None, + momentum_buffer: MomentumBuffer | None = None, eta: float = 1.0, norm_threshold: float = 0.0, use_original_formulation: bool = False, diff --git a/src/diffusers/guiders/auto_guidance.py b/src/diffusers/guiders/auto_guidance.py index b7f62e2f4a6e..d6b6d3c492f0 100644 --- a/src/diffusers/guiders/auto_guidance.py +++ b/src/diffusers/guiders/auto_guidance.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import math -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any import torch @@ -36,10 +38,10 @@ class AutoGuidance(BaseGuidance): The scale parameter for classifier-free guidance. Higher values result in stronger conditioning on the text prompt, while lower values allow for more freedom in generation. Higher values may lead to saturation and deterioration of image quality. - auto_guidance_layers (`int` or `List[int]`, *optional*): + auto_guidance_layers (`int` or `list[int]`, *optional*): The layer indices to apply skip layer guidance to. Can be a single integer or a list of integers. If not provided, `skip_layer_config` must be provided. - auto_guidance_config (`LayerSkipConfig` or `List[LayerSkipConfig]`, *optional*): + auto_guidance_config (`LayerSkipConfig` or `list[LayerSkipConfig]`, *optional*): The configuration for the skip layer guidance. Can be a single `LayerSkipConfig` or a list of `LayerSkipConfig`. If not provided, `skip_layer_guidance_layers` must be provided. dropout (`float`, *optional*): @@ -65,9 +67,9 @@ class AutoGuidance(BaseGuidance): def __init__( self, guidance_scale: float = 7.5, - auto_guidance_layers: Optional[Union[int, List[int]]] = None, - auto_guidance_config: Union[LayerSkipConfig, List[LayerSkipConfig], Dict[str, Any]] = None, - dropout: Optional[float] = None, + auto_guidance_layers: int | list[int] | None = None, + auto_guidance_config: LayerSkipConfig | list[LayerSkipConfig] | dict[str, Any] = None, + dropout: float | None = None, guidance_rescale: float = 0.0, use_original_formulation: bool = False, start: float = 0.0, @@ -133,7 +135,7 @@ def cleanup_models(self, denoiser: torch.nn.Module) -> None: registry = HookRegistry.check_if_exists_or_initialize(denoiser) registry.remove_hook(name, recurse=True) - def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: + def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]: tuple_indices = [0] if self.num_conditions == 1 else [0, 1] data_batches = [] for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions): @@ -142,8 +144,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> return data_batches def prepare_inputs_from_block_state( - self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]] - ) -> List["BlockState"]: + self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]] + ) -> list["BlockState"]: tuple_indices = [0] if self.num_conditions == 1 else [0, 1] data_batches = [] for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions): @@ -151,7 +153,7 @@ def prepare_inputs_from_block_state( data_batches.append(data_batch) return data_batches - def forward(self, pred_cond: torch.Tensor, pred_uncond: Optional[torch.Tensor] = None) -> GuiderOutput: + def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput: pred = None if not self._is_ag_enabled(): diff --git a/src/diffusers/guiders/classifier_free_guidance.py b/src/diffusers/guiders/classifier_free_guidance.py index 5e55d4d869c1..a2180a626bfc 100644 --- a/src/diffusers/guiders/classifier_free_guidance.py +++ b/src/diffusers/guiders/classifier_free_guidance.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import math -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING import torch @@ -91,7 +93,7 @@ def __init__( self.guidance_rescale = guidance_rescale self.use_original_formulation = use_original_formulation - def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: + def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]: tuple_indices = [0] if self.num_conditions == 1 else [0, 1] data_batches = [] for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions): @@ -100,8 +102,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> return data_batches def prepare_inputs_from_block_state( - self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]] - ) -> List["BlockState"]: + self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]] + ) -> list["BlockState"]: tuple_indices = [0] if self.num_conditions == 1 else [0, 1] data_batches = [] for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions): @@ -109,7 +111,7 @@ def prepare_inputs_from_block_state( data_batches.append(data_batch) return data_batches - def forward(self, pred_cond: torch.Tensor, pred_uncond: Optional[torch.Tensor] = None) -> GuiderOutput: + def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput: pred = None if not self._is_cfg_enabled(): diff --git a/src/diffusers/guiders/classifier_free_zero_star_guidance.py b/src/diffusers/guiders/classifier_free_zero_star_guidance.py index 23b492e51b02..dd71c7537cac 100644 --- a/src/diffusers/guiders/classifier_free_zero_star_guidance.py +++ b/src/diffusers/guiders/classifier_free_zero_star_guidance.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import math -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING import torch @@ -77,7 +79,7 @@ def __init__( self.guidance_rescale = guidance_rescale self.use_original_formulation = use_original_formulation - def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: + def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]: tuple_indices = [0] if self.num_conditions == 1 else [0, 1] data_batches = [] for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions): @@ -86,8 +88,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> return data_batches def prepare_inputs_from_block_state( - self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]] - ) -> List["BlockState"]: + self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]] + ) -> list["BlockState"]: tuple_indices = [0] if self.num_conditions == 1 else [0, 1] data_batches = [] for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions): @@ -95,7 +97,7 @@ def prepare_inputs_from_block_state( data_batches.append(data_batch) return data_batches - def forward(self, pred_cond: torch.Tensor, pred_uncond: Optional[torch.Tensor] = None) -> GuiderOutput: + def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput: pred = None # YiYi Notes: add default behavior for self._enabled == False diff --git a/src/diffusers/guiders/frequency_decoupled_guidance.py b/src/diffusers/guiders/frequency_decoupled_guidance.py index 4ec6e2d36da9..b92ddf2c03f9 100644 --- a/src/diffusers/guiders/frequency_decoupled_guidance.py +++ b/src/diffusers/guiders/frequency_decoupled_guidance.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import math -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING import torch @@ -37,7 +39,7 @@ build_laplacian_pyramid_func = None -def project(v0: torch.Tensor, v1: torch.Tensor, upcast_to_double: bool = True) -> Tuple[torch.Tensor, torch.Tensor]: +def project(v0: torch.Tensor, v1: torch.Tensor, upcast_to_double: bool = True) -> tuple[torch.Tensor, torch.Tensor]: """ Project vector v0 onto vector v1, returning the parallel and orthogonal components of v0. Implementation from paper (Algorithm 2). @@ -58,7 +60,7 @@ def project(v0: torch.Tensor, v1: torch.Tensor, upcast_to_double: bool = True) - return v0_parallel, v0_orthogonal -def build_image_from_pyramid(pyramid: List[torch.Tensor]) -> torch.Tensor: +def build_image_from_pyramid(pyramid: list[torch.Tensor]) -> torch.Tensor: """ Recovers the data space latents from the Laplacian pyramid frequency space. Implementation from the paper (Algorithm 2). @@ -99,19 +101,19 @@ class FrequencyDecoupledGuidance(BaseGuidance): paper. By default, we use the diffusers-native implementation that has been in the codebase for a long time. Args: - guidance_scales (`List[float]`, defaults to `[10.0, 5.0]`): + guidance_scales (`list[float]`, defaults to `[10.0, 5.0]`): The scale parameter for frequency-decoupled guidance for each frequency component, listed from highest frequency level to lowest. Higher values result in stronger conditioning on the text prompt, while lower values allow for more freedom in generation. Higher values may lead to saturation and deterioration of image quality. The FDG authors recommend using higher guidance scales for higher frequency components and lower guidance scales for lower frequency components (so `guidance_scales` should typically be sorted in descending order). - guidance_rescale (`float` or `List[float]`, defaults to `0.0`): + guidance_rescale (`float` or `list[float]`, defaults to `0.0`): The rescale factor applied to the noise predictions. This is used to improve image quality and fix overexposure. Based on Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). If a list is supplied, it should be the same length as `guidance_scales`. - parallel_weights (`float` or `List[float]`, *optional*): + parallel_weights (`float` or `list[float]`, *optional*): Optional weights for the parallel component of each frequency component of the projected CFG shift. If not set, the weights will default to `1.0` for all components, which corresponds to using the normal CFG shift (that is, equal weights for the parallel and orthogonal components). If set, a value in `[0, 1]` is @@ -120,10 +122,10 @@ class FrequencyDecoupledGuidance(BaseGuidance): Whether to use the original formulation of classifier-free guidance as proposed in the paper. By default, we use the diffusers-native implementation that has been in the codebase for a long time. See [~guiders.classifier_free_guidance.ClassifierFreeGuidance] for more details. - start (`float` or `List[float]`, defaults to `0.0`): + start (`float` or `list[float]`, defaults to `0.0`): The fraction of the total number of denoising steps after which guidance starts. If a list is supplied, it should be the same length as `guidance_scales`. - stop (`float` or `List[float]`, defaults to `1.0`): + stop (`float` or `list[float]`, defaults to `1.0`): The fraction of the total number of denoising steps after which guidance stops. If a list is supplied, it should be the same length as `guidance_scales`. guidance_rescale_space (`str`, defaults to `"data"`): @@ -141,12 +143,12 @@ class FrequencyDecoupledGuidance(BaseGuidance): @register_to_config def __init__( self, - guidance_scales: Union[List[float], Tuple[float]] = [10.0, 5.0], - guidance_rescale: Union[float, List[float], Tuple[float]] = 0.0, - parallel_weights: Optional[Union[float, List[float], Tuple[float]]] = None, + guidance_scales: list[float] | tuple[float] = [10.0, 5.0], + guidance_rescale: float | list[float] | tuple[float] = 0.0, + parallel_weights: float | list[float] | tuple[float] | None = None, use_original_formulation: bool = False, - start: Union[float, List[float], Tuple[float]] = 0.0, - stop: Union[float, List[float], Tuple[float]] = 1.0, + start: float | list[float] | tuple[float] = 0.0, + stop: float | list[float] | tuple[float] = 1.0, guidance_rescale_space: str = "data", upcast_to_double: bool = True, enabled: bool = True, @@ -218,7 +220,7 @@ def __init__( f"({len(self.guidance_scales)})" ) - def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: + def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]: tuple_indices = [0] if self.num_conditions == 1 else [0, 1] data_batches = [] for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions): @@ -227,8 +229,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> return data_batches def prepare_inputs_from_block_state( - self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]] - ) -> List["BlockState"]: + self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]] + ) -> list["BlockState"]: tuple_indices = [0] if self.num_conditions == 1 else [0, 1] data_batches = [] for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions): @@ -236,7 +238,7 @@ def prepare_inputs_from_block_state( data_batches.append(data_batch) return data_batches - def forward(self, pred_cond: torch.Tensor, pred_uncond: Optional[torch.Tensor] = None) -> GuiderOutput: + def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput: pred = None if not self._is_fdg_enabled(): diff --git a/src/diffusers/guiders/guider_utils.py b/src/diffusers/guiders/guider_utils.py index 6c328328fc3b..7be68424c345 100644 --- a/src/diffusers/guiders/guider_utils.py +++ b/src/diffusers/guiders/guider_utils.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import os -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any import torch from huggingface_hub.utils import validate_hf_hub_args @@ -51,7 +53,7 @@ def __init__(self, start: float = 0.0, stop: float = 1.0, enabled: bool = True): self._num_inference_steps: int = None self._timestep: torch.LongTensor = None self._count_prepared = 0 - self._input_fields: Dict[str, Union[str, Tuple[str, str]]] = None + self._input_fields: dict[str, str | tuple[str, str]] = None self._enabled = enabled if not (0.0 <= start < 1.0): @@ -101,11 +103,11 @@ def set_state(self, step: int, num_inference_steps: int, timestep: torch.LongTen self._timestep = timestep self._count_prepared = 0 - def get_state(self) -> Dict[str, Any]: + def get_state(self) -> dict[str, Any]: """ Returns the current state of the guidance technique as a dictionary. The state variables will be included in the __repr__ method. Returns: - `Dict[str, Any]`: A dictionary containing the current state variables including: + `dict[str, Any]`: A dictionary containing the current state variables including: - step: Current inference step - num_inference_steps: Total number of inference steps - timestep: Current timestep tensor @@ -163,15 +165,15 @@ def cleanup_models(self, denoiser: torch.nn.Module) -> None: """ pass - def prepare_inputs(self, data: "BlockState") -> List["BlockState"]: + def prepare_inputs(self, data: "BlockState") -> list["BlockState"]: raise NotImplementedError("BaseGuidance::prepare_inputs must be implemented in subclasses.") def prepare_inputs_from_block_state( - self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]] - ) -> List["BlockState"]: + self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]] + ) -> list["BlockState"]: raise NotImplementedError("BaseGuidance::prepare_inputs_from_block_state must be implemented in subclasses.") - def __call__(self, data: List["BlockState"]) -> Any: + def __call__(self, data: list["BlockState"]) -> Any: if not all(hasattr(d, "noise_pred") for d in data): raise ValueError("Expected all data to have `noise_pred` attribute.") if len(data) != self.num_conditions: @@ -199,7 +201,7 @@ def num_conditions(self) -> int: @classmethod def _prepare_batch( cls, - data: Dict[str, Tuple[torch.Tensor, torch.Tensor]], + data: dict[str, tuple[torch.Tensor, torch.Tensor]], tuple_index: int, identifier: str, ) -> "BlockState": @@ -208,7 +210,7 @@ def _prepare_batch( `BaseGuidance` class. It prepares the batch based on the provided tuple index. Args: - input_fields (`Dict[str, Union[str, Tuple[str, str]]]`): + input_fields (`dict[str, str | tuple[str, str]]`): A dictionary where the keys are the names of the fields that will be used to store the data once it is prepared with `prepare_inputs`. The values can be either a string or a tuple of length 2, which is used to look up the required data provided for preparation. If a string is provided, it will be used as the @@ -242,7 +244,7 @@ def _prepare_batch( @classmethod def _prepare_batch_from_block_state( cls, - input_fields: Dict[str, Union[str, Tuple[str, str]]], + input_fields: dict[str, str | tuple[str, str]], data: "BlockState", tuple_index: int, identifier: str, @@ -252,7 +254,7 @@ def _prepare_batch_from_block_state( `BaseGuidance` class. It prepares the batch based on the provided tuple index. Args: - input_fields (`Dict[str, Union[str, Tuple[str, str]]]`): + input_fields (`dict[str, str | tuple[str, str]]`): A dictionary where the keys are the names of the fields that will be used to store the data once it is prepared with `prepare_inputs`. The values can be either a string or a tuple of length 2, which is used to look up the required data provided for preparation. If a string is provided, it will be used as the @@ -288,8 +290,8 @@ def _prepare_batch_from_block_state( @validate_hf_hub_args def from_pretrained( cls, - pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None, - subfolder: Optional[str] = None, + pretrained_model_name_or_path: str | os.PathLike | None = None, + subfolder: str | None = None, return_unused_kwargs=False, **kwargs, ) -> Self: @@ -308,14 +310,14 @@ def from_pretrained( The subfolder location of a model file within a larger model repository on the Hub or locally. return_unused_kwargs (`bool`, *optional*, defaults to `False`): Whether kwargs that are not consumed by the Python class should be returned or not. - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. output_loading_info(`bool`, *optional*, defaults to `False`): @@ -345,7 +347,7 @@ def from_pretrained( ) return cls.from_config(config, return_unused_kwargs=return_unused_kwargs, **kwargs) - def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs): + def save_pretrained(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs): """ Save a guider configuration object to a directory so that it can be reloaded using the [`~BaseGuidance.from_pretrained`] class method. @@ -357,7 +359,7 @@ def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the repository you want to push to with `repo_id` (will default to the name of `save_directory` in your namespace). - kwargs (`Dict[str, Any]`, *optional*): + kwargs (`dict[str, Any]`, *optional*): Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. """ self.save_config(save_directory=save_directory, push_to_hub=push_to_hub, **kwargs) @@ -365,8 +367,8 @@ def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: class GuiderOutput(BaseOutput): pred: torch.Tensor - pred_cond: Optional[torch.Tensor] - pred_uncond: Optional[torch.Tensor] + pred_cond: torch.Tensor | None + pred_uncond: torch.Tensor | None def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): diff --git a/src/diffusers/guiders/magnitude_aware_guidance.py b/src/diffusers/guiders/magnitude_aware_guidance.py index b81cf0d3a1f9..e83545fd889a 100644 --- a/src/diffusers/guiders/magnitude_aware_guidance.py +++ b/src/diffusers/guiders/magnitude_aware_guidance.py @@ -13,7 +13,7 @@ # limitations under the License. import math -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING import torch @@ -71,7 +71,7 @@ def __init__( self.guidance_rescale = guidance_rescale self.use_original_formulation = use_original_formulation - def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: + def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]: tuple_indices = [0] if self.num_conditions == 1 else [0, 1] data_batches = [] for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions): @@ -80,8 +80,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> return data_batches def prepare_inputs_from_block_state( - self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]] - ) -> List["BlockState"]: + self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]] + ) -> list["BlockState"]: tuple_indices = [0] if self.num_conditions == 1 else [0, 1] data_batches = [] for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions): @@ -89,7 +89,7 @@ def prepare_inputs_from_block_state( data_batches.append(data_batch) return data_batches - def forward(self, pred_cond: torch.Tensor, pred_uncond: Optional[torch.Tensor] = None) -> GuiderOutput: + def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput: pred = None if not self._is_mambo_g_enabled(): diff --git a/src/diffusers/guiders/perturbed_attention_guidance.py b/src/diffusers/guiders/perturbed_attention_guidance.py index f233e90ca410..904d319ec3bb 100644 --- a/src/diffusers/guiders/perturbed_attention_guidance.py +++ b/src/diffusers/guiders/perturbed_attention_guidance.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import math -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any import torch @@ -58,10 +60,10 @@ class PerturbedAttentionGuidance(BaseGuidance): The fraction of the total number of denoising steps after which perturbed attention guidance starts. perturbed_guidance_stop (`float`, defaults to `0.2`): The fraction of the total number of denoising steps after which perturbed attention guidance stops. - perturbed_guidance_layers (`int` or `List[int]`, *optional*): + perturbed_guidance_layers (`int` or `list[int]`, *optional*): The layer indices to apply perturbed attention guidance to. Can be a single integer or a list of integers. If not provided, `perturbed_guidance_config` must be provided. - perturbed_guidance_config (`LayerSkipConfig` or `List[LayerSkipConfig]`, *optional*): + perturbed_guidance_config (`LayerSkipConfig` or `list[LayerSkipConfig]`, *optional*): The configuration for the perturbed attention guidance. Can be a single `LayerSkipConfig` or a list of `LayerSkipConfig`. If not provided, `perturbed_guidance_layers` must be provided. guidance_rescale (`float`, defaults to `0.0`): @@ -92,8 +94,8 @@ def __init__( perturbed_guidance_scale: float = 2.8, perturbed_guidance_start: float = 0.01, perturbed_guidance_stop: float = 0.2, - perturbed_guidance_layers: Optional[Union[int, List[int]]] = None, - perturbed_guidance_config: Union[LayerSkipConfig, List[LayerSkipConfig], Dict[str, Any]] = None, + perturbed_guidance_layers: int | list[int] | None = None, + perturbed_guidance_config: LayerSkipConfig | list[LayerSkipConfig] | dict[str, Any] = None, guidance_rescale: float = 0.0, use_original_formulation: bool = False, start: float = 0.0, @@ -169,7 +171,7 @@ def cleanup_models(self, denoiser: torch.nn.Module) -> None: registry.remove_hook(hook_name, recurse=True) # Copied from diffusers.guiders.skip_layer_guidance.SkipLayerGuidance.prepare_inputs - def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: + def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]: if self.num_conditions == 1: tuple_indices = [0] input_predictions = ["pred_cond"] @@ -188,8 +190,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> return data_batches def prepare_inputs_from_block_state( - self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]] - ) -> List["BlockState"]: + self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]] + ) -> list["BlockState"]: if self.num_conditions == 1: tuple_indices = [0] input_predictions = ["pred_cond"] @@ -211,8 +213,8 @@ def prepare_inputs_from_block_state( def forward( self, pred_cond: torch.Tensor, - pred_uncond: Optional[torch.Tensor] = None, - pred_cond_skip: Optional[torch.Tensor] = None, + pred_uncond: torch.Tensor | None = None, + pred_cond_skip: torch.Tensor | None = None, ) -> GuiderOutput: pred = None diff --git a/src/diffusers/guiders/skip_layer_guidance.py b/src/diffusers/guiders/skip_layer_guidance.py index e6109300d99c..cb7e85e179d2 100644 --- a/src/diffusers/guiders/skip_layer_guidance.py +++ b/src/diffusers/guiders/skip_layer_guidance.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import math -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any import torch @@ -64,11 +66,11 @@ class SkipLayerGuidance(BaseGuidance): The fraction of the total number of denoising steps after which skip layer guidance starts. skip_layer_guidance_stop (`float`, defaults to `0.2`): The fraction of the total number of denoising steps after which skip layer guidance stops. - skip_layer_guidance_layers (`int` or `List[int]`, *optional*): + skip_layer_guidance_layers (`int` or `list[int]`, *optional*): The layer indices to apply skip layer guidance to. Can be a single integer or a list of integers. If not provided, `skip_layer_config` must be provided. The recommended values are `[7, 8, 9]` for Stable Diffusion 3.5 Medium. - skip_layer_config (`LayerSkipConfig` or `List[LayerSkipConfig]`, *optional*): + skip_layer_config (`LayerSkipConfig` or `list[LayerSkipConfig]`, *optional*): The configuration for the skip layer guidance. Can be a single `LayerSkipConfig` or a list of `LayerSkipConfig`. If not provided, `skip_layer_guidance_layers` must be provided. guidance_rescale (`float`, defaults to `0.0`): @@ -94,8 +96,8 @@ def __init__( skip_layer_guidance_scale: float = 2.8, skip_layer_guidance_start: float = 0.01, skip_layer_guidance_stop: float = 0.2, - skip_layer_guidance_layers: Optional[Union[int, List[int]]] = None, - skip_layer_config: Union[LayerSkipConfig, List[LayerSkipConfig], Dict[str, Any]] = None, + skip_layer_guidance_layers: int | list[int] | None = None, + skip_layer_config: LayerSkipConfig | list[LayerSkipConfig] | dict[str, Any] = None, guidance_rescale: float = 0.0, use_original_formulation: bool = False, start: float = 0.0, @@ -165,7 +167,7 @@ def cleanup_models(self, denoiser: torch.nn.Module) -> None: for hook_name in self._skip_layer_hook_names: registry.remove_hook(hook_name, recurse=True) - def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: + def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]: if self.num_conditions == 1: tuple_indices = [0] input_predictions = ["pred_cond"] @@ -184,8 +186,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> return data_batches def prepare_inputs_from_block_state( - self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]] - ) -> List["BlockState"]: + self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]] + ) -> list["BlockState"]: if self.num_conditions == 1: tuple_indices = [0] input_predictions = ["pred_cond"] @@ -206,8 +208,8 @@ def prepare_inputs_from_block_state( def forward( self, pred_cond: torch.Tensor, - pred_uncond: Optional[torch.Tensor] = None, - pred_cond_skip: Optional[torch.Tensor] = None, + pred_uncond: torch.Tensor | None = None, + pred_cond_skip: torch.Tensor | None = None, ) -> GuiderOutput: pred = None diff --git a/src/diffusers/guiders/smoothed_energy_guidance.py b/src/diffusers/guiders/smoothed_energy_guidance.py index 6c3906e820e0..4767607421de 100644 --- a/src/diffusers/guiders/smoothed_energy_guidance.py +++ b/src/diffusers/guiders/smoothed_energy_guidance.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import math -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING import torch @@ -54,11 +56,11 @@ class SmoothedEnergyGuidance(BaseGuidance): The fraction of the total number of denoising steps after which smoothed energy guidance starts. seg_guidance_stop (`float`, defaults to `1.0`): The fraction of the total number of denoising steps after which smoothed energy guidance stops. - seg_guidance_layers (`int` or `List[int]`, *optional*): + seg_guidance_layers (`int` or `list[int]`, *optional*): The layer indices to apply smoothed energy guidance to. Can be a single integer or a list of integers. If not provided, `seg_guidance_config` must be provided. The recommended values are `[7, 8, 9]` for Stable Diffusion 3.5 Medium. - seg_guidance_config (`SmoothedEnergyGuidanceConfig` or `List[SmoothedEnergyGuidanceConfig]`, *optional*): + seg_guidance_config (`SmoothedEnergyGuidanceConfig` or `list[SmoothedEnergyGuidanceConfig]`, *optional*): The configuration for the smoothed energy layer guidance. Can be a single `SmoothedEnergyGuidanceConfig` or a list of `SmoothedEnergyGuidanceConfig`. If not provided, `seg_guidance_layers` must be provided. guidance_rescale (`float`, defaults to `0.0`): @@ -86,8 +88,8 @@ def __init__( seg_blur_threshold_inf: float = 9999.0, seg_guidance_start: float = 0.0, seg_guidance_stop: float = 1.0, - seg_guidance_layers: Optional[Union[int, List[int]]] = None, - seg_guidance_config: Union[SmoothedEnergyGuidanceConfig, List[SmoothedEnergyGuidanceConfig]] = None, + seg_guidance_layers: int | list[int] | None = None, + seg_guidance_config: SmoothedEnergyGuidanceConfig | list[SmoothedEnergyGuidanceConfig] = None, guidance_rescale: float = 0.0, use_original_formulation: bool = False, start: float = 0.0, @@ -154,7 +156,7 @@ def cleanup_models(self, denoiser: torch.nn.Module): for hook_name in self._seg_layer_hook_names: registry.remove_hook(hook_name, recurse=True) - def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: + def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]: if self.num_conditions == 1: tuple_indices = [0] input_predictions = ["pred_cond"] @@ -173,8 +175,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> return data_batches def prepare_inputs_from_block_state( - self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]] - ) -> List["BlockState"]: + self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]] + ) -> list["BlockState"]: if self.num_conditions == 1: tuple_indices = [0] input_predictions = ["pred_cond"] @@ -195,8 +197,8 @@ def prepare_inputs_from_block_state( def forward( self, pred_cond: torch.Tensor, - pred_uncond: Optional[torch.Tensor] = None, - pred_cond_seg: Optional[torch.Tensor] = None, + pred_uncond: torch.Tensor | None = None, + pred_cond_seg: torch.Tensor | None = None, ) -> GuiderOutput: pred = None diff --git a/src/diffusers/guiders/tangential_classifier_free_guidance.py b/src/diffusers/guiders/tangential_classifier_free_guidance.py index 76899c6e8494..c8911f4a69d9 100644 --- a/src/diffusers/guiders/tangential_classifier_free_guidance.py +++ b/src/diffusers/guiders/tangential_classifier_free_guidance.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import math -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING import torch @@ -66,7 +68,7 @@ def __init__( self.guidance_rescale = guidance_rescale self.use_original_formulation = use_original_formulation - def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]: + def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]: tuple_indices = [0] if self.num_conditions == 1 else [0, 1] data_batches = [] for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions): @@ -75,8 +77,8 @@ def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> return data_batches def prepare_inputs_from_block_state( - self, data: "BlockState", input_fields: Dict[str, Union[str, Tuple[str, str]]] - ) -> List["BlockState"]: + self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]] + ) -> list["BlockState"]: tuple_indices = [0] if self.num_conditions == 1 else [0, 1] data_batches = [] for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions): @@ -84,7 +86,7 @@ def prepare_inputs_from_block_state( data_batches.append(data_batch) return data_batches - def forward(self, pred_cond: torch.Tensor, pred_uncond: Optional[torch.Tensor] = None) -> GuiderOutput: + def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput: pred = None if not self._is_tcfg_enabled(): diff --git a/src/diffusers/hooks/__init__.py b/src/diffusers/hooks/__init__.py index eb12b8a52a1e..23c8bc92b2f1 100644 --- a/src/diffusers/hooks/__init__.py +++ b/src/diffusers/hooks/__init__.py @@ -23,6 +23,7 @@ from .hooks import HookRegistry, ModelHook from .layer_skip import LayerSkipConfig, apply_layer_skip from .layerwise_casting import apply_layerwise_casting, apply_layerwise_casting_hook + from .mag_cache import MagCacheConfig, apply_mag_cache from .pyramid_attention_broadcast import PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast from .smoothed_energy_guidance_utils import SmoothedEnergyGuidanceConfig from .taylorseer_cache import TaylorSeerCacheConfig, apply_taylorseer_cache diff --git a/src/diffusers/hooks/_common.py b/src/diffusers/hooks/_common.py index ca7934e5c313..fa7ab770da6d 100644 --- a/src/diffusers/hooks/_common.py +++ b/src/diffusers/hooks/_common.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional - import torch from ..models.attention import AttentionModuleMixin, FeedForward, LuminaFeedForward @@ -23,7 +21,13 @@ _ATTENTION_CLASSES = (Attention, MochiAttention, AttentionModuleMixin) _FEEDFORWARD_CLASSES = (FeedForward, LuminaFeedForward) -_SPATIAL_TRANSFORMER_BLOCK_IDENTIFIERS = ("blocks", "transformer_blocks", "single_transformer_blocks", "layers") +_SPATIAL_TRANSFORMER_BLOCK_IDENTIFIERS = ( + "blocks", + "transformer_blocks", + "single_transformer_blocks", + "layers", + "visual_transformer_blocks", +) _TEMPORAL_TRANSFORMER_BLOCK_IDENTIFIERS = ("temporal_transformer_blocks",) _CROSS_TRANSFORMER_BLOCK_IDENTIFIERS = ("blocks", "transformer_blocks", "layers") @@ -44,12 +48,13 @@ torch.nn.ConvTranspose2d, torch.nn.ConvTranspose3d, torch.nn.Linear, + torch.nn.Embedding, # TODO(aryan): look into torch.nn.LayerNorm, torch.nn.GroupNorm later, seems to be causing some issues with CogVideoX # because of double invocation of the same norm layer in CogVideoXLayerNorm ) -def _get_submodule_from_fqn(module: torch.nn.Module, fqn: str) -> Optional[torch.nn.Module]: +def _get_submodule_from_fqn(module: torch.nn.Module, fqn: str) -> torch.nn.Module | None: for submodule_name, submodule in module.named_modules(): if submodule_name == fqn: return submodule diff --git a/src/diffusers/hooks/_helpers.py b/src/diffusers/hooks/_helpers.py index da7313cb4737..0267dd481a89 100644 --- a/src/diffusers/hooks/_helpers.py +++ b/src/diffusers/hooks/_helpers.py @@ -14,7 +14,7 @@ import inspect from dataclasses import dataclass -from typing import Any, Callable, Dict, Type +from typing import Any, Callable, Type @dataclass @@ -26,9 +26,10 @@ class AttentionProcessorMetadata: class TransformerBlockMetadata: return_hidden_states_index: int = None return_encoder_hidden_states_index: int = None + hidden_states_argument_name: str = "hidden_states" _cls: Type = None - _cached_parameter_indices: Dict[str, int] = None + _cached_parameter_indices: dict[str, int] = None def _get_parameter_from_args_kwargs(self, identifier: str, args=(), kwargs=None): kwargs = kwargs or {} @@ -169,7 +170,7 @@ def _register_attention_processors_metadata(): def _register_transformer_blocks_metadata(): - from ..models.attention import BasicTransformerBlock + from ..models.attention import BasicTransformerBlock, JointTransformerBlock from ..models.transformers.cogvideox_transformer_3d import CogVideoXBlock from ..models.transformers.transformer_bria import BriaTransformerBlock from ..models.transformers.transformer_cogview4 import CogView4TransformerBlock @@ -184,6 +185,7 @@ def _register_transformer_blocks_metadata(): HunyuanImageSingleTransformerBlock, HunyuanImageTransformerBlock, ) + from ..models.transformers.transformer_kandinsky import Kandinsky5TransformerDecoderBlock from ..models.transformers.transformer_ltx import LTXVideoTransformerBlock from ..models.transformers.transformer_mochi import MochiTransformerBlock from ..models.transformers.transformer_qwenimage import QwenImageTransformerBlock @@ -331,6 +333,24 @@ def _register_transformer_blocks_metadata(): ), ) + TransformerBlockRegistry.register( + model_class=JointTransformerBlock, + metadata=TransformerBlockMetadata( + return_hidden_states_index=1, + return_encoder_hidden_states_index=0, + ), + ) + + # Kandinsky 5.0 (Kandinsky5TransformerDecoderBlock) + TransformerBlockRegistry.register( + model_class=Kandinsky5TransformerDecoderBlock, + metadata=TransformerBlockMetadata( + return_hidden_states_index=0, + return_encoder_hidden_states_index=None, + hidden_states_argument_name="visual_embed", + ), + ) + # fmt: off def _skip_attention___ret___hidden_states(self, *args, **kwargs): diff --git a/src/diffusers/hooks/context_parallel.py b/src/diffusers/hooks/context_parallel.py index 6491d17b4f46..6130be2b8290 100644 --- a/src/diffusers/hooks/context_parallel.py +++ b/src/diffusers/hooks/context_parallel.py @@ -11,12 +11,14 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - +import copy +import functools import inspect from dataclasses import dataclass -from typing import Dict, List, Type, Union +from typing import Type import torch +import torch.distributed as dist if torch.distributed.is_available(): @@ -27,9 +29,10 @@ ContextParallelInput, ContextParallelModelPlan, ContextParallelOutput, + gather_size_by_comm, ) from ..utils import get_logger -from ..utils.torch_utils import unwrap_module +from ..utils.torch_utils import maybe_allow_in_graph, unwrap_module from .hooks import HookRegistry, ModelHook @@ -42,7 +45,7 @@ # TODO(aryan): consolidate with ._helpers.TransformerBlockMetadata @dataclass class ModuleForwardMetadata: - cached_parameter_indices: Dict[str, int] = None + cached_parameter_indices: dict[str, int] = None _cls: Type = None def _get_parameter_from_args_kwargs(self, identifier: str, args=(), kwargs=None): @@ -78,7 +81,7 @@ def _get_parameter_from_args_kwargs(self, identifier: str, args=(), kwargs=None) def apply_context_parallel( module: torch.nn.Module, parallel_config: ContextParallelConfig, - plan: Dict[str, ContextParallelModelPlan], + plan: dict[str, ContextParallelModelPlan], ) -> None: """Apply context parallel on a model.""" logger.debug(f"Applying context parallel with CP mesh: {parallel_config._mesh} and plan: {plan}") @@ -107,7 +110,7 @@ def apply_context_parallel( registry.register_hook(hook, hook_name) -def remove_context_parallel(module: torch.nn.Module, plan: Dict[str, ContextParallelModelPlan]) -> None: +def remove_context_parallel(module: torch.nn.Module, plan: dict[str, ContextParallelModelPlan]) -> None: for module_id, cp_model_plan in plan.items(): submodule = _get_submodule_by_name(module, module_id) if not isinstance(submodule, list): @@ -208,6 +211,10 @@ def _prepare_cp_input(self, x: torch.Tensor, cp_input: ContextParallelInput) -> ) return x else: + if self.parallel_config.ulysses_anything: + return PartitionAnythingSharder.shard_anything( + x, cp_input.split_dim, self.parallel_config._flattened_mesh + ) return EquipartitionSharder.shard(x, cp_input.split_dim, self.parallel_config._flattened_mesh) @@ -233,7 +240,14 @@ def post_forward(self, module, output): for i, cpm in enumerate(self.metadata): if cpm is None: continue - output[i] = EquipartitionSharder.unshard(output[i], cpm.gather_dim, self.parallel_config._flattened_mesh) + if self.parallel_config.ulysses_anything: + output[i] = PartitionAnythingSharder.unshard_anything( + output[i], cpm.gather_dim, self.parallel_config._flattened_mesh + ) + else: + output[i] = EquipartitionSharder.unshard( + output[i], cpm.gather_dim, self.parallel_config._flattened_mesh + ) return output[0] if is_tensor else tuple(output) @@ -274,13 +288,80 @@ def unshard(cls, tensor: torch.Tensor, dim: int, mesh: torch.distributed.device_ return tensor -def _get_submodule_by_name(model: torch.nn.Module, name: str) -> Union[torch.nn.Module, List[torch.nn.Module]]: +class AllGatherAnythingFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, tensor: torch.Tensor, dim: int, group: dist.device_mesh.DeviceMesh): + ctx.dim = dim + ctx.group = group + ctx.world_size = dist.get_world_size(group) + ctx.rank = dist.get_rank(group) + gathered_tensor = _all_gather_anything(tensor, dim, group) + return gathered_tensor + + @staticmethod + def backward(ctx, grad_output): + # NOTE: We use `tensor_split` instead of chunk, because the `chunk` + # function may return fewer than the specified number of chunks! + grad_splits = torch.tensor_split(grad_output, ctx.world_size, dim=ctx.dim) + return grad_splits[ctx.rank], None, None + + +class PartitionAnythingSharder: + @classmethod + def shard_anything( + cls, tensor: torch.Tensor, dim: int, mesh: torch.distributed.device_mesh.DeviceMesh + ) -> torch.Tensor: + assert tensor.size()[dim] >= mesh.size(), ( + f"Cannot shard tensor of size {tensor.size()} along dim {dim} across mesh of size {mesh.size()}." + ) + # NOTE: We use `tensor_split` instead of chunk, because the `chunk` + # function may return fewer than the specified number of chunks! + return tensor.tensor_split(mesh.size(), dim=dim)[dist.get_rank(mesh.get_group())] + + @classmethod + def unshard_anything( + cls, tensor: torch.Tensor, dim: int, mesh: torch.distributed.device_mesh.DeviceMesh + ) -> torch.Tensor: + tensor = tensor.contiguous() + tensor = AllGatherAnythingFunction.apply(tensor, dim, mesh.get_group()) + return tensor + + +@functools.lru_cache(maxsize=64) +def _fill_gather_shapes(shape: tuple[int], gather_dims: tuple[int], dim: int, world_size: int) -> list[list[int]]: + gather_shapes = [] + for i in range(world_size): + rank_shape = list(copy.deepcopy(shape)) + rank_shape[dim] = gather_dims[i] + gather_shapes.append(rank_shape) + return gather_shapes + + +@maybe_allow_in_graph +def _all_gather_anything(tensor: torch.Tensor, dim: int, group: dist.device_mesh.DeviceMesh) -> torch.Tensor: + world_size = dist.get_world_size(group=group) + + tensor = tensor.contiguous() + shape = tensor.shape + rank_dim = shape[dim] + gather_dims = gather_size_by_comm(rank_dim, group) + + gather_shapes = _fill_gather_shapes(tuple(shape), tuple(gather_dims), dim, world_size) + + gathered_tensors = [torch.empty(shape, device=tensor.device, dtype=tensor.dtype) for shape in gather_shapes] + + dist.all_gather(gathered_tensors, tensor, group=group) + gathered_tensor = torch.cat(gathered_tensors, dim=dim) + return gathered_tensor + + +def _get_submodule_by_name(model: torch.nn.Module, name: str) -> torch.nn.Module | list[torch.nn.Module]: if name.count("*") > 1: raise ValueError("Wildcard '*' can only be used once in the name") return _find_submodule_by_name(model, name) -def _find_submodule_by_name(model: torch.nn.Module, name: str) -> Union[torch.nn.Module, List[torch.nn.Module]]: +def _find_submodule_by_name(model: torch.nn.Module, name: str) -> torch.nn.Module | list[torch.nn.Module]: if name == "": return model first_atom, remaining_name = name.split(".", 1) if "." in name else (name, "") diff --git a/src/diffusers/hooks/faster_cache.py b/src/diffusers/hooks/faster_cache.py index a01afeffdb95..682cebe36c7d 100644 --- a/src/diffusers/hooks/faster_cache.py +++ b/src/diffusers/hooks/faster_cache.py @@ -14,7 +14,7 @@ import re from dataclasses import dataclass -from typing import Any, Callable, List, Optional, Tuple +from typing import Any, Callable import torch @@ -60,7 +60,7 @@ class FasterCacheConfig: Calculate the attention states every `N` iterations. If this is set to `N`, the attention computation will be skipped `N - 1` times (i.e., cached attention states will be reused) before computing the new attention states again. - spatial_attention_timestep_skip_range (`Tuple[float, float]`, defaults to `(-1, 681)`): + spatial_attention_timestep_skip_range (`tuple[float, float]`, defaults to `(-1, 681)`): The timestep range within which the spatial attention computation can be skipped without a significant loss in quality. This is to be determined by the user based on the underlying model. The first value in the tuple is the lower bound and the second value is the upper bound. Typically, diffusion timesteps for @@ -68,17 +68,17 @@ class FasterCacheConfig: timestep 0). For the default values, this would mean that the spatial attention computation skipping will be applicable only after denoising timestep 681 is reached, and continue until the end of the denoising process. - temporal_attention_timestep_skip_range (`Tuple[float, float]`, *optional*, defaults to `None`): + temporal_attention_timestep_skip_range (`tuple[float, float]`, *optional*, defaults to `None`): The timestep range within which the temporal attention computation can be skipped without a significant loss in quality. This is to be determined by the user based on the underlying model. The first value in the tuple is the lower bound and the second value is the upper bound. Typically, diffusion timesteps for denoising are in the reversed range of 0 to 1000 (i.e. denoising starts at timestep 1000 and ends at timestep 0). - low_frequency_weight_update_timestep_range (`Tuple[int, int]`, defaults to `(99, 901)`): + low_frequency_weight_update_timestep_range (`tuple[int, int]`, defaults to `(99, 901)`): The timestep range within which the low frequency weight scaling update is applied. The first value in the tuple is the lower bound and the second value is the upper bound of the timestep range. The callback function for the update is called only within this range. - high_frequency_weight_update_timestep_range (`Tuple[int, int]`, defaults to `(-1, 301)`): + high_frequency_weight_update_timestep_range (`tuple[int, int]`, defaults to `(-1, 301)`): The timestep range within which the high frequency weight scaling update is applied. The first value in the tuple is the lower bound and the second value is the upper bound of the timestep range. The callback function for the update is called only within this range. @@ -92,15 +92,15 @@ class FasterCacheConfig: Process the unconditional branch every `N` iterations. If this is set to `N`, the unconditional branch computation will be skipped `N - 1` times (i.e., cached unconditional branch states will be reused) before computing the new unconditional branch states again. - unconditional_batch_timestep_skip_range (`Tuple[float, float]`, defaults to `(-1, 641)`): + unconditional_batch_timestep_skip_range (`tuple[float, float]`, defaults to `(-1, 641)`): The timestep range within which the unconditional branch computation can be skipped without a significant loss in quality. This is to be determined by the user based on the underlying model. The first value in the tuple is the lower bound and the second value is the upper bound. - spatial_attention_block_identifiers (`Tuple[str, ...]`, defaults to `("blocks.*attn1", "transformer_blocks.*attn1", "single_transformer_blocks.*attn1")`): + spatial_attention_block_identifiers (`tuple[str, ...]`, defaults to `("blocks.*attn1", "transformer_blocks.*attn1", "single_transformer_blocks.*attn1")`): The identifiers to match the spatial attention blocks in the model. If the name of the block contains any of these identifiers, FasterCache will be applied to that block. This can either be the full layer names, partial layer names, or regex patterns. Matching will always be done using a regex match. - temporal_attention_block_identifiers (`Tuple[str, ...]`, defaults to `("temporal_transformer_blocks.*attn1",)`): + temporal_attention_block_identifiers (`tuple[str, ...]`, defaults to `("temporal_transformer_blocks.*attn1",)`): The identifiers to match the temporal attention blocks in the model. If the name of the block contains any of these identifiers, FasterCache will be applied to that block. This can either be the full layer names, partial layer names, or regex patterns. Matching will always be done using a regex match. @@ -123,7 +123,7 @@ class FasterCacheConfig: is_guidance_distilled (`bool`, defaults to `False`): Whether the model is guidance distilled or not. If the model is guidance distilled, FasterCache will not be applied at the denoiser-level to skip the unconditional branch computation (as there is none). - _unconditional_conditional_input_kwargs_identifiers (`List[str]`, defaults to `("hidden_states", "encoder_hidden_states", "timestep", "attention_mask", "encoder_attention_mask")`): + _unconditional_conditional_input_kwargs_identifiers (`list[str]`, defaults to `("hidden_states", "encoder_hidden_states", "timestep", "attention_mask", "encoder_attention_mask")`): The identifiers to match the input kwargs that contain the batchwise-concatenated unconditional and conditional inputs. If the name of the input kwargs contains any of these identifiers, FasterCache will split the inputs into unconditional and conditional branches. This must be a list of exact input kwargs @@ -133,14 +133,14 @@ class FasterCacheConfig: # In the paper and codebase, they hardcode these values to 2. However, it can be made configurable # after some testing. We default to 2 if these parameters are not provided. spatial_attention_block_skip_range: int = 2 - temporal_attention_block_skip_range: Optional[int] = None + temporal_attention_block_skip_range: int | None = None - spatial_attention_timestep_skip_range: Tuple[int, int] = (-1, 681) - temporal_attention_timestep_skip_range: Tuple[int, int] = (-1, 681) + spatial_attention_timestep_skip_range: tuple[int, int] = (-1, 681) + temporal_attention_timestep_skip_range: tuple[int, int] = (-1, 681) # Indicator functions for low/high frequency as mentioned in Equation 11 of the paper - low_frequency_weight_update_timestep_range: Tuple[int, int] = (99, 901) - high_frequency_weight_update_timestep_range: Tuple[int, int] = (-1, 301) + low_frequency_weight_update_timestep_range: tuple[int, int] = (99, 901) + high_frequency_weight_update_timestep_range: tuple[int, int] = (-1, 301) # ⍺1 and ⍺2 as mentioned in Equation 11 of the paper alpha_low_frequency: float = 1.1 @@ -148,10 +148,10 @@ class FasterCacheConfig: # n as described in CFG-Cache explanation in the paper - dependent on the model unconditional_batch_skip_range: int = 5 - unconditional_batch_timestep_skip_range: Tuple[int, int] = (-1, 641) + unconditional_batch_timestep_skip_range: tuple[int, int] = (-1, 641) - spatial_attention_block_identifiers: Tuple[str, ...] = _SPATIAL_ATTENTION_BLOCK_IDENTIFIERS - temporal_attention_block_identifiers: Tuple[str, ...] = _TEMPORAL_ATTENTION_BLOCK_IDENTIFIERS + spatial_attention_block_identifiers: tuple[str, ...] = _SPATIAL_ATTENTION_BLOCK_IDENTIFIERS + temporal_attention_block_identifiers: tuple[str, ...] = _TEMPORAL_ATTENTION_BLOCK_IDENTIFIERS attention_weight_callback: Callable[[torch.nn.Module], float] = None low_frequency_weight_callback: Callable[[torch.nn.Module], float] = None @@ -162,7 +162,7 @@ class FasterCacheConfig: current_timestep_callback: Callable[[], int] = None - _unconditional_conditional_input_kwargs_identifiers: List[str] = _UNCOND_COND_INPUT_KWARGS_IDENTIFIERS + _unconditional_conditional_input_kwargs_identifiers: list[str] = _UNCOND_COND_INPUT_KWARGS_IDENTIFIERS def __repr__(self) -> str: return ( @@ -209,7 +209,7 @@ class FasterCacheBlockState: def __init__(self) -> None: self.iteration: int = 0 self.batch_size: int = None - self.cache: Tuple[torch.Tensor, torch.Tensor] = None + self.cache: tuple[torch.Tensor, torch.Tensor] = None def reset(self): self.iteration = 0 @@ -223,10 +223,10 @@ class FasterCacheDenoiserHook(ModelHook): def __init__( self, unconditional_batch_skip_range: int, - unconditional_batch_timestep_skip_range: Tuple[int, int], + unconditional_batch_timestep_skip_range: tuple[int, int], tensor_format: str, is_guidance_distilled: bool, - uncond_cond_input_kwargs_identifiers: List[str], + uncond_cond_input_kwargs_identifiers: list[str], current_timestep_callback: Callable[[], int], low_frequency_weight_callback: Callable[[torch.nn.Module], torch.Tensor], high_frequency_weight_callback: Callable[[torch.nn.Module], torch.Tensor], @@ -252,7 +252,7 @@ def initialize_hook(self, module): return module @staticmethod - def _get_cond_input(input: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + def _get_cond_input(input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: # Note: this method assumes that the input tensor is batchwise-concatenated with unconditional inputs # followed by conditional inputs. _, cond = input.chunk(2, dim=0) @@ -371,7 +371,7 @@ class FasterCacheBlockHook(ModelHook): def __init__( self, block_skip_range: int, - timestep_skip_range: Tuple[int, int], + timestep_skip_range: tuple[int, int], is_guidance_distilled: bool, weight_callback: Callable[[torch.nn.Module], float], current_timestep_callback: Callable[[], int], diff --git a/src/diffusers/hooks/first_block_cache.py b/src/diffusers/hooks/first_block_cache.py index 862d44059301..685ccd383674 100644 --- a/src/diffusers/hooks/first_block_cache.py +++ b/src/diffusers/hooks/first_block_cache.py @@ -13,7 +13,6 @@ # limitations under the License. from dataclasses import dataclass -from typing import Tuple, Union import torch @@ -53,9 +52,9 @@ class FBCSharedBlockState(BaseState): def __init__(self) -> None: super().__init__() - self.head_block_output: Union[torch.Tensor, Tuple[torch.Tensor, ...]] = None + self.head_block_output: torch.Tensor | tuple[torch.Tensor, ...] = None self.head_block_residual: torch.Tensor = None - self.tail_block_residuals: Union[torch.Tensor, Tuple[torch.Tensor, ...]] = None + self.tail_block_residuals: torch.Tensor | tuple[torch.Tensor, ...] = None self.should_compute: bool = True def reset(self): diff --git a/src/diffusers/hooks/group_offloading.py b/src/diffusers/hooks/group_offloading.py index 47f1f4199615..891ac28455af 100644 --- a/src/diffusers/hooks/group_offloading.py +++ b/src/diffusers/hooks/group_offloading.py @@ -17,7 +17,7 @@ from contextlib import contextmanager, nullcontext from dataclasses import dataclass, replace from enum import Enum -from typing import Dict, List, Optional, Set, Tuple, Union +from typing import Set import safetensors.torch import torch @@ -56,31 +56,31 @@ class GroupOffloadingConfig: non_blocking: bool record_stream: bool low_cpu_mem_usage: bool - num_blocks_per_group: Optional[int] = None - offload_to_disk_path: Optional[str] = None - stream: Optional[Union[torch.cuda.Stream, torch.Stream]] = None - block_modules: Optional[List[str]] = None - exclude_kwargs: Optional[List[str]] = None - module_prefix: Optional[str] = "" + num_blocks_per_group: int | None = None + offload_to_disk_path: str | None = None + stream: torch.cuda.Stream | torch.Stream | None = None + block_modules: list[str] | None = None + exclude_kwargs: list[str] | None = None + module_prefix: str = "" class ModuleGroup: def __init__( self, - modules: List[torch.nn.Module], + modules: list[torch.nn.Module], offload_device: torch.device, onload_device: torch.device, offload_leader: torch.nn.Module, - onload_leader: Optional[torch.nn.Module] = None, - parameters: Optional[List[torch.nn.Parameter]] = None, - buffers: Optional[List[torch.Tensor]] = None, + onload_leader: torch.nn.Module | None = None, + parameters: list[torch.nn.Parameter] | None = None, + buffers: list[torch.Tensor] | None = None, non_blocking: bool = False, - stream: Union[torch.cuda.Stream, torch.Stream, None] = None, - record_stream: Optional[bool] = False, + stream: torch.cuda.Stream | torch.Stream | None = None, + record_stream: bool | None = False, low_cpu_mem_usage: bool = False, onload_self: bool = True, - offload_to_disk_path: Optional[str] = None, - group_id: Optional[Union[int, str]] = None, + offload_to_disk_path: str | None = None, + group_id: int | str | None = None, ) -> None: self.modules = modules self.offload_device = offload_device @@ -287,7 +287,7 @@ class GroupOffloadingHook(ModelHook): def __init__(self, group: ModuleGroup, *, config: GroupOffloadingConfig) -> None: self.group = group - self.next_group: Optional[ModuleGroup] = None + self.next_group: ModuleGroup | None = None self.config = config def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module: @@ -307,6 +307,17 @@ def pre_forward(self, module: torch.nn.Module, *args, **kwargs): if self.group.onload_leader == module: if self.group.onload_self: self.group.onload_() + else: + # onload_self=False means this group relies on prefetching from a previous group. + # However, for conditionally-executed modules (e.g. patch_short/patch_mid/patch_long in Helios), + # the prefetch chain may not cover them if they were absent during the first forward pass + # when the execution order was traced. In that case, their weights remain on offload_device, + # so we fall back to a synchronous onload here. + params = [p for m in self.group.modules for p in m.parameters()] + list(self.group.parameters) + if params and params[0].device == self.group.offload_device: + self.group.onload_() + if self.group.stream is not None: + self.group.stream.synchronize() should_onload_next_group = self.next_group is not None and not self.next_group.onload_self if should_onload_next_group: @@ -359,7 +370,7 @@ class LazyPrefetchGroupOffloadingHook(ModelHook): _is_stateful = False def __init__(self): - self.execution_order: List[Tuple[str, torch.nn.Module]] = [] + self.execution_order: list[tuple[str, torch.nn.Module]] = [] self._layer_execution_tracker_module_names = set() def initialize_hook(self, module): @@ -463,17 +474,17 @@ def pre_forward(self, module, *args, **kwargs): def apply_group_offloading( module: torch.nn.Module, - onload_device: Union[str, torch.device], - offload_device: Union[str, torch.device] = torch.device("cpu"), - offload_type: Union[str, GroupOffloadingType] = "block_level", - num_blocks_per_group: Optional[int] = None, + onload_device: str | torch.device, + offload_device: str | torch.device = torch.device("cpu"), + offload_type: str | GroupOffloadingType = "block_level", + num_blocks_per_group: int | None = None, non_blocking: bool = False, use_stream: bool = False, record_stream: bool = False, low_cpu_mem_usage: bool = False, - offload_to_disk_path: Optional[str] = None, - block_modules: Optional[List[str]] = None, - exclude_kwargs: Optional[List[str]] = None, + offload_to_disk_path: str | None = None, + block_modules: list[str] | None = None, + exclude_kwargs: list[str] | None = None, ) -> None: r""" Applies group offloading to the internal layers of a torch.nn.Module. To understand what group offloading is, and @@ -531,10 +542,10 @@ def apply_group_offloading( If True, the CPU memory usage is minimized by pinning tensors on-the-fly instead of pre-pinning them. This option only matters when using streamed CPU offloading (i.e. `use_stream=True`). This can be useful when the CPU memory is a bottleneck but may counteract the benefits of using streams. - block_modules (`List[str]`, *optional*): + block_modules (`list[str]`, *optional*): List of module names that should be treated as blocks for offloading. If provided, only these modules will be considered for block-level offloading. If not provided, the default block detection logic will be used. - exclude_kwargs (`List[str]`, *optional*): + exclude_kwargs (`list[str]`, *optional*): List of kwarg keys that should not be processed by send_to_device. This is useful for mutable state like caching lists that need to maintain their object identity across forward passes. If not provided, will be inferred from the module's `_skip_keys` attribute if it exists. @@ -844,7 +855,7 @@ def _apply_lazy_group_offloading_hook( def _gather_parameters_with_no_group_offloading_parent( module: torch.nn.Module, modules_with_group_offloading: Set[str] -) -> List[torch.nn.Parameter]: +) -> list[torch.nn.Parameter]: parameters = [] for name, parameter in module.named_parameters(): has_parent_with_group_offloading = False @@ -862,7 +873,7 @@ def _gather_parameters_with_no_group_offloading_parent( def _gather_buffers_with_no_group_offloading_parent( module: torch.nn.Module, modules_with_group_offloading: Set[str] -) -> List[torch.Tensor]: +) -> list[torch.Tensor]: buffers = [] for name, buffer in module.named_buffers(): has_parent_with_group_offloading = False @@ -878,7 +889,7 @@ def _gather_buffers_with_no_group_offloading_parent( return buffers -def _find_parent_module_in_module_dict(name: str, module_dict: Dict[str, torch.nn.Module]) -> str: +def _find_parent_module_in_module_dict(name: str, module_dict: dict[str, torch.nn.Module]) -> str: atoms = name.split(".") while len(atoms) > 0: parent_name = ".".join(atoms) @@ -902,7 +913,7 @@ def _raise_error_if_accelerate_model_or_sequential_hook_present(module: torch.nn ) -def _get_top_level_group_offload_hook(module: torch.nn.Module) -> Optional[GroupOffloadingHook]: +def _get_top_level_group_offload_hook(module: torch.nn.Module) -> GroupOffloadingHook | None: for submodule in module.modules(): if hasattr(submodule, "_diffusers_hook"): group_offloading_hook = submodule._diffusers_hook.get_hook(_GROUP_OFFLOADING) diff --git a/src/diffusers/hooks/hooks.py b/src/diffusers/hooks/hooks.py index 6e097e5882a0..2d575b85427c 100644 --- a/src/diffusers/hooks/hooks.py +++ b/src/diffusers/hooks/hooks.py @@ -13,7 +13,7 @@ # limitations under the License. import functools -from typing import Any, Dict, Optional, Tuple +from typing import Any import torch @@ -86,19 +86,19 @@ def deinitalize_hook(self, module: torch.nn.Module) -> torch.nn.Module: """ return module - def pre_forward(self, module: torch.nn.Module, *args, **kwargs) -> Tuple[Tuple[Any], Dict[str, Any]]: + def pre_forward(self, module: torch.nn.Module, *args, **kwargs) -> tuple[tuple[Any], dict[str, Any]]: r""" Hook that is executed just before the forward method of the model. Args: module (`torch.nn.Module`): The module whose forward pass will be executed just after this event. - args (`Tuple[Any]`): + args (`tuple[Any]`): The positional arguments passed to the module. - kwargs (`Dict[Str, Any]`): + kwargs (`dict[Str, Any]`): The keyword arguments passed to the module. Returns: - `Tuple[Tuple[Any], Dict[Str, Any]]`: + `tuple[tuple[Any], dict[Str, Any]]`: A tuple with the treated `args` and `kwargs`. """ return args, kwargs @@ -168,7 +168,7 @@ class HookRegistry: def __init__(self, module_ref: torch.nn.Module) -> None: super().__init__() - self.hooks: Dict[str, ModelHook] = {} + self.hooks: dict[str, ModelHook] = {} self._module_ref = module_ref self._hook_order = [] @@ -214,7 +214,7 @@ def new_forward(module, *args, **kwargs): self._hook_order.append(name) self._fn_refs.append(fn_ref) - def get_hook(self, name: str) -> Optional[ModelHook]: + def get_hook(self, name: str) -> ModelHook | None: return self.hooks.get(name, None) def remove_hook(self, name: str, recurse: bool = True) -> None: @@ -265,7 +265,7 @@ def check_if_exists_or_initialize(cls, module: torch.nn.Module) -> "HookRegistry module._diffusers_hook = cls(module) return module._diffusers_hook - def _set_context(self, name: Optional[str] = None) -> None: + def _set_context(self, name: str | None = None) -> None: for hook_name in reversed(self._hook_order): hook = self.hooks[hook_name] if hook._is_stateful: diff --git a/src/diffusers/hooks/layer_skip.py b/src/diffusers/hooks/layer_skip.py index 0ce02e987d09..112edfa2f79b 100644 --- a/src/diffusers/hooks/layer_skip.py +++ b/src/diffusers/hooks/layer_skip.py @@ -14,7 +14,7 @@ import math from dataclasses import asdict, dataclass -from typing import Callable, List, Optional +from typing import Callable import torch @@ -43,7 +43,7 @@ class LayerSkipConfig: Configuration for skipping internal transformer blocks when executing a transformer model. Args: - indices (`List[int]`): + indices (`list[int]`): The indices of the layer to skip. This is typically the first layer in the transformer block. fqn (`str`, defaults to `"auto"`): The fully qualified name identifying the stack of transformer blocks. Typically, this is @@ -63,7 +63,7 @@ class LayerSkipConfig: skipped layers are fully retained, which is equivalent to not skipping any layers. """ - indices: List[int] + indices: list[int] fqn: str = "auto" skip_attention: bool = True skip_attention_scores: bool = False @@ -196,7 +196,7 @@ def apply_layer_skip(module: torch.nn.Module, config: LayerSkipConfig) -> None: _apply_layer_skip_hook(module, config) -def _apply_layer_skip_hook(module: torch.nn.Module, config: LayerSkipConfig, name: Optional[str] = None) -> None: +def _apply_layer_skip_hook(module: torch.nn.Module, config: LayerSkipConfig, name: str | None = None) -> None: name = name or _LAYER_SKIP_HOOK if config.skip_attention and config.skip_attention_scores: diff --git a/src/diffusers/hooks/layerwise_casting.py b/src/diffusers/hooks/layerwise_casting.py index a036ad37dc2f..1edff7805fcc 100644 --- a/src/diffusers/hooks/layerwise_casting.py +++ b/src/diffusers/hooks/layerwise_casting.py @@ -13,7 +13,7 @@ # limitations under the License. import re -from typing import Optional, Tuple, Type, Union +from typing import Type import torch @@ -102,8 +102,8 @@ def apply_layerwise_casting( module: torch.nn.Module, storage_dtype: torch.dtype, compute_dtype: torch.dtype, - skip_modules_pattern: Union[str, Tuple[str, ...]] = "auto", - skip_modules_classes: Optional[Tuple[Type[torch.nn.Module], ...]] = None, + skip_modules_pattern: str | tuple[str, ...] = "auto", + skip_modules_classes: tuple[Type[torch.nn.Module], ...] | None = None, non_blocking: bool = False, ) -> None: r""" @@ -137,12 +137,12 @@ def apply_layerwise_casting( The dtype to cast the module to before/after the forward pass for storage. compute_dtype (`torch.dtype`): The dtype to cast the module to during the forward pass for computation. - skip_modules_pattern (`Tuple[str, ...]`, defaults to `"auto"`): + skip_modules_pattern (`tuple[str, ...]`, defaults to `"auto"`): A list of patterns to match the names of the modules to skip during the layerwise casting process. If set to `"auto"`, the default patterns are used. If set to `None`, no modules are skipped. If set to `None` alongside `skip_modules_classes` being `None`, the layerwise casting is applied directly to the module instead of its internal submodules. - skip_modules_classes (`Tuple[Type[torch.nn.Module], ...]`, defaults to `None`): + skip_modules_classes (`tuple[Type[torch.nn.Module], ...]`, defaults to `None`): A list of module classes to skip during the layerwise casting process. non_blocking (`bool`, defaults to `False`): If `True`, the weight casting operations are non-blocking. @@ -169,8 +169,8 @@ def _apply_layerwise_casting( module: torch.nn.Module, storage_dtype: torch.dtype, compute_dtype: torch.dtype, - skip_modules_pattern: Optional[Tuple[str, ...]] = None, - skip_modules_classes: Optional[Tuple[Type[torch.nn.Module], ...]] = None, + skip_modules_pattern: tuple[str, ...] | None = None, + skip_modules_classes: tuple[Type[torch.nn.Module], ...] | None = None, non_blocking: bool = False, _prefix: str = "", ) -> None: diff --git a/src/diffusers/hooks/mag_cache.py b/src/diffusers/hooks/mag_cache.py new file mode 100644 index 000000000000..d28cd2d793b6 --- /dev/null +++ b/src/diffusers/hooks/mag_cache.py @@ -0,0 +1,468 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch + +from ..utils import get_logger +from ..utils.torch_utils import unwrap_module +from ._common import _ALL_TRANSFORMER_BLOCK_IDENTIFIERS +from ._helpers import TransformerBlockRegistry +from .hooks import BaseState, HookRegistry, ModelHook, StateManager + + +logger = get_logger(__name__) # pylint: disable=invalid-name + +_MAG_CACHE_LEADER_BLOCK_HOOK = "mag_cache_leader_block_hook" +_MAG_CACHE_BLOCK_HOOK = "mag_cache_block_hook" + +# Default Mag Ratios for Flux models (Dev/Schnell) are provided for convenience. +# Users must explicitly pass these to the config if using Flux. +# Reference: https://github.com/Zehong-Ma/MagCache +FLUX_MAG_RATIOS = torch.tensor( + [1.0] + + [ + 1.21094, + 1.11719, + 1.07812, + 1.0625, + 1.03906, + 1.03125, + 1.03906, + 1.02344, + 1.03125, + 1.02344, + 0.98047, + 1.01562, + 1.00781, + 1.0, + 1.00781, + 1.0, + 1.00781, + 1.0, + 1.0, + 0.99609, + 0.99609, + 0.98047, + 0.98828, + 0.96484, + 0.95703, + 0.93359, + 0.89062, + ] +) + + +def nearest_interp(src_array: torch.Tensor, target_length: int) -> torch.Tensor: + """ + Interpolate the source array to the target length using nearest neighbor interpolation. + """ + src_length = len(src_array) + if target_length == 1: + return src_array[-1:] + + scale = (src_length - 1) / (target_length - 1) + grid = torch.arange(target_length, device=src_array.device, dtype=torch.float32) + mapped_indices = torch.round(grid * scale).long() + return src_array[mapped_indices] + + +@dataclass +class MagCacheConfig: + r""" + Configuration for [MagCache](https://github.com/Zehong-Ma/MagCache). + + Args: + threshold (`float`, defaults to `0.06`): + The threshold for the accumulated error. If the accumulated error is below this threshold, the block + computation is skipped. A higher threshold allows for more aggressive skipping (faster) but may degrade + quality. + max_skip_steps (`int`, defaults to `3`): + The maximum number of consecutive steps that can be skipped (K in the paper). + retention_ratio (`float`, defaults to `0.2`): + The fraction of initial steps during which skipping is disabled to ensure stability. For example, if + `num_inference_steps` is 28 and `retention_ratio` is 0.2, the first 6 steps will never be skipped. + num_inference_steps (`int`, defaults to `28`): + The number of inference steps used in the pipeline. This is required to interpolate `mag_ratios` correctly. + mag_ratios (`torch.Tensor`, *optional*): + The pre-computed magnitude ratios for the model. These are checkpoint-dependent. If not provided, you must + set `calibrate=True` to calculate them for your specific model. For Flux models, you can use + `diffusers.hooks.mag_cache.FLUX_MAG_RATIOS`. + calibrate (`bool`, defaults to `False`): + If True, enables calibration mode. In this mode, no blocks are skipped. Instead, the hook calculates the + magnitude ratios for the current run and logs them at the end. Use this to obtain `mag_ratios` for new + models or schedulers. + """ + + threshold: float = 0.06 + max_skip_steps: int = 3 + retention_ratio: float = 0.2 + num_inference_steps: int = 28 + mag_ratios: Optional[Union[torch.Tensor, List[float]]] = None + calibrate: bool = False + + def __post_init__(self): + # User MUST provide ratios OR enable calibration. + if self.mag_ratios is None and not self.calibrate: + raise ValueError( + " `mag_ratios` must be provided for MagCache inference because these ratios are model-dependent.\n" + "To get them for your model:\n" + "1. Initialize `MagCacheConfig(calibrate=True, ...)`\n" + "2. Run inference on your model once.\n" + "3. Copy the printed ratios array and pass it to `mag_ratios` in the config.\n" + "For Flux models, you can import `FLUX_MAG_RATIOS` from `diffusers.hooks.mag_cache`." + ) + + if not self.calibrate and self.mag_ratios is not None: + if not torch.is_tensor(self.mag_ratios): + self.mag_ratios = torch.tensor(self.mag_ratios) + + if len(self.mag_ratios) != self.num_inference_steps: + logger.debug( + f"Interpolating mag_ratios from length {len(self.mag_ratios)} to {self.num_inference_steps}" + ) + self.mag_ratios = nearest_interp(self.mag_ratios, self.num_inference_steps) + + +class MagCacheState(BaseState): + def __init__(self) -> None: + super().__init__() + # Cache for the residual (output - input) from the *previous* timestep + self.previous_residual: torch.Tensor = None + + # State inputs/outputs for the current forward pass + self.head_block_input: Union[torch.Tensor, Tuple[torch.Tensor, ...]] = None + self.should_compute: bool = True + + # MagCache accumulators + self.accumulated_ratio: float = 1.0 + self.accumulated_err: float = 0.0 + self.accumulated_steps: int = 0 + + # Current step counter (timestep index) + self.step_index: int = 0 + + # Calibration storage + self.calibration_ratios: List[float] = [] + + def reset(self): + self.previous_residual = None + self.should_compute = True + self.accumulated_ratio = 1.0 + self.accumulated_err = 0.0 + self.accumulated_steps = 0 + self.step_index = 0 + self.calibration_ratios = [] + + +class MagCacheHeadHook(ModelHook): + _is_stateful = True + + def __init__(self, state_manager: StateManager, config: MagCacheConfig): + self.state_manager = state_manager + self.config = config + self._metadata = None + + def initialize_hook(self, module): + unwrapped_module = unwrap_module(module) + self._metadata = TransformerBlockRegistry.get(unwrapped_module.__class__) + return module + + @torch.compiler.disable + def new_forward(self, module: torch.nn.Module, *args, **kwargs): + if self.state_manager._current_context is None: + self.state_manager.set_context("inference") + + arg_name = self._metadata.hidden_states_argument_name + hidden_states = self._metadata._get_parameter_from_args_kwargs(arg_name, args, kwargs) + + state: MagCacheState = self.state_manager.get_state() + state.head_block_input = hidden_states + + should_compute = True + + if self.config.calibrate: + # Never skip during calibration + should_compute = True + else: + # MagCache Logic + current_step = state.step_index + if current_step >= len(self.config.mag_ratios): + current_scale = 1.0 + else: + current_scale = self.config.mag_ratios[current_step] + + retention_step = int(self.config.retention_ratio * self.config.num_inference_steps + 0.5) + + if current_step >= retention_step: + state.accumulated_ratio *= current_scale + state.accumulated_steps += 1 + state.accumulated_err += abs(1.0 - state.accumulated_ratio) + + if ( + state.previous_residual is not None + and state.accumulated_err <= self.config.threshold + and state.accumulated_steps <= self.config.max_skip_steps + ): + should_compute = False + else: + state.accumulated_ratio = 1.0 + state.accumulated_steps = 0 + state.accumulated_err = 0.0 + + state.should_compute = should_compute + + if not should_compute: + logger.debug(f"MagCache: Skipping step {state.step_index}") + # Apply MagCache: Output = Input + Previous Residual + + output = hidden_states + res = state.previous_residual + + if res.device != output.device: + res = res.to(output.device) + + # Attempt to apply residual handling shape mismatches (e.g., text+image vs image only) + if res.shape == output.shape: + output = output + res + elif ( + output.ndim == 3 + and res.ndim == 3 + and output.shape[0] == res.shape[0] + and output.shape[2] == res.shape[2] + ): + # Assuming concatenation where image part is at the end (standard in Flux/SD3) + diff = output.shape[1] - res.shape[1] + if diff > 0: + output = output.clone() + output[:, diff:, :] = output[:, diff:, :] + res + else: + logger.warning( + f"MagCache: Dimension mismatch. Input {output.shape}, Residual {res.shape}. " + "Cannot apply residual safely. Returning input without residual." + ) + else: + logger.warning( + f"MagCache: Dimension mismatch. Input {output.shape}, Residual {res.shape}. " + "Cannot apply residual safely. Returning input without residual." + ) + + if self._metadata.return_encoder_hidden_states_index is not None: + original_encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs( + "encoder_hidden_states", args, kwargs + ) + max_idx = max( + self._metadata.return_hidden_states_index, self._metadata.return_encoder_hidden_states_index + ) + ret_list = [None] * (max_idx + 1) + ret_list[self._metadata.return_hidden_states_index] = output + ret_list[self._metadata.return_encoder_hidden_states_index] = original_encoder_hidden_states + return tuple(ret_list) + else: + return output + + else: + # Compute original forward + output = self.fn_ref.original_forward(*args, **kwargs) + return output + + def reset_state(self, module): + self.state_manager.reset() + return module + + +class MagCacheBlockHook(ModelHook): + def __init__(self, state_manager: StateManager, is_tail: bool = False, config: MagCacheConfig = None): + super().__init__() + self.state_manager = state_manager + self.is_tail = is_tail + self.config = config + self._metadata = None + + def initialize_hook(self, module): + unwrapped_module = unwrap_module(module) + self._metadata = TransformerBlockRegistry.get(unwrapped_module.__class__) + return module + + @torch.compiler.disable + def new_forward(self, module: torch.nn.Module, *args, **kwargs): + if self.state_manager._current_context is None: + self.state_manager.set_context("inference") + state: MagCacheState = self.state_manager.get_state() + + if not state.should_compute: + arg_name = self._metadata.hidden_states_argument_name + hidden_states = self._metadata._get_parameter_from_args_kwargs(arg_name, args, kwargs) + + if self.is_tail: + # Still need to advance step index even if we skip + self._advance_step(state) + + if self._metadata.return_encoder_hidden_states_index is not None: + encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs( + "encoder_hidden_states", args, kwargs + ) + max_idx = max( + self._metadata.return_hidden_states_index, self._metadata.return_encoder_hidden_states_index + ) + ret_list = [None] * (max_idx + 1) + ret_list[self._metadata.return_hidden_states_index] = hidden_states + ret_list[self._metadata.return_encoder_hidden_states_index] = encoder_hidden_states + return tuple(ret_list) + + return hidden_states + + output = self.fn_ref.original_forward(*args, **kwargs) + + if self.is_tail: + # Calculate residual for next steps + if isinstance(output, tuple): + out_hidden = output[self._metadata.return_hidden_states_index] + else: + out_hidden = output + + in_hidden = state.head_block_input + + if in_hidden is None: + return output + + # Determine residual + if out_hidden.shape == in_hidden.shape: + residual = out_hidden - in_hidden + elif out_hidden.ndim == 3 and in_hidden.ndim == 3 and out_hidden.shape[2] == in_hidden.shape[2]: + diff = in_hidden.shape[1] - out_hidden.shape[1] + if diff == 0: + residual = out_hidden - in_hidden + else: + residual = out_hidden - in_hidden # Fallback to matching tail + else: + # Fallback for completely mismatched shapes + residual = out_hidden + + if self.config.calibrate: + self._perform_calibration_step(state, residual) + + state.previous_residual = residual + self._advance_step(state) + + return output + + def _perform_calibration_step(self, state: MagCacheState, current_residual: torch.Tensor): + if state.previous_residual is None: + # First step has no previous residual to compare against. + # log 1.0 as a neutral starting point. + ratio = 1.0 + else: + # MagCache Calibration Formula: mean(norm(curr) / norm(prev)) + # norm(dim=-1) gives magnitude of each token vector + curr_norm = torch.linalg.norm(current_residual.float(), dim=-1) + prev_norm = torch.linalg.norm(state.previous_residual.float(), dim=-1) + + # Avoid division by zero + ratio = (curr_norm / (prev_norm + 1e-8)).mean().item() + + state.calibration_ratios.append(ratio) + + def _advance_step(self, state: MagCacheState): + state.step_index += 1 + if state.step_index >= self.config.num_inference_steps: + # End of inference loop + if self.config.calibrate: + print("\n[MagCache] Calibration Complete. Copy these values to MagCacheConfig(mag_ratios=...):") + print(f"{state.calibration_ratios}\n") + logger.info(f"MagCache Calibration Results: {state.calibration_ratios}") + + # Reset state + state.step_index = 0 + state.accumulated_ratio = 1.0 + state.accumulated_steps = 0 + state.accumulated_err = 0.0 + state.previous_residual = None + state.calibration_ratios = [] + + +def apply_mag_cache(module: torch.nn.Module, config: MagCacheConfig) -> None: + """ + Applies MagCache to a given module (typically a Transformer). + + Args: + module (`torch.nn.Module`): + The module to apply MagCache to. + config (`MagCacheConfig`): + The configuration for MagCache. + """ + # Initialize registry on the root module so the Pipeline can set context. + HookRegistry.check_if_exists_or_initialize(module) + + state_manager = StateManager(MagCacheState, (), {}) + remaining_blocks = [] + + for name, submodule in module.named_children(): + if name not in _ALL_TRANSFORMER_BLOCK_IDENTIFIERS or not isinstance(submodule, torch.nn.ModuleList): + continue + for index, block in enumerate(submodule): + remaining_blocks.append((f"{name}.{index}", block)) + + if not remaining_blocks: + logger.warning("MagCache: No transformer blocks found to apply hooks.") + return + + # Handle single-block models + if len(remaining_blocks) == 1: + name, block = remaining_blocks[0] + logger.info(f"MagCache: Applying Head+Tail Hooks to single block '{name}'") + _apply_mag_cache_block_hook(block, state_manager, config, is_tail=True) + _apply_mag_cache_head_hook(block, state_manager, config) + return + + head_block_name, head_block = remaining_blocks.pop(0) + tail_block_name, tail_block = remaining_blocks.pop(-1) + + logger.info(f"MagCache: Applying Head Hook to {head_block_name}") + _apply_mag_cache_head_hook(head_block, state_manager, config) + + for name, block in remaining_blocks: + _apply_mag_cache_block_hook(block, state_manager, config) + + logger.info(f"MagCache: Applying Tail Hook to {tail_block_name}") + _apply_mag_cache_block_hook(tail_block, state_manager, config, is_tail=True) + + +def _apply_mag_cache_head_hook(block: torch.nn.Module, state_manager: StateManager, config: MagCacheConfig) -> None: + registry = HookRegistry.check_if_exists_or_initialize(block) + + # Automatically remove existing hook to allow re-application (e.g. switching modes) + if registry.get_hook(_MAG_CACHE_LEADER_BLOCK_HOOK) is not None: + registry.remove_hook(_MAG_CACHE_LEADER_BLOCK_HOOK) + + hook = MagCacheHeadHook(state_manager, config) + registry.register_hook(hook, _MAG_CACHE_LEADER_BLOCK_HOOK) + + +def _apply_mag_cache_block_hook( + block: torch.nn.Module, + state_manager: StateManager, + config: MagCacheConfig, + is_tail: bool = False, +) -> None: + registry = HookRegistry.check_if_exists_or_initialize(block) + + # Automatically remove existing hook to allow re-application + if registry.get_hook(_MAG_CACHE_BLOCK_HOOK) is not None: + registry.remove_hook(_MAG_CACHE_BLOCK_HOOK) + + hook = MagCacheBlockHook(state_manager, is_tail, config) + registry.register_hook(hook, _MAG_CACHE_BLOCK_HOOK) diff --git a/src/diffusers/hooks/pyramid_attention_broadcast.py b/src/diffusers/hooks/pyramid_attention_broadcast.py index 12d6aa0616e9..ed5bd24dea01 100644 --- a/src/diffusers/hooks/pyramid_attention_broadcast.py +++ b/src/diffusers/hooks/pyramid_attention_broadcast.py @@ -14,7 +14,7 @@ import re from dataclasses import dataclass -from typing import Any, Callable, Optional, Tuple, Union +from typing import Any, Callable import torch @@ -54,34 +54,34 @@ class PyramidAttentionBroadcastConfig: The number of times a specific cross-attention broadcast is skipped before computing the attention states to re-use. If this is set to the value `N`, the attention computation will be skipped `N - 1` times (i.e., old attention states will be reused) before computing the new attention states again. - spatial_attention_timestep_skip_range (`Tuple[int, int]`, defaults to `(100, 800)`): + spatial_attention_timestep_skip_range (`tuple[int, int]`, defaults to `(100, 800)`): The range of timesteps to skip in the spatial attention layer. The attention computations will be conditionally skipped if the current timestep is within the specified range. - temporal_attention_timestep_skip_range (`Tuple[int, int]`, defaults to `(100, 800)`): + temporal_attention_timestep_skip_range (`tuple[int, int]`, defaults to `(100, 800)`): The range of timesteps to skip in the temporal attention layer. The attention computations will be conditionally skipped if the current timestep is within the specified range. - cross_attention_timestep_skip_range (`Tuple[int, int]`, defaults to `(100, 800)`): + cross_attention_timestep_skip_range (`tuple[int, int]`, defaults to `(100, 800)`): The range of timesteps to skip in the cross-attention layer. The attention computations will be conditionally skipped if the current timestep is within the specified range. - spatial_attention_block_identifiers (`Tuple[str, ...]`): + spatial_attention_block_identifiers (`tuple[str, ...]`): The identifiers to match against the layer names to determine if the layer is a spatial attention layer. - temporal_attention_block_identifiers (`Tuple[str, ...]`): + temporal_attention_block_identifiers (`tuple[str, ...]`): The identifiers to match against the layer names to determine if the layer is a temporal attention layer. - cross_attention_block_identifiers (`Tuple[str, ...]`): + cross_attention_block_identifiers (`tuple[str, ...]`): The identifiers to match against the layer names to determine if the layer is a cross-attention layer. """ - spatial_attention_block_skip_range: Optional[int] = None - temporal_attention_block_skip_range: Optional[int] = None - cross_attention_block_skip_range: Optional[int] = None + spatial_attention_block_skip_range: int | None = None + temporal_attention_block_skip_range: int | None = None + cross_attention_block_skip_range: int | None = None - spatial_attention_timestep_skip_range: Tuple[int, int] = (100, 800) - temporal_attention_timestep_skip_range: Tuple[int, int] = (100, 800) - cross_attention_timestep_skip_range: Tuple[int, int] = (100, 800) + spatial_attention_timestep_skip_range: tuple[int, int] = (100, 800) + temporal_attention_timestep_skip_range: tuple[int, int] = (100, 800) + cross_attention_timestep_skip_range: tuple[int, int] = (100, 800) - spatial_attention_block_identifiers: Tuple[str, ...] = _SPATIAL_TRANSFORMER_BLOCK_IDENTIFIERS - temporal_attention_block_identifiers: Tuple[str, ...] = _TEMPORAL_TRANSFORMER_BLOCK_IDENTIFIERS - cross_attention_block_identifiers: Tuple[str, ...] = _CROSS_TRANSFORMER_BLOCK_IDENTIFIERS + spatial_attention_block_identifiers: tuple[str, ...] = _SPATIAL_TRANSFORMER_BLOCK_IDENTIFIERS + temporal_attention_block_identifiers: tuple[str, ...] = _TEMPORAL_TRANSFORMER_BLOCK_IDENTIFIERS + cross_attention_block_identifiers: tuple[str, ...] = _CROSS_TRANSFORMER_BLOCK_IDENTIFIERS current_timestep_callback: Callable[[], int] = None @@ -141,7 +141,7 @@ class PyramidAttentionBroadcastHook(ModelHook): _is_stateful = True def __init__( - self, timestep_skip_range: Tuple[int, int], block_skip_range: int, current_timestep_callback: Callable[[], int] + self, timestep_skip_range: tuple[int, int], block_skip_range: int, current_timestep_callback: Callable[[], int] ) -> None: super().__init__() @@ -191,7 +191,7 @@ def apply_pyramid_attention_broadcast(module: torch.nn.Module, config: PyramidAt Args: module (`torch.nn.Module`): The module to apply Pyramid Attention Broadcast to. - config (`Optional[PyramidAttentionBroadcastConfig]`, `optional`, defaults to `None`): + config (`PyramidAttentionBroadcastConfig | None`, `optional`, defaults to `None`): The configuration to use for Pyramid Attention Broadcast. Example: @@ -288,8 +288,8 @@ def _apply_pyramid_attention_broadcast_on_attention_class( def _apply_pyramid_attention_broadcast_hook( - module: Union[Attention, MochiAttention], - timestep_skip_range: Tuple[int, int], + module: Attention | MochiAttention, + timestep_skip_range: tuple[int, int], block_skip_range: int, current_timestep_callback: Callable[[], int], ): @@ -299,7 +299,7 @@ def _apply_pyramid_attention_broadcast_hook( Args: module (`torch.nn.Module`): The module to apply Pyramid Attention Broadcast to. - timestep_skip_range (`Tuple[int, int]`): + timestep_skip_range (`tuple[int, int]`): The range of timesteps to skip in the attention layer. The attention computations will be conditionally skipped if the current timestep is within the specified range. block_skip_range (`int`): diff --git a/src/diffusers/hooks/smoothed_energy_guidance_utils.py b/src/diffusers/hooks/smoothed_energy_guidance_utils.py index 622f60764762..f413b6376d7c 100644 --- a/src/diffusers/hooks/smoothed_energy_guidance_utils.py +++ b/src/diffusers/hooks/smoothed_energy_guidance_utils.py @@ -14,7 +14,6 @@ import math from dataclasses import asdict, dataclass -from typing import List, Optional import torch import torch.nn.functional as F @@ -35,21 +34,21 @@ class SmoothedEnergyGuidanceConfig: Configuration for skipping internal transformer blocks when executing a transformer model. Args: - indices (`List[int]`): + indices (`list[int]`): The indices of the layer to skip. This is typically the first layer in the transformer block. fqn (`str`, defaults to `"auto"`): The fully qualified name identifying the stack of transformer blocks. Typically, this is `transformer_blocks`, `single_transformer_blocks`, `blocks`, `layers`, or `temporal_transformer_blocks`. For automatic detection, set this to `"auto"`. "auto" only works on DiT models. For UNet models, you must provide the correct fqn. - _query_proj_identifiers (`List[str]`, defaults to `None`): + _query_proj_identifiers (`list[str]`, defaults to `None`): The identifiers for the query projection layers. Typically, these are `to_q`, `query`, or `q_proj`. If `None`, `to_q` is used by default. """ - indices: List[int] + indices: list[int] fqn: str = "auto" - _query_proj_identifiers: List[str] = None + _query_proj_identifiers: list[str] = None def to_dict(self): return asdict(self) @@ -73,7 +72,7 @@ def post_forward(self, module: torch.nn.Module, output: torch.Tensor) -> torch.T def _apply_smoothed_energy_guidance_hook( - module: torch.nn.Module, config: SmoothedEnergyGuidanceConfig, blur_sigma: float, name: Optional[str] = None + module: torch.nn.Module, config: SmoothedEnergyGuidanceConfig, blur_sigma: float, name: str | None = None ) -> None: name = name or _SMOOTHED_ENERGY_GUIDANCE_HOOK diff --git a/src/diffusers/hooks/taylorseer_cache.py b/src/diffusers/hooks/taylorseer_cache.py index 7cad9f4fa161..303155105e71 100644 --- a/src/diffusers/hooks/taylorseer_cache.py +++ b/src/diffusers/hooks/taylorseer_cache.py @@ -1,7 +1,6 @@ import math import re from dataclasses import dataclass -from typing import Dict, List, Optional, Tuple import torch import torch.nn as nn @@ -51,12 +50,12 @@ class TaylorSeerCacheConfig: Data type used for storing and computing Taylor series factors. Lower precision reduces memory but may affect stability; higher precision improves accuracy at the cost of more memory. - skip_predict_identifiers (`List[str]`, *optional*, defaults to `None`): + skip_predict_identifiers (`list[str]`, *optional*, defaults to `None`): Regex patterns (using `re.fullmatch`) for module names to place as "skip" in "cache" mode. In this mode, the module computes fully during initial or refresh steps but returns a zero tensor (matching recorded shape) during prediction steps to skip computation cheaply. - cache_identifiers (`List[str]`, *optional*, defaults to `None`): + cache_identifiers (`list[str]`, *optional*, defaults to `None`): Regex patterns (using `re.fullmatch`) for module names to place in Taylor-series caching mode, where outputs are approximated and cached for reuse. @@ -82,11 +81,11 @@ def forward(x): cache_interval: int = 5 disable_cache_before_step: int = 3 - disable_cache_after_step: Optional[int] = None + disable_cache_after_step: int | None = None max_order: int = 1 - taylor_factors_dtype: Optional[torch.dtype] = torch.bfloat16 - skip_predict_identifiers: Optional[List[str]] = None - cache_identifiers: Optional[List[str]] = None + taylor_factors_dtype: torch.dtype | None = torch.bfloat16 + skip_predict_identifiers: list[str] | None = None + cache_identifiers: list[str] | None = None use_lite_mode: bool = False def __repr__(self) -> str: @@ -106,7 +105,7 @@ def __repr__(self) -> str: class TaylorSeerState: def __init__( self, - taylor_factors_dtype: Optional[torch.dtype] = torch.bfloat16, + taylor_factors_dtype: torch.dtype | None = torch.bfloat16, max_order: int = 1, is_inactive: bool = False, ): @@ -114,11 +113,11 @@ def __init__( self.max_order = max_order self.is_inactive = is_inactive - self.module_dtypes: Tuple[torch.dtype, ...] = () - self.last_update_step: Optional[int] = None - self.taylor_factors: Dict[int, Dict[int, torch.Tensor]] = {} - self.inactive_shapes: Optional[Tuple[Tuple[int, ...], ...]] = None - self.device: Optional[torch.device] = None + self.module_dtypes: tuple[torch.dtype, ...] = () + self.last_update_step: int | None = None + self.taylor_factors: dict[int, dict[int, torch.Tensor]] = {} + self.inactive_shapes: tuple[tuple[int, ...], ...] | None = None + self.device: torch.device | None = None self.current_step: int = -1 def reset(self) -> None: @@ -130,7 +129,7 @@ def reset(self) -> None: def update( self, - outputs: Tuple[torch.Tensor, ...], + outputs: tuple[torch.Tensor, ...], ) -> None: self.module_dtypes = tuple(output.dtype for output in outputs) self.device = outputs[0].device @@ -139,7 +138,7 @@ def update( self.inactive_shapes = tuple(output.shape for output in outputs) else: for i, features in enumerate(outputs): - new_factors: Dict[int, torch.Tensor] = {0: features} + new_factors: dict[int, torch.Tensor] = {0: features} is_first_update = self.last_update_step is None if not is_first_update: delta_step = self.current_step - self.last_update_step @@ -160,7 +159,7 @@ def update( self.last_update_step = self.current_step @torch.compiler.disable - def predict(self) -> List[torch.Tensor]: + def predict(self) -> list[torch.Tensor]: if self.last_update_step is None: raise ValueError("Cannot predict without prior initialization/update.") @@ -204,7 +203,7 @@ def __init__( disable_cache_before_step: int, taylor_factors_dtype: torch.dtype, state_manager: StateManager, - disable_cache_after_step: Optional[int] = None, + disable_cache_after_step: int | None = None, ): super().__init__() self.cache_interval = cache_interval @@ -245,7 +244,7 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): return outputs_list[0] if len(outputs_list) == 1 else tuple(outputs_list) -def _resolve_patterns(config: TaylorSeerCacheConfig) -> Tuple[List[str], List[str]]: +def _resolve_patterns(config: TaylorSeerCacheConfig) -> tuple[list[str], list[str]]: """ Resolve effective inactive and active pattern lists from config + templates. """ diff --git a/src/diffusers/hooks/utils.py b/src/diffusers/hooks/utils.py index c5260eeebe1f..da9e398fdf29 100644 --- a/src/diffusers/hooks/utils.py +++ b/src/diffusers/hooks/utils.py @@ -21,8 +21,8 @@ def _get_identifiable_transformer_blocks_in_module(module: torch.nn.Module): module_list_with_transformer_blocks = [] for name, submodule in module.named_modules(): name_endswith_identifier = any(name.endswith(identifier) for identifier in _ALL_TRANSFORMER_BLOCK_IDENTIFIERS) - is_modulelist = isinstance(submodule, torch.nn.ModuleList) - if name_endswith_identifier and is_modulelist: + is_ModuleList = isinstance(submodule, torch.nn.ModuleList) + if name_endswith_identifier and is_ModuleList: module_list_with_transformer_blocks.append((name, submodule)) return module_list_with_transformer_blocks diff --git a/src/diffusers/image_processor.py b/src/diffusers/image_processor.py index abd0a25819f5..57cde8d12338 100644 --- a/src/diffusers/image_processor.py +++ b/src/diffusers/image_processor.py @@ -14,7 +14,6 @@ import math import warnings -from typing import List, Optional, Tuple, Union import numpy as np import PIL.Image @@ -26,14 +25,9 @@ from .utils import CONFIG_NAME, PIL_INTERPOLATION, deprecate -PipelineImageInput = Union[ - PIL.Image.Image, - np.ndarray, - torch.Tensor, - List[PIL.Image.Image], - List[np.ndarray], - List[torch.Tensor], -] +PipelineImageInput = ( + PIL.Image.Image | np.ndarray | torch.Tensor | list[PIL.Image.Image] | list[np.ndarray] | list[torch.Tensor] +) PipelineDepthInput = PipelineImageInput @@ -47,7 +41,7 @@ def is_valid_image(image) -> bool: - A 2D or 3D `np.ndarray` or `torch.Tensor` (grayscale or color image). Args: - image (`Union[PIL.Image.Image, np.ndarray, torch.Tensor]`): + image (`PIL.Image.Image | np.ndarray | torch.Tensor`): The image to validate. It can be a PIL image, a NumPy array, or a torch tensor. Returns: @@ -68,7 +62,7 @@ def is_valid_image_imagelist(images): - A list of valid images. Args: - images (`Union[np.ndarray, torch.Tensor, PIL.Image.Image, List]`): + images (`np.ndarray | torch.Tensor | PIL.Image.Image | list`): The image(s) to check. Can be a batch of images (4D tensor/array), a single image, or a list of valid images. @@ -116,7 +110,7 @@ def __init__( vae_scale_factor: int = 8, vae_latent_channels: int = 4, resample: str = "lanczos", - reducing_gap: int = None, + reducing_gap: int | None = None, do_normalize: bool = True, do_binarize: bool = False, do_convert_rgb: bool = False, @@ -131,7 +125,7 @@ def __init__( ) @staticmethod - def numpy_to_pil(images: np.ndarray) -> List[PIL.Image.Image]: + def numpy_to_pil(images: np.ndarray) -> list[PIL.Image.Image]: r""" Convert a numpy image or a batch of images to a PIL image. @@ -140,7 +134,7 @@ def numpy_to_pil(images: np.ndarray) -> List[PIL.Image.Image]: The image array to convert to PIL format. Returns: - `List[PIL.Image.Image]`: + `list[PIL.Image.Image]`: A list of PIL images. """ if images.ndim == 3: @@ -155,12 +149,12 @@ def numpy_to_pil(images: np.ndarray) -> List[PIL.Image.Image]: return pil_images @staticmethod - def pil_to_numpy(images: Union[List[PIL.Image.Image], PIL.Image.Image]) -> np.ndarray: + def pil_to_numpy(images: list[PIL.Image.Image] | PIL.Image.Image) -> np.ndarray: r""" Convert a PIL image or a list of PIL images to NumPy arrays. Args: - images (`PIL.Image.Image` or `List[PIL.Image.Image]`): + images (`PIL.Image.Image` or `list[PIL.Image.Image]`): The PIL image or list of images to convert to NumPy format. Returns: @@ -210,7 +204,7 @@ def pt_to_numpy(images: torch.Tensor) -> np.ndarray: return images @staticmethod - def normalize(images: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + def normalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor: r""" Normalize an image array to [-1,1]. @@ -225,7 +219,7 @@ def normalize(images: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torc return 2.0 * images - 1.0 @staticmethod - def denormalize(images: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + def denormalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor: r""" Denormalize an image array to [0,1]. @@ -467,11 +461,11 @@ def _resize_and_crop( def resize( self, - image: Union[PIL.Image.Image, np.ndarray, torch.Tensor], + image: PIL.Image.Image | np.ndarray | torch.Tensor, height: int, width: int, resize_mode: str = "default", # "default", "fill", "crop" - ) -> Union[PIL.Image.Image, np.ndarray, torch.Tensor]: + ) -> PIL.Image.Image | np.ndarray | torch.Tensor: """ Resize image. @@ -544,7 +538,7 @@ def binarize(self, image: PIL.Image.Image) -> PIL.Image.Image: return image def _denormalize_conditionally( - self, images: torch.Tensor, do_denormalize: Optional[List[bool]] = None + self, images: torch.Tensor, do_denormalize: list[bool] | None = None ) -> torch.Tensor: r""" Denormalize a batch of images based on a condition list. @@ -552,7 +546,7 @@ def _denormalize_conditionally( Args: images (`torch.Tensor`): The input image tensor. - do_denormalize (`Optional[List[bool]`, *optional*, defaults to `None`): + do_denormalize (`Optional[list[bool]`, *optional*, defaults to `None`): A list of booleans indicating whether to denormalize each image in the batch. If `None`, will use the value of `do_normalize` in the `VaeImageProcessor` config. """ @@ -565,25 +559,25 @@ def _denormalize_conditionally( def get_default_height_width( self, - image: Union[PIL.Image.Image, np.ndarray, torch.Tensor], - height: Optional[int] = None, - width: Optional[int] = None, - ) -> Tuple[int, int]: + image: PIL.Image.Image | np.ndarray | torch.Tensor, + height: int | None = None, + width: int | None = None, + ) -> tuple[int, int]: r""" Returns the height and width of the image, downscaled to the next integer multiple of `vae_scale_factor`. Args: - image (`Union[PIL.Image.Image, np.ndarray, torch.Tensor]`): + image (`PIL.Image.Image | np.ndarray | torch.Tensor`): The image input, which can be a PIL image, NumPy array, or PyTorch tensor. If it is a NumPy array, it should have shape `[batch, height, width]` or `[batch, height, width, channels]`. If it is a PyTorch tensor, it should have shape `[batch, channels, height, width]`. - height (`Optional[int]`, *optional*, defaults to `None`): + height (`int | None`, *optional*, defaults to `None`): The height of the preprocessed image. If `None`, the height of the `image` input will be used. - width (`Optional[int]`, *optional*, defaults to `None`): + width (`int | None`, *optional*, defaults to `None`): The width of the preprocessed image. If `None`, the width of the `image` input will be used. Returns: - `Tuple[int, int]`: + `tuple[int, int]`: A tuple containing the height and width, both resized to the nearest integer multiple of `vae_scale_factor`. """ @@ -613,10 +607,10 @@ def get_default_height_width( def preprocess( self, image: PipelineImageInput, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, resize_mode: str = "default", # "default", "fill", "crop" - crops_coords: Optional[Tuple[int, int, int, int]] = None, + crops_coords: tuple[int, int, int, int] | None = None, ) -> torch.Tensor: """ Preprocess the image input. @@ -638,7 +632,7 @@ def preprocess( image to fit within the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, cropping the excess. Note that resize_mode `fill` and `crop` are only supported for PIL image input. - crops_coords (`List[Tuple[int, int, int, int]]`, *optional*, defaults to `None`): + crops_coords (`list[tuple[int, int, int, int]]`, *optional*, defaults to `None`): The crop coordinates for each image in the batch. If `None`, will not crop the image. Returns: @@ -745,8 +739,8 @@ def postprocess( self, image: torch.Tensor, output_type: str = "pil", - do_denormalize: Optional[List[bool]] = None, - ) -> Union[PIL.Image.Image, np.ndarray, torch.Tensor]: + do_denormalize: list[bool] | None = None, + ) -> PIL.Image.Image | np.ndarray | torch.Tensor: """ Postprocess the image output from tensor to `output_type`. @@ -755,7 +749,7 @@ def postprocess( The image input, should be a pytorch tensor with shape `B x C x H x W`. output_type (`str`, *optional*, defaults to `pil`): The output type of the image, can be one of `pil`, `np`, `pt`, `latent`. - do_denormalize (`List[bool]`, *optional*, defaults to `None`): + do_denormalize (`list[bool]`, *optional*, defaults to `None`): Whether to denormalize the image to [0,1]. If `None`, will use the value of `do_normalize` in the `VaeImageProcessor` config. @@ -796,7 +790,7 @@ def apply_overlay( mask: PIL.Image.Image, init_image: PIL.Image.Image, image: PIL.Image.Image, - crop_coords: Optional[Tuple[int, int, int, int]] = None, + crop_coords: tuple[int, int, int, int] | None = None, ) -> PIL.Image.Image: r""" Applies an overlay of the mask and the inpainted image on the original image. @@ -808,7 +802,7 @@ def apply_overlay( The original image to which the overlay is applied. image (`PIL.Image.Image`): The image to overlay onto the original. - crop_coords (`Tuple[int, int, int, int]`, *optional*): + crop_coords (`tuple[int, int, int, int]`, *optional*): Coordinates to crop the image. If provided, the image will be cropped accordingly. Returns: @@ -853,7 +847,7 @@ def __init__( vae_scale_factor: int = 8, vae_latent_channels: int = 4, resample: str = "lanczos", - reducing_gap: int = None, + reducing_gap: int | None = None, do_normalize: bool = True, do_binarize: bool = False, do_convert_grayscale: bool = False, @@ -887,11 +881,11 @@ def __init__( def preprocess( self, image: PIL.Image.Image, - mask: PIL.Image.Image = None, - height: int = None, - width: int = None, - padding_mask_crop: Optional[int] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + mask: PIL.Image.Image | None = None, + height: int | None = None, + width: int | None = None, + padding_mask_crop: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: """ Preprocess the image and mask. """ @@ -944,10 +938,10 @@ def postprocess( self, image: torch.Tensor, output_type: str = "pil", - original_image: Optional[PIL.Image.Image] = None, - original_mask: Optional[PIL.Image.Image] = None, - crops_coords: Optional[Tuple[int, int, int, int]] = None, - ) -> Tuple[PIL.Image.Image, PIL.Image.Image]: + original_image: PIL.Image.Image | None = None, + original_mask: PIL.Image.Image | None = None, + crops_coords: tuple[int, int, int, int] | None = None, + ) -> tuple[PIL.Image.Image, PIL.Image.Image]: """ Postprocess the image, optionally apply mask overlay """ @@ -998,7 +992,7 @@ def __init__( super().__init__() @staticmethod - def numpy_to_pil(images: np.ndarray) -> List[PIL.Image.Image]: + def numpy_to_pil(images: np.ndarray) -> list[PIL.Image.Image]: r""" Convert a NumPy image or a batch of images to a list of PIL images. @@ -1007,7 +1001,7 @@ def numpy_to_pil(images: np.ndarray) -> List[PIL.Image.Image]: The input NumPy array of images, which can be a single image or a batch. Returns: - `List[PIL.Image.Image]`: + `list[PIL.Image.Image]`: A list of PIL images converted from the input NumPy array. """ if images.ndim == 3: @@ -1022,12 +1016,12 @@ def numpy_to_pil(images: np.ndarray) -> List[PIL.Image.Image]: return pil_images @staticmethod - def depth_pil_to_numpy(images: Union[List[PIL.Image.Image], PIL.Image.Image]) -> np.ndarray: + def depth_pil_to_numpy(images: list[PIL.Image.Image] | PIL.Image.Image) -> np.ndarray: r""" Convert a PIL image or a list of PIL images to NumPy arrays. Args: - images (`Union[List[PIL.Image.Image], PIL.Image.Image]`): + images (`list[PIL.Image.Image, PIL.Image.Image]`): The input image or list of images to be converted. Returns: @@ -1042,7 +1036,7 @@ def depth_pil_to_numpy(images: Union[List[PIL.Image.Image], PIL.Image.Image]) -> return images @staticmethod - def rgblike_to_depthmap(image: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + def rgblike_to_depthmap(image: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor: r""" Convert an RGB-like depth image to a depth map. """ @@ -1079,7 +1073,7 @@ def rgblike_to_depthmap(image: Union[np.ndarray, torch.Tensor]) -> Union[np.ndar else: raise TypeError("Input image must be a torch.Tensor or np.ndarray") - def numpy_to_depth(self, images: np.ndarray) -> List[PIL.Image.Image]: + def numpy_to_depth(self, images: np.ndarray) -> list[PIL.Image.Image]: r""" Convert a NumPy depth image or a batch of images to a list of PIL images. @@ -1088,7 +1082,7 @@ def numpy_to_depth(self, images: np.ndarray) -> List[PIL.Image.Image]: The input NumPy array of depth images, which can be a single image or a batch. Returns: - `List[PIL.Image.Image]`: + `list[PIL.Image.Image]`: A list of PIL images converted from the input NumPy depth images. """ if images.ndim == 3: @@ -1111,8 +1105,8 @@ def postprocess( self, image: torch.Tensor, output_type: str = "pil", - do_denormalize: Optional[List[bool]] = None, - ) -> Union[PIL.Image.Image, np.ndarray, torch.Tensor]: + do_denormalize: list[bool] | None = None, + ) -> PIL.Image.Image | np.ndarray | torch.Tensor: """ Postprocess the image output from tensor to `output_type`. @@ -1121,7 +1115,7 @@ def postprocess( The image input, should be a pytorch tensor with shape `B x C x H x W`. output_type (`str`, *optional*, defaults to `pil`): The output type of the image, can be one of `pil`, `np`, `pt`, `latent`. - do_denormalize (`List[bool]`, *optional*, defaults to `None`): + do_denormalize (`list[bool]`, *optional*, defaults to `None`): Whether to denormalize the image to [0,1]. If `None`, will use the value of `do_normalize` in the `VaeImageProcessor` config. @@ -1159,29 +1153,29 @@ def postprocess( def preprocess( self, - rgb: Union[torch.Tensor, PIL.Image.Image, np.ndarray], - depth: Union[torch.Tensor, PIL.Image.Image, np.ndarray], - height: Optional[int] = None, - width: Optional[int] = None, - target_res: Optional[int] = None, + rgb: torch.Tensor | PIL.Image.Image | np.ndarray, + depth: torch.Tensor | PIL.Image.Image | np.ndarray, + height: int | None = None, + width: int | None = None, + target_res: int | None = None, ) -> torch.Tensor: r""" Preprocess the image input. Accepted formats are PIL images, NumPy arrays, or PyTorch tensors. Args: - rgb (`Union[torch.Tensor, PIL.Image.Image, np.ndarray]`): + rgb (`torch.Tensor | PIL.Image.Image | np.ndarray`): The RGB input image, which can be a single image or a batch. - depth (`Union[torch.Tensor, PIL.Image.Image, np.ndarray]`): + depth (`torch.Tensor | PIL.Image.Image | np.ndarray`): The depth input image, which can be a single image or a batch. - height (`Optional[int]`, *optional*, defaults to `None`): + height (`int | None`, *optional*, defaults to `None`): The desired height of the processed image. If `None`, defaults to the height of the input image. - width (`Optional[int]`, *optional*, defaults to `None`): + width (`int | None`, *optional*, defaults to `None`): The desired width of the processed image. If `None`, defaults to the width of the input image. - target_res (`Optional[int]`, *optional*, defaults to `None`): + target_res (`int | None`, *optional*, defaults to `None`): Target resolution for resizing the images. If specified, overrides height and width. Returns: - `Tuple[torch.Tensor, torch.Tensor]`: + `tuple[torch.Tensor, torch.Tensor]`: A tuple containing the processed RGB and depth images as PyTorch tensors. """ supported_formats = (PIL.Image.Image, np.ndarray, torch.Tensor) @@ -1419,7 +1413,7 @@ def __init__( ) @staticmethod - def classify_height_width_bin(height: int, width: int, ratios: dict) -> Tuple[int, int]: + def classify_height_width_bin(height: int, width: int, ratios: dict) -> tuple[int, int]: r""" Returns the binned height and width based on the aspect ratio. @@ -1429,7 +1423,7 @@ def classify_height_width_bin(height: int, width: int, ratios: dict) -> Tuple[in ratios (`dict`): A dictionary where keys are aspect ratios and values are tuples of (height, width). Returns: - `Tuple[int, int]`: The closest binned height and width. + `tuple[int, int]`: The closest binned height and width. """ ar = float(height / width) closest_ratio = min(ratios.keys(), key=lambda ratio: abs(float(ratio) - ar)) diff --git a/src/diffusers/loaders/__init__.py b/src/diffusers/loaders/__init__.py index ace4e8543a1c..ed0d2a07336f 100644 --- a/src/diffusers/loaders/__init__.py +++ b/src/diffusers/loaders/__init__.py @@ -67,6 +67,7 @@ def text_encoder_attn_modules(text_encoder): "SD3LoraLoaderMixin", "AuraFlowLoraLoaderMixin", "StableDiffusionXLLoraLoaderMixin", + "LTX2LoraLoaderMixin", "LTXVideoLoraLoaderMixin", "LoraLoaderMixin", "FluxLoraLoaderMixin", @@ -77,6 +78,7 @@ def text_encoder_attn_modules(text_encoder): "SanaLoraLoaderMixin", "Lumina2LoraLoaderMixin", "WanLoraLoaderMixin", + "HeliosLoraLoaderMixin", "KandinskyLoraLoaderMixin", "HiDreamImageLoraLoaderMixin", "SkyReelsV2LoraLoaderMixin", @@ -117,10 +119,12 @@ def text_encoder_attn_modules(text_encoder): CogView4LoraLoaderMixin, Flux2LoraLoaderMixin, FluxLoraLoaderMixin, + HeliosLoraLoaderMixin, HiDreamImageLoraLoaderMixin, HunyuanVideoLoraLoaderMixin, KandinskyLoraLoaderMixin, LoraLoaderMixin, + LTX2LoraLoaderMixin, LTXVideoLoraLoaderMixin, Lumina2LoraLoaderMixin, Mochi1LoraLoaderMixin, diff --git a/src/diffusers/loaders/ip_adapter.py b/src/diffusers/loaders/ip_adapter.py index dca4758ba038..13bb44e4a2a6 100644 --- a/src/diffusers/loaders/ip_adapter.py +++ b/src/diffusers/loaders/ip_adapter.py @@ -13,7 +13,7 @@ # limitations under the License. from pathlib import Path -from typing import Dict, List, Optional, Union +from typing import List, Union import torch import torch.nn.functional as F @@ -57,15 +57,15 @@ class IPAdapterMixin: @validate_hf_hub_args def load_ip_adapter( self, - pretrained_model_name_or_path_or_dict: Union[str, List[str], Dict[str, torch.Tensor]], - subfolder: Union[str, List[str]], - weight_name: Union[str, List[str]], - image_encoder_folder: Optional[str] = "image_encoder", + pretrained_model_name_or_path_or_dict: str | list[str] | dict[str, torch.Tensor], + subfolder: str | list[str], + weight_name: str | list[str], + image_encoder_folder: str | None = "image_encoder", **kwargs, ): """ Parameters: - pretrained_model_name_or_path_or_dict (`str` or `List[str]` or `os.PathLike` or `List[os.PathLike]` or `dict` or `List[dict]`): + pretrained_model_name_or_path_or_dict (`str` or `list[str]` or `os.PathLike` or `list[os.PathLike]` or `dict` or `list[dict]`): Can be either: - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on @@ -74,10 +74,10 @@ def load_ip_adapter( with [`ModelMixin.save_pretrained`]. - A [torch state dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - subfolder (`str` or `List[str]`): + subfolder (`str` or `list[str]`): The subfolder location of a model file within a larger model repository on the Hub or locally. If a list is passed, it should have the same length as `weight_name`. - weight_name (`str` or `List[str]`): + weight_name (`str` or `list[str]`): The name of the weight file to load. If a list is passed, it should have the same length as `subfolder`. image_encoder_folder (`str`, *optional*, defaults to `image_encoder`): @@ -87,14 +87,14 @@ def load_ip_adapter( `image_encoder_folder="image_encoder"`. If the image encoder is located in a folder other than `subfolder`, you should pass the path to the folder that contains image encoder weights, for example, `image_encoder_folder="different_subfolder/image_encoder"`. - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. local_files_only (`bool`, *optional*, defaults to `False`): @@ -358,14 +358,14 @@ class ModularIPAdapterMixin: @validate_hf_hub_args def load_ip_adapter( self, - pretrained_model_name_or_path_or_dict: Union[str, List[str], Dict[str, torch.Tensor]], - subfolder: Union[str, List[str]], - weight_name: Union[str, List[str]], + pretrained_model_name_or_path_or_dict: str | list[str] | dict[str, torch.Tensor], + subfolder: str | list[str], + weight_name: str | list[str], **kwargs, ): """ Parameters: - pretrained_model_name_or_path_or_dict (`str` or `List[str]` or `os.PathLike` or `List[os.PathLike]` or `dict` or `List[dict]`): + pretrained_model_name_or_path_or_dict (`str` or `list[str]` or `os.PathLike` or `list[os.PathLike]` or `dict` or `list[dict]`): Can be either: - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on @@ -374,20 +374,20 @@ def load_ip_adapter( with [`ModelMixin.save_pretrained`]. - A [torch state dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - subfolder (`str` or `List[str]`): + subfolder (`str` or `list[str]`): The subfolder location of a model file within a larger model repository on the Hub or locally. If a list is passed, it should have the same length as `weight_name`. - weight_name (`str` or `List[str]`): + weight_name (`str` or `list[str]`): The name of the weight file to load. If a list is passed, it should have the same length as `subfolder`. - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. local_files_only (`bool`, *optional*, defaults to `False`): @@ -608,17 +608,17 @@ class FluxIPAdapterMixin: @validate_hf_hub_args def load_ip_adapter( self, - pretrained_model_name_or_path_or_dict: Union[str, List[str], Dict[str, torch.Tensor]], - weight_name: Union[str, List[str]], - subfolder: Optional[Union[str, List[str]]] = "", - image_encoder_pretrained_model_name_or_path: Optional[str] = "image_encoder", - image_encoder_subfolder: Optional[str] = "", + pretrained_model_name_or_path_or_dict: str | list[str] | dict[str, torch.Tensor], + weight_name: str | list[str], + subfolder: str | list[str] | None = "", + image_encoder_pretrained_model_name_or_path: str | None = "image_encoder", + image_encoder_subfolder: str | None = "", image_encoder_dtype: torch.dtype = torch.float16, **kwargs, ): """ Parameters: - pretrained_model_name_or_path_or_dict (`str` or `List[str]` or `os.PathLike` or `List[os.PathLike]` or `dict` or `List[dict]`): + pretrained_model_name_or_path_or_dict (`str` or `list[str]` or `os.PathLike` or `list[os.PathLike]` or `dict` or `list[dict]`): Can be either: - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on @@ -627,10 +627,10 @@ def load_ip_adapter( with [`ModelMixin.save_pretrained`]. - A [torch state dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - subfolder (`str` or `List[str]`): + subfolder (`str` or `list[str]`): The subfolder location of a model file within a larger model repository on the Hub or locally. If a list is passed, it should have the same length as `weight_name`. - weight_name (`str` or `List[str]`): + weight_name (`str` or `list[str]`): The name of the weight file to load. If a list is passed, it should have the same length as `weight_name`. image_encoder_pretrained_model_name_or_path (`str`, *optional*, defaults to `./image_encoder`): @@ -640,14 +640,14 @@ def load_ip_adapter( hosted on the Hub. - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved with [`ModelMixin.save_pretrained`]. - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. local_files_only (`bool`, *optional*, defaults to `False`): @@ -797,13 +797,13 @@ def load_ip_adapter( # load ip-adapter into transformer self.transformer._load_ip_adapter_weights(state_dicts, low_cpu_mem_usage=low_cpu_mem_usage) - def set_ip_adapter_scale(self, scale: Union[float, List[float], List[List[float]]]): + def set_ip_adapter_scale(self, scale: float | list[float] | list[list[float]]): """ Set IP-Adapter scales per-transformer block. Input `scale` could be a single config or a list of configs for granular control over each IP-Adapter behavior. A config can be a float or a list. - `float` is converted to list and repeated for the number of blocks and the number of IP adapters. `List[float]` - length match the number of blocks, it is repeated for each IP adapter. `List[List[float]]` must match the + `float` is converted to list and repeated for the number of blocks and the number of IP adapters. `list[float]` + length match the number of blocks, it is repeated for each IP adapter. `list[list[float]]` must match the number of IP adapters and each must match the number of blocks. Example: @@ -918,10 +918,10 @@ def is_ip_adapter_active(self) -> bool: @validate_hf_hub_args def load_ip_adapter( self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], weight_name: str = "ip-adapter.safetensors", - subfolder: Optional[str] = None, - image_encoder_folder: Optional[str] = "image_encoder", + subfolder: str | None = None, + image_encoder_folder: str | None = "image_encoder", **kwargs, ) -> None: """ @@ -947,13 +947,13 @@ def load_ip_adapter( `image_encoder_folder="image_encoder"`. If the image encoder is located in a folder other than `subfolder`, you should pass the path to the folder that contains image encoder weights, for example, `image_encoder_folder="different_subfolder/image_encoder"`. - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. local_files_only (`bool`, *optional*, defaults to `False`): diff --git a/src/diffusers/loaders/lora_base.py b/src/diffusers/loaders/lora_base.py index 3d75a7d875a4..5b5579664b55 100644 --- a/src/diffusers/loaders/lora_base.py +++ b/src/diffusers/loaders/lora_base.py @@ -11,13 +11,14 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations import copy import inspect import json import os from pathlib import Path -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import safetensors import torch @@ -77,7 +78,7 @@ def fuse_text_encoder_lora(text_encoder, lora_scale=1.0, safe_fusing=False, adap Controls how much to influence the outputs with the LoRA parameters. safe_fusing (`bool`, defaults to `False`): Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them. - adapter_names (`List[str]` or `str`): + adapter_names (`list[str]` or `str`): The names of the adapters to use. """ merge_kwargs = {"safe_merge": safe_fusing} @@ -116,20 +117,20 @@ def unfuse_text_encoder_lora(text_encoder): def set_adapters_for_text_encoder( - adapter_names: Union[List[str], str], - text_encoder: Optional["PreTrainedModel"] = None, # noqa: F821 - text_encoder_weights: Optional[Union[float, List[float], List[None]]] = None, + adapter_names: list[str] | str, + text_encoder: "PreTrainedModel" | None = None, # noqa: F821 + text_encoder_weights: float | list[float] | list[None] | None = None, ): """ Sets the adapter layers for the text encoder. Args: - adapter_names (`List[str]` or `str`): + adapter_names (`list[str]` or `str`): The names of the adapters to use. text_encoder (`torch.nn.Module`, *optional*): The text encoder module to set the adapter layers for. If `None`, it will try to get the `text_encoder` attribute. - text_encoder_weights (`List[float]`, *optional*): + text_encoder_weights (`list[float]`, *optional*): The weights to use for the text encoder. If `None`, the weights are set to `1.0` for all the adapters. """ if text_encoder is None: @@ -159,7 +160,7 @@ def process_weights(adapter_names, weights): set_weights_and_activate_adapters(text_encoder, adapter_names, text_encoder_weights) -def disable_lora_for_text_encoder(text_encoder: Optional["PreTrainedModel"] = None): +def disable_lora_for_text_encoder(text_encoder: "PreTrainedModel" | None = None): """ Disables the LoRA layers for the text encoder. @@ -173,7 +174,7 @@ def disable_lora_for_text_encoder(text_encoder: Optional["PreTrainedModel"] = No set_adapter_layers(text_encoder, enabled=False) -def enable_lora_for_text_encoder(text_encoder: Optional["PreTrainedModel"] = None): +def enable_lora_for_text_encoder(text_encoder: "PreTrainedModel" | None = None): """ Enables the LoRA layers for the text encoder. @@ -535,10 +536,10 @@ def unload_lora_weights(self): def fuse_lora( self, - components: List[str] = [], + components: list[str] = [], lora_scale: float = 1.0, safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, + adapter_names: list[str] | None = None, **kwargs, ): r""" @@ -547,12 +548,12 @@ def fuse_lora( > [!WARNING] > This is an experimental API. Args: - components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into. + components: (`list[str]`): list of LoRA-injectable components to fuse the LoRAs into. lora_scale (`float`, defaults to 1.0): Controls how much to influence the outputs with the LoRA parameters. safe_fusing (`bool`, defaults to `False`): Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them. - adapter_names (`List[str]`, *optional*): + adapter_names (`list[str]`, *optional*): Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused. Example: @@ -619,7 +620,7 @@ def fuse_lora( self._merged_adapters = self._merged_adapters | merged_adapter_names - def unfuse_lora(self, components: List[str] = [], **kwargs): + def unfuse_lora(self, components: list[str] = [], **kwargs): r""" Reverses the effect of [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). @@ -627,7 +628,7 @@ def unfuse_lora(self, components: List[str] = [], **kwargs): > [!WARNING] > This is an experimental API. Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. + components (`list[str]`): list of LoRA-injectable components to unfuse LoRA from. unfuse_unet (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. unfuse_text_encoder (`bool`, defaults to `True`): Whether to unfuse the text encoder LoRA parameters. If the text encoder wasn't monkey-patched with the @@ -674,16 +675,16 @@ def unfuse_lora(self, components: List[str] = [], **kwargs): def set_adapters( self, - adapter_names: Union[List[str], str], - adapter_weights: Optional[Union[float, Dict, List[float], List[Dict]]] = None, + adapter_names: list[str] | str, + adapter_weights: float | dict | list[float] | list[dict] | None = None, ): """ Set the currently active adapters for use in the pipeline. Args: - adapter_names (`List[str]` or `str`): + adapter_names (`list[str]` or `str`): The names of the adapters to use. - adapter_weights (`Union[List[float], float]`, *optional*): + adapter_weights (`list[float, float]`, *optional*): The adapter(s) weights to use with the UNet. If `None`, the weights are set to `1.0` for all the adapters. @@ -835,12 +836,12 @@ def enable_lora(self): elif issubclass(model.__class__, PreTrainedModel): enable_lora_for_text_encoder(model) - def delete_adapters(self, adapter_names: Union[List[str], str]): + def delete_adapters(self, adapter_names: list[str] | str): """ Delete an adapter's LoRA layers from the pipeline. Args: - adapter_names (`Union[List[str], str]`): + adapter_names (`list[str, str]`): The names of the adapters to delete. Example: @@ -873,7 +874,7 @@ def delete_adapters(self, adapter_names: Union[List[str], str]): for adapter_name in adapter_names: delete_adapter_layers(model, adapter_name) - def get_active_adapters(self) -> List[str]: + def get_active_adapters(self) -> list[str]: """ Gets the list of the current active adapters. @@ -906,7 +907,7 @@ def get_active_adapters(self) -> List[str]: return active_adapters - def get_list_adapters(self) -> Dict[str, List[str]]: + def get_list_adapters(self) -> dict[str, list[str]]: """ Gets the current list of all available adapters in the pipeline. """ @@ -928,7 +929,7 @@ def get_list_adapters(self) -> Dict[str, List[str]]: return set_adapters - def set_lora_device(self, adapter_names: List[str], device: Union[torch.device, str, int]) -> None: + def set_lora_device(self, adapter_names: list[str], device: torch.device | str | int) -> None: """ Moves the LoRAs listed in `adapter_names` to a target device. Useful for offloading the LoRA to the CPU in case you want to load multiple adapters and free some GPU memory. @@ -955,9 +956,9 @@ def set_lora_device(self, adapter_names: List[str], device: Union[torch.device, ``` Args: - adapter_names (`List[str]`): - List of adapters to send device to. - device (`Union[torch.device, str, int]`): + adapter_names (`list[str]`): + list of adapters to send device to. + device (`torch.device | str | int`): Device to send the adapters to. Can be either a torch device, a str or an integer. """ if not USE_PEFT_BACKEND: @@ -1007,13 +1008,13 @@ def pack_weights(layers, prefix): @staticmethod def write_lora_layers( - state_dict: Dict[str, torch.Tensor], + state_dict: dict[str, torch.Tensor], save_directory: str, is_main_process: bool, weight_name: str, save_function: Callable, safe_serialization: bool, - lora_adapter_metadata: Optional[dict] = None, + lora_adapter_metadata: dict | None = None, ): """Writes the state dict of the LoRA layers (optionally with metadata) to disk.""" if os.path.isfile(save_directory): @@ -1059,9 +1060,9 @@ def save_function(weights, filename): @classmethod def _save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - lora_layers: Dict[str, Dict[str, Union[torch.nn.Module, torch.Tensor]]], - lora_metadata: Dict[str, Optional[dict]], + save_directory: str | os.PathLike, + lora_layers: dict[str, dict[str, torch.nn.Module | torch.Tensor]], + lora_metadata: dict[str, dict | None], is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index 2e87f757c352..0895d5223e13 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -13,7 +13,6 @@ # limitations under the License. import re -from typing import List import torch @@ -857,7 +856,7 @@ def _convert(original_key, diffusers_key, state_dict, new_state_dict): ) state_dict = {k: v for k, v in state_dict.items() if not k.startswith("text_encoders.t5xxl.transformer.")} - has_diffb = any("diff_b" in k and k.startswith(("lora_unet_", "lora_te_")) for k in state_dict) + has_diffb = any("diff_b" in k and k.startswith(("lora_unet_", "lora_te_", "lora_te1_")) for k in state_dict) if has_diffb: zero_status_diff_b = state_dict_all_zero(state_dict, ".diff_b") if zero_status_diff_b: @@ -896,7 +895,7 @@ def _convert(original_key, diffusers_key, state_dict, new_state_dict): state_dict = { _custom_replace(k, limit_substrings): v for k, v in state_dict.items() - if k.startswith(("lora_unet_", "lora_te_")) + if k.startswith(("lora_unet_", "lora_te_", "lora_te1_")) } if any("text_projection" in k for k in state_dict): @@ -1021,7 +1020,7 @@ def handle_qkv(sds_sd, ait_sd, sds_key, ait_keys, dims=None): return new_state_dict -def _custom_replace(key: str, substrings: List[str]) -> str: +def _custom_replace(key: str, substrings: list[str]) -> str: # Replaces the "."s with "_"s upto the `substrings`. # Example: # lora_unet.foo.bar.lora_A.weight -> lora_unet_foo_bar.lora_A.weight @@ -2140,6 +2139,54 @@ def _convert_non_diffusers_ltxv_lora_to_diffusers(state_dict, non_diffusers_pref return converted_state_dict +def _convert_non_diffusers_ltx2_lora_to_diffusers(state_dict, non_diffusers_prefix="diffusion_model"): + # Remove the prefix + state_dict = {k: v for k, v in state_dict.items() if k.startswith(f"{non_diffusers_prefix}.")} + converted_state_dict = {k.removeprefix(f"{non_diffusers_prefix}."): v for k, v in state_dict.items()} + + if non_diffusers_prefix == "diffusion_model": + rename_dict = { + "patchify_proj": "proj_in", + "audio_patchify_proj": "audio_proj_in", + "av_ca_video_scale_shift_adaln_single": "av_cross_attn_video_scale_shift", + "av_ca_a2v_gate_adaln_single": "av_cross_attn_video_a2v_gate", + "av_ca_audio_scale_shift_adaln_single": "av_cross_attn_audio_scale_shift", + "av_ca_v2a_gate_adaln_single": "av_cross_attn_audio_v2a_gate", + "scale_shift_table_a2v_ca_video": "video_a2v_cross_attn_scale_shift_table", + "scale_shift_table_a2v_ca_audio": "audio_a2v_cross_attn_scale_shift_table", + "q_norm": "norm_q", + "k_norm": "norm_k", + } + else: + rename_dict = {"aggregate_embed": "text_proj_in"} + + # Apply renaming + renamed_state_dict = {} + for key, value in converted_state_dict.items(): + new_key = key[:] + for old_pattern, new_pattern in rename_dict.items(): + new_key = new_key.replace(old_pattern, new_pattern) + renamed_state_dict[new_key] = value + + # Handle adaln_single -> time_embed and audio_adaln_single -> audio_time_embed + final_state_dict = {} + for key, value in renamed_state_dict.items(): + if key.startswith("adaln_single."): + new_key = key.replace("adaln_single.", "time_embed.") + final_state_dict[new_key] = value + elif key.startswith("audio_adaln_single."): + new_key = key.replace("audio_adaln_single.", "audio_time_embed.") + final_state_dict[new_key] = value + else: + final_state_dict[key] = value + + # Add transformer prefix + prefix = "transformer" if non_diffusers_prefix == "diffusion_model" else "connectors" + final_state_dict = {f"{prefix}.{k}": v for k, v in final_state_dict.items()} + + return final_state_dict + + def _convert_non_diffusers_qwen_lora_to_diffusers(state_dict): has_diffusion_model = any(k.startswith("diffusion_model.") for k in state_dict) if has_diffusion_model: @@ -2273,8 +2320,22 @@ def _convert_non_diffusers_flux2_lora_to_diffusers(state_dict): prefix = "diffusion_model." original_state_dict = {k[len(prefix) :]: v for k, v in state_dict.items()} - num_double_layers = 8 - num_single_layers = 48 + has_lora_down_up = any("lora_down" in k or "lora_up" in k for k in original_state_dict.keys()) + if has_lora_down_up: + temp_state_dict = {} + for k, v in original_state_dict.items(): + new_key = k.replace("lora_down", "lora_A").replace("lora_up", "lora_B") + temp_state_dict[new_key] = v + original_state_dict = temp_state_dict + + num_double_layers = 0 + num_single_layers = 0 + for key in original_state_dict.keys(): + if key.startswith("single_blocks."): + num_single_layers = max(num_single_layers, int(key.split(".")[1]) + 1) + elif key.startswith("double_blocks."): + num_double_layers = max(num_double_layers, int(key.split(".")[1]) + 1) + lora_keys = ("lora_A", "lora_B") attn_types = ("img_attn", "txt_attn") @@ -2283,13 +2344,15 @@ def _convert_non_diffusers_flux2_lora_to_diffusers(state_dict): attn_prefix = f"single_transformer_blocks.{sl}.attn" for lora_key in lora_keys: - converted_state_dict[f"{attn_prefix}.to_qkv_mlp_proj.{lora_key}.weight"] = original_state_dict.pop( - f"{single_block_prefix}.linear1.{lora_key}.weight" - ) + linear1_key = f"{single_block_prefix}.linear1.{lora_key}.weight" + if linear1_key in original_state_dict: + converted_state_dict[f"{attn_prefix}.to_qkv_mlp_proj.{lora_key}.weight"] = original_state_dict.pop( + linear1_key + ) - converted_state_dict[f"{attn_prefix}.to_out.{lora_key}.weight"] = original_state_dict.pop( - f"{single_block_prefix}.linear2.{lora_key}.weight" - ) + linear2_key = f"{single_block_prefix}.linear2.{lora_key}.weight" + if linear2_key in original_state_dict: + converted_state_dict[f"{attn_prefix}.to_out.{lora_key}.weight"] = original_state_dict.pop(linear2_key) for dl in range(num_double_layers): transformer_block_prefix = f"transformer_blocks.{dl}" @@ -2298,6 +2361,10 @@ def _convert_non_diffusers_flux2_lora_to_diffusers(state_dict): for attn_type in attn_types: attn_prefix = f"{transformer_block_prefix}.attn" qkv_key = f"double_blocks.{dl}.{attn_type}.qkv.{lora_key}.weight" + + if qkv_key not in original_state_dict: + continue + fused_qkv_weight = original_state_dict.pop(qkv_key) if lora_key == "lora_A": @@ -2329,8 +2396,9 @@ def _convert_non_diffusers_flux2_lora_to_diffusers(state_dict): for org_proj, diff_proj in proj_mappings: for lora_key in lora_keys: original_key = f"double_blocks.{dl}.{org_proj}.{lora_key}.weight" - diffusers_key = f"{transformer_block_prefix}.{diff_proj}.{lora_key}.weight" - converted_state_dict[diffusers_key] = original_state_dict.pop(original_key) + if original_key in original_state_dict: + diffusers_key = f"{transformer_block_prefix}.{diff_proj}.{lora_key}.weight" + converted_state_dict[diffusers_key] = original_state_dict.pop(original_key) mlp_mappings = [ ("img_mlp.0", "ff.linear_in"), @@ -2341,8 +2409,27 @@ def _convert_non_diffusers_flux2_lora_to_diffusers(state_dict): for org_mlp, diff_mlp in mlp_mappings: for lora_key in lora_keys: original_key = f"double_blocks.{dl}.{org_mlp}.{lora_key}.weight" - diffusers_key = f"{transformer_block_prefix}.{diff_mlp}.{lora_key}.weight" - converted_state_dict[diffusers_key] = original_state_dict.pop(original_key) + if original_key in original_state_dict: + diffusers_key = f"{transformer_block_prefix}.{diff_mlp}.{lora_key}.weight" + converted_state_dict[diffusers_key] = original_state_dict.pop(original_key) + + extra_mappings = { + "img_in": "x_embedder", + "txt_in": "context_embedder", + "time_in.in_layer": "time_guidance_embed.timestep_embedder.linear_1", + "time_in.out_layer": "time_guidance_embed.timestep_embedder.linear_2", + "final_layer.linear": "proj_out", + "final_layer.adaLN_modulation.1": "norm_out.linear", + "single_stream_modulation.lin": "single_stream_modulation.linear", + "double_stream_modulation_img.lin": "double_stream_modulation_img.linear", + "double_stream_modulation_txt.lin": "double_stream_modulation_txt.linear", + } + + for org_key, diff_key in extra_mappings.items(): + for lora_key in lora_keys: + original_key = f"{org_key}.{lora_key}.weight" + if original_key in original_state_dict: + converted_state_dict[f"{diff_key}.{lora_key}.weight"] = original_state_dict.pop(original_key) if len(original_state_dict) > 0: raise ValueError(f"`original_state_dict` should be empty at this point but has {original_state_dict.keys()=}.") @@ -2432,6 +2519,13 @@ def normalize_out_key(k: str) -> str: if has_default: state_dict = {k.replace("default.", ""): v for k, v in state_dict.items()} + # Normalize ZImage-specific dot-separated module names to underscore form so they + # match the diffusers model parameter names (context_refiner, noise_refiner). + state_dict = { + k.replace("context.refiner.", "context_refiner.").replace("noise.refiner.", "noise_refiner."): v + for k, v in state_dict.items() + } + converted_state_dict = {} all_keys = list(state_dict.keys()) down_key = ".lora_down.weight" @@ -2442,19 +2536,18 @@ def normalize_out_key(k: str) -> str: has_non_diffusers_lora_id = any(down_key in k or up_key in k for k in all_keys) has_diffusers_lora_id = any(a_key in k or b_key in k for k in all_keys) - if has_non_diffusers_lora_id: - - def get_alpha_scales(down_weight, alpha_key): - rank = down_weight.shape[0] - alpha = state_dict.pop(alpha_key).item() - scale = alpha / rank # LoRA is scaled by 'alpha / rank' in forward pass, so we need to scale it back here - scale_down = scale - scale_up = 1.0 - while scale_down * 2 < scale_up: - scale_down *= 2 - scale_up /= 2 - return scale_down, scale_up + def get_alpha_scales(down_weight, alpha_key): + rank = down_weight.shape[0] + alpha = state_dict.pop(alpha_key).item() + scale = alpha / rank # LoRA is scaled by 'alpha / rank' in forward pass, so we need to scale it back here + scale_down = scale + scale_up = 1.0 + while scale_down * 2 < scale_up: + scale_down *= 2 + scale_up /= 2 + return scale_down, scale_up + if has_non_diffusers_lora_id: for k in all_keys: if k.endswith(down_key): diffusers_down_key = k.replace(down_key, ".lora_A.weight") @@ -2467,13 +2560,69 @@ def get_alpha_scales(down_weight, alpha_key): converted_state_dict[diffusers_down_key] = down_weight * scale_down converted_state_dict[diffusers_up_key] = up_weight * scale_up - # Already in diffusers format (lora_A/lora_B), just pop + # Already in diffusers format (lora_A/lora_B), apply alpha scaling and pop. elif has_diffusers_lora_id: for k in all_keys: - if a_key in k or b_key in k: - converted_state_dict[k] = state_dict.pop(k) - elif ".alpha" in k: + if k.endswith(a_key): + diffusers_up_key = k.replace(a_key, b_key) + alpha_key = k.replace(a_key, ".alpha") + + down_weight = state_dict.pop(k) + up_weight = state_dict.pop(diffusers_up_key) + scale_down, scale_up = get_alpha_scales(down_weight, alpha_key) + converted_state_dict[k] = down_weight * scale_down + converted_state_dict[diffusers_up_key] = up_weight * scale_up + + # Handle dot-format LoRA keys: ".lora.down.weight" / ".lora.up.weight". + # Some external ZImage trainers (e.g. Anime-Z) use dots instead of underscores in + # lora weight names and also include redundant keys: + # - "qkv.lora.*" duplicates individual "to.q/k/v.lora.*" keys → skip qkv + # - "out.lora.*" duplicates "to_out.0.lora.*" keys → skip bare out + # - "to.q/k/v.lora.*" → normalise to "to_q/k/v.lora_A/B.weight" + lora_dot_down_key = ".lora.down.weight" + lora_dot_up_key = ".lora.up.weight" + has_lora_dot_format = any(lora_dot_down_key in k for k in state_dict) + + if has_lora_dot_format: + dot_keys = list(state_dict.keys()) + for k in dot_keys: + if lora_dot_down_key not in k: + continue + if k not in state_dict: + continue # already popped by a prior iteration + + base = k[: -len(lora_dot_down_key)] + + # Skip combined "qkv" projection — individual to.q/k/v keys are also present. + if base.endswith(".qkv"): state_dict.pop(k) + state_dict.pop(k.replace(lora_dot_down_key, lora_dot_up_key), None) + state_dict.pop(base + ".alpha", None) + continue + + # Skip bare "out.lora.*" — "to_out.0.lora.*" covers the same projection. + if re.search(r"\.out$", base) and ".to_out" not in base: + state_dict.pop(k) + state_dict.pop(k.replace(lora_dot_down_key, lora_dot_up_key), None) + continue + + # Normalise "to.q/k/v" → "to_q/k/v" for the diffusers output key. + norm_k = re.sub( + r"\.to\.([qkv])" + re.escape(lora_dot_down_key) + r"$", + r".to_\1" + lora_dot_down_key, + k, + ) + norm_base = norm_k[: -len(lora_dot_down_key)] + alpha_key = norm_base + ".alpha" + + diffusers_down = norm_k.replace(lora_dot_down_key, ".lora_A.weight") + diffusers_up = norm_k.replace(lora_dot_down_key, ".lora_B.weight") + + down_weight = state_dict.pop(k) + up_weight = state_dict.pop(k.replace(lora_dot_down_key, lora_dot_up_key)) + scale_down, scale_up = get_alpha_scales(down_weight, alpha_key) + converted_state_dict[diffusers_down] = down_weight * scale_down + converted_state_dict[diffusers_up] = up_weight * scale_up if len(state_dict) > 0: raise ValueError(f"`state_dict` should be empty at this point but has {state_dict.keys()=}") diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py index 03a2fe9f3f8e..5d10f596f2e6 100644 --- a/src/diffusers/loaders/lora_pipeline.py +++ b/src/diffusers/loaders/lora_pipeline.py @@ -13,7 +13,7 @@ # limitations under the License. import os -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import torch from huggingface_hub.utils import validate_hf_hub_args @@ -48,6 +48,7 @@ _convert_non_diffusers_flux2_lora_to_diffusers, _convert_non_diffusers_hidream_lora_to_diffusers, _convert_non_diffusers_lora_to_diffusers, + _convert_non_diffusers_ltx2_lora_to_diffusers, _convert_non_diffusers_ltxv_lora_to_diffusers, _convert_non_diffusers_lumina2_lora_to_diffusers, _convert_non_diffusers_qwen_lora_to_diffusers, @@ -74,6 +75,7 @@ TEXT_ENCODER_NAME = "text_encoder" UNET_NAME = "unet" TRANSFORMER_NAME = "transformer" +LTX2_CONNECTOR_NAME = "connectors" _MODULE_NAME_TO_ATTRIBUTE_MAP_FLUX = {"x_embedder": "in_channels"} @@ -139,8 +141,8 @@ class StableDiffusionLoraLoaderMixin(LoraBaseMixin): def load_lora_weights( self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, hotswap: bool = False, **kwargs, ): @@ -212,7 +214,7 @@ def load_lora_weights( is_correct_format = all("lora" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint.") + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") self.load_lora_into_unet( state_dict, @@ -242,7 +244,7 @@ def load_lora_weights( @validate_hf_hub_args def lora_state_dict( cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs, ): r""" @@ -262,14 +264,14 @@ def lora_state_dict( - A [torch state dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. local_files_only (`bool`, *optional*, defaults to `False`): @@ -369,7 +371,7 @@ def load_lora_into_unet( A standard state dict containing the lora layer parameters. The keys can either be indexed directly into the unet or prefixed with an additional `unet` which can be used to distinguish between text encoder lora layers. - network_alphas (`Dict[str, float]`): + network_alphas (`dict[str, float]`): The value of the network alpha used for stable learning and preventing underflow. This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning). @@ -431,7 +433,7 @@ def load_lora_into_text_encoder( state_dict (`dict`): A standard state dict containing the lora layer parameters. The key should be prefixed with an additional `text_encoder` to distinguish between unet lora layers. - network_alphas (`Dict[str, float]`): + network_alphas (`dict[str, float]`): The value of the network alpha used for stable learning and preventing underflow. This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning). @@ -471,9 +473,9 @@ def load_lora_into_text_encoder( @classmethod def save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - unet_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, - text_encoder_lora_layers: Dict[str, torch.nn.Module] = None, + save_directory: str | os.PathLike, + unet_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, + text_encoder_lora_layers: dict[str, torch.nn.Module] = None, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, @@ -487,9 +489,9 @@ def save_lora_weights( Arguments: save_directory (`str` or `os.PathLike`): Directory to save LoRA parameters to. Will be created if it doesn't exist. - unet_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): + unet_lora_layers (`dict[str, torch.nn.Module]` or `dict[str, torch.Tensor]`): State dict of the LoRA layers corresponding to the `unet`. - text_encoder_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): + text_encoder_lora_layers (`dict[str, torch.nn.Module]` or `dict[str, torch.Tensor]`): State dict of the LoRA layers corresponding to the `text_encoder`. Must explicitly pass the text encoder LoRA state dict because it comes from 🤗 Transformers. is_main_process (`bool`, *optional*, defaults to `True`): @@ -533,10 +535,10 @@ def save_lora_weights( def fuse_lora( self, - components: List[str] = ["unet", "text_encoder"], + components: list[str] = ["unet", "text_encoder"], lora_scale: float = 1.0, safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, + adapter_names: list[str] | None = None, **kwargs, ): r""" @@ -545,12 +547,12 @@ def fuse_lora( > [!WARNING] > This is an experimental API. Args: - components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into. + components: (`list[str]`): list of LoRA-injectable components to fuse the LoRAs into. lora_scale (`float`, defaults to 1.0): Controls how much to influence the outputs with the LoRA parameters. safe_fusing (`bool`, defaults to `False`): Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them. - adapter_names (`List[str]`, *optional*): + adapter_names (`list[str]`, *optional*): Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused. Example: @@ -574,7 +576,7 @@ def fuse_lora( **kwargs, ) - def unfuse_lora(self, components: List[str] = ["unet", "text_encoder"], **kwargs): + def unfuse_lora(self, components: list[str] = ["unet", "text_encoder"], **kwargs): r""" Reverses the effect of [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). @@ -582,7 +584,7 @@ def unfuse_lora(self, components: List[str] = ["unet", "text_encoder"], **kwargs > [!WARNING] > This is an experimental API. Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. + components (`list[str]`): list of LoRA-injectable components to unfuse LoRA from. unfuse_unet (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters. unfuse_text_encoder (`bool`, defaults to `True`): Whether to unfuse the text encoder LoRA parameters. If the text encoder wasn't monkey-patched with the @@ -604,8 +606,8 @@ class StableDiffusionXLLoraLoaderMixin(LoraBaseMixin): def load_lora_weights( self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, hotswap: bool = False, **kwargs, ): @@ -639,7 +641,7 @@ def load_lora_weights( is_correct_format = all("lora" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint.") + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") self.load_lora_into_unet( state_dict, @@ -681,7 +683,7 @@ def load_lora_weights( # Copied from diffusers.loaders.lora_pipeline.StableDiffusionLoraLoaderMixin.lora_state_dict def lora_state_dict( cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs, ): r""" @@ -701,14 +703,14 @@ def lora_state_dict( - A [torch state dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. local_files_only (`bool`, *optional*, defaults to `False`): @@ -809,7 +811,7 @@ def load_lora_into_unet( A standard state dict containing the lora layer parameters. The keys can either be indexed directly into the unet or prefixed with an additional `unet` which can be used to distinguish between text encoder lora layers. - network_alphas (`Dict[str, float]`): + network_alphas (`dict[str, float]`): The value of the network alpha used for stable learning and preventing underflow. This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning). @@ -872,7 +874,7 @@ def load_lora_into_text_encoder( state_dict (`dict`): A standard state dict containing the lora layer parameters. The key should be prefixed with an additional `text_encoder` to distinguish between unet lora layers. - network_alphas (`Dict[str, float]`): + network_alphas (`dict[str, float]`): The value of the network alpha used for stable learning and preventing underflow. This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning). @@ -912,10 +914,10 @@ def load_lora_into_text_encoder( @classmethod def save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - unet_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, - text_encoder_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, - text_encoder_2_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + save_directory: str | os.PathLike, + unet_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, + text_encoder_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, + text_encoder_2_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, @@ -959,10 +961,10 @@ def save_lora_weights( def fuse_lora( self, - components: List[str] = ["unet", "text_encoder", "text_encoder_2"], + components: list[str] = ["unet", "text_encoder", "text_encoder_2"], lora_scale: float = 1.0, safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, + adapter_names: list[str] | None = None, **kwargs, ): r""" @@ -976,7 +978,7 @@ def fuse_lora( **kwargs, ) - def unfuse_lora(self, components: List[str] = ["unet", "text_encoder", "text_encoder_2"], **kwargs): + def unfuse_lora(self, components: list[str] = ["unet", "text_encoder", "text_encoder_2"], **kwargs): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ @@ -1000,7 +1002,7 @@ class SD3LoraLoaderMixin(LoraBaseMixin): @validate_hf_hub_args def lora_state_dict( cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs, ): r""" @@ -1052,7 +1054,7 @@ def lora_state_dict( def load_lora_weights( self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], adapter_name=None, hotswap: bool = False, **kwargs, @@ -1079,7 +1081,7 @@ def load_lora_weights( is_correct_format = all("lora" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint.") + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") self.load_lora_into_transformer( state_dict, @@ -1168,7 +1170,7 @@ def load_lora_into_text_encoder( state_dict (`dict`): A standard state dict containing the lora layer parameters. The key should be prefixed with an additional `text_encoder` to distinguish between unet lora layers. - network_alphas (`Dict[str, float]`): + network_alphas (`dict[str, float]`): The value of the network alpha used for stable learning and preventing underflow. This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning). @@ -1209,10 +1211,10 @@ def load_lora_into_text_encoder( # Copied from diffusers.loaders.lora_pipeline.StableDiffusionXLLoraLoaderMixin.save_lora_weights with unet->transformer def save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, - text_encoder_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, - text_encoder_2_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, + text_encoder_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, + text_encoder_2_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, @@ -1257,10 +1259,10 @@ def save_lora_weights( # Copied from diffusers.loaders.lora_pipeline.StableDiffusionXLLoraLoaderMixin.fuse_lora with unet->transformer def fuse_lora( self, - components: List[str] = ["transformer", "text_encoder", "text_encoder_2"], + components: list[str] = ["transformer", "text_encoder", "text_encoder_2"], lora_scale: float = 1.0, safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, + adapter_names: list[str] | None = None, **kwargs, ): r""" @@ -1275,7 +1277,7 @@ def fuse_lora( ) # Copied from diffusers.loaders.lora_pipeline.StableDiffusionXLLoraLoaderMixin.unfuse_lora with unet->transformer - def unfuse_lora(self, components: List[str] = ["transformer", "text_encoder", "text_encoder_2"], **kwargs): + def unfuse_lora(self, components: list[str] = ["transformer", "text_encoder", "text_encoder_2"], **kwargs): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ @@ -1295,7 +1297,7 @@ class AuraFlowLoraLoaderMixin(LoraBaseMixin): # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.lora_state_dict def lora_state_dict( cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs, ): r""" @@ -1348,8 +1350,8 @@ def lora_state_dict( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights def load_lora_weights( self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, hotswap: bool = False, **kwargs, ): @@ -1375,7 +1377,7 @@ def load_lora_weights( is_correct_format = all("lora" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint.") + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") self.load_lora_into_transformer( state_dict, @@ -1423,13 +1425,13 @@ def load_lora_into_transformer( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights def save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, safe_serialization: bool = True, - transformer_lora_adapter_metadata: Optional[dict] = None, + transformer_lora_adapter_metadata: dict | None = None, ): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. @@ -1457,10 +1459,10 @@ def save_lora_weights( # Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.fuse_lora def fuse_lora( self, - components: List[str] = ["transformer"], + components: list[str] = ["transformer"], lora_scale: float = 1.0, safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, + adapter_names: list[str] | None = None, **kwargs, ): r""" @@ -1475,7 +1477,7 @@ def fuse_lora( ) # Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.unfuse_lora - def unfuse_lora(self, components: List[str] = ["transformer", "text_encoder"], **kwargs): + def unfuse_lora(self, components: list[str] = ["transformer", "text_encoder"], **kwargs): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ @@ -1499,7 +1501,7 @@ class FluxLoraLoaderMixin(LoraBaseMixin): @validate_hf_hub_args def lora_state_dict( cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], return_alphas: bool = False, **kwargs, ): @@ -1622,8 +1624,8 @@ def lora_state_dict( def load_lora_weights( self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, hotswap: bool = False, **kwargs, ): @@ -1657,7 +1659,7 @@ def load_lora_weights( ) if not (has_lora_keys or has_norm_keys): - raise ValueError("Invalid LoRA checkpoint.") + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") transformer_lora_state_dict = { k: state_dict.get(k) @@ -1761,7 +1763,7 @@ def _load_norm_into_transformer( transformer, prefix=None, discard_original_layers=False, - ) -> Dict[str, torch.Tensor]: + ) -> dict[str, torch.Tensor]: # Remove prefix if present prefix = prefix or cls.transformer_name for key in list(state_dict.keys()): @@ -1830,7 +1832,7 @@ def load_lora_into_text_encoder( state_dict (`dict`): A standard state dict containing the lora layer parameters. The key should be prefixed with an additional `text_encoder` to distinguish between unet lora layers. - network_alphas (`Dict[str, float]`): + network_alphas (`dict[str, float]`): The value of the network alpha used for stable learning and preventing underflow. This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning). @@ -1871,9 +1873,9 @@ def load_lora_into_text_encoder( # Copied from diffusers.loaders.lora_pipeline.StableDiffusionLoraLoaderMixin.save_lora_weights with unet->transformer def save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, - text_encoder_lora_layers: Dict[str, torch.nn.Module] = None, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, + text_encoder_lora_layers: dict[str, torch.nn.Module] = None, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, @@ -1887,9 +1889,9 @@ def save_lora_weights( Arguments: save_directory (`str` or `os.PathLike`): Directory to save LoRA parameters to. Will be created if it doesn't exist. - transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): + transformer_lora_layers (`dict[str, torch.nn.Module]` or `dict[str, torch.Tensor]`): State dict of the LoRA layers corresponding to the `transformer`. - text_encoder_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): + text_encoder_lora_layers (`dict[str, torch.nn.Module]` or `dict[str, torch.Tensor]`): State dict of the LoRA layers corresponding to the `text_encoder`. Must explicitly pass the text encoder LoRA state dict because it comes from 🤗 Transformers. is_main_process (`bool`, *optional*, defaults to `True`): @@ -1933,10 +1935,10 @@ def save_lora_weights( def fuse_lora( self, - components: List[str] = ["transformer"], + components: list[str] = ["transformer"], lora_scale: float = 1.0, safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, + adapter_names: list[str] | None = None, **kwargs, ): r""" @@ -1963,7 +1965,7 @@ def fuse_lora( **kwargs, ) - def unfuse_lora(self, components: List[str] = ["transformer", "text_encoder"], **kwargs): + def unfuse_lora(self, components: list[str] = ["transformer", "text_encoder"], **kwargs): r""" Reverses the effect of [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). @@ -1971,7 +1973,7 @@ def unfuse_lora(self, components: List[str] = ["transformer", "text_encoder"], * > [!WARNING] > This is an experimental API. Args: - components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from. + components (`list[str]`): list of LoRA-injectable components to unfuse LoRA from. """ transformer = getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer if hasattr(transformer, "_transformer_norm_layers") and transformer._transformer_norm_layers: @@ -2320,7 +2322,7 @@ def load_lora_into_text_encoder( state_dict (`dict`): A standard state dict containing the lora layer parameters. The key should be prefixed with an additional `text_encoder` to distinguish between unet lora layers. - network_alphas (`Dict[str, float]`): + network_alphas (`dict[str, float]`): The value of the network alpha used for stable learning and preventing underflow. This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning). @@ -2360,9 +2362,9 @@ def load_lora_into_text_encoder( @classmethod def save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - text_encoder_lora_layers: Dict[str, torch.nn.Module] = None, - transformer_lora_layers: Dict[str, torch.nn.Module] = None, + save_directory: str | os.PathLike, + text_encoder_lora_layers: dict[str, torch.nn.Module] = None, + transformer_lora_layers: dict[str, torch.nn.Module] = None, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, @@ -2374,9 +2376,9 @@ def save_lora_weights( Arguments: save_directory (`str` or `os.PathLike`): Directory to save LoRA parameters to. Will be created if it doesn't exist. - unet_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): + unet_lora_layers (`dict[str, torch.nn.Module]` or `dict[str, torch.Tensor]`): State dict of the LoRA layers corresponding to the `unet`. - text_encoder_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`): + text_encoder_lora_layers (`dict[str, torch.nn.Module]` or `dict[str, torch.Tensor]`): State dict of the LoRA layers corresponding to the `text_encoder`. Must explicitly pass the text encoder LoRA state dict because it comes from 🤗 Transformers. is_main_process (`bool`, *optional*, defaults to `True`): @@ -2425,7 +2427,7 @@ class CogVideoXLoraLoaderMixin(LoraBaseMixin): # Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.lora_state_dict def lora_state_dict( cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs, ): r""" @@ -2477,8 +2479,8 @@ def lora_state_dict( def load_lora_weights( self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, hotswap: bool = False, **kwargs, ): @@ -2504,7 +2506,7 @@ def load_lora_weights( is_correct_format = all("lora" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint.") + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") self.load_lora_into_transformer( state_dict, @@ -2551,13 +2553,13 @@ def load_lora_into_transformer( @classmethod def save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, safe_serialization: bool = True, - transformer_lora_adapter_metadata: Optional[dict] = None, + transformer_lora_adapter_metadata: dict | None = None, ): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. @@ -2584,10 +2586,10 @@ def save_lora_weights( def fuse_lora( self, - components: List[str] = ["transformer"], + components: list[str] = ["transformer"], lora_scale: float = 1.0, safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, + adapter_names: list[str] | None = None, **kwargs, ): r""" @@ -2601,7 +2603,7 @@ def fuse_lora( **kwargs, ) - def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): + def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ @@ -2621,7 +2623,7 @@ class Mochi1LoraLoaderMixin(LoraBaseMixin): # Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.lora_state_dict def lora_state_dict( cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs, ): r""" @@ -2674,8 +2676,8 @@ def lora_state_dict( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights def load_lora_weights( self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, hotswap: bool = False, **kwargs, ): @@ -2701,7 +2703,7 @@ def load_lora_weights( is_correct_format = all("lora" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint.") + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") self.load_lora_into_transformer( state_dict, @@ -2749,13 +2751,13 @@ def load_lora_into_transformer( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights def save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, safe_serialization: bool = True, - transformer_lora_adapter_metadata: Optional[dict] = None, + transformer_lora_adapter_metadata: dict | None = None, ): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. @@ -2783,10 +2785,10 @@ def save_lora_weights( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora def fuse_lora( self, - components: List[str] = ["transformer"], + components: list[str] = ["transformer"], lora_scale: float = 1.0, safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, + adapter_names: list[str] | None = None, **kwargs, ): r""" @@ -2801,7 +2803,7 @@ def fuse_lora( ) # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora - def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): + def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ @@ -2820,7 +2822,7 @@ class LTXVideoLoraLoaderMixin(LoraBaseMixin): @validate_hf_hub_args def lora_state_dict( cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs, ): r""" @@ -2877,8 +2879,8 @@ def lora_state_dict( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights def load_lora_weights( self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, hotswap: bool = False, **kwargs, ): @@ -2904,7 +2906,7 @@ def load_lora_weights( is_correct_format = all("lora" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint.") + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") self.load_lora_into_transformer( state_dict, @@ -2952,13 +2954,13 @@ def load_lora_into_transformer( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights def save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, safe_serialization: bool = True, - transformer_lora_adapter_metadata: Optional[dict] = None, + transformer_lora_adapter_metadata: dict | None = None, ): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. @@ -2986,10 +2988,10 @@ def save_lora_weights( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora def fuse_lora( self, - components: List[str] = ["transformer"], + components: list[str] = ["transformer"], lora_scale: float = 1.0, safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, + adapter_names: list[str] | None = None, **kwargs, ): r""" @@ -3004,7 +3006,234 @@ def fuse_lora( ) # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora - def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): + def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs): + r""" + See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. + """ + super().unfuse_lora(components=components, **kwargs) + + +class LTX2LoraLoaderMixin(LoraBaseMixin): + r""" + Load LoRA layers into [`LTX2VideoTransformer3DModel`]. Specific to [`LTX2Pipeline`]. + """ + + _lora_loadable_modules = ["transformer", "connectors"] + transformer_name = TRANSFORMER_NAME + connectors_name = LTX2_CONNECTOR_NAME + + @classmethod + @validate_hf_hub_args + def lora_state_dict( + cls, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + **kwargs, + ): + r""" + See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details. + """ + # Load the main state dict first which has the LoRA layers for either of + # transformer and text encoder or both. + cache_dir = kwargs.pop("cache_dir", None) + force_download = kwargs.pop("force_download", False) + proxies = kwargs.pop("proxies", None) + local_files_only = kwargs.pop("local_files_only", None) + token = kwargs.pop("token", None) + revision = kwargs.pop("revision", None) + subfolder = kwargs.pop("subfolder", None) + weight_name = kwargs.pop("weight_name", None) + use_safetensors = kwargs.pop("use_safetensors", None) + return_lora_metadata = kwargs.pop("return_lora_metadata", False) + + allow_pickle = False + if use_safetensors is None: + use_safetensors = True + allow_pickle = True + + user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"} + + state_dict, metadata = _fetch_state_dict( + pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict, + weight_name=weight_name, + use_safetensors=use_safetensors, + local_files_only=local_files_only, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + token=token, + revision=revision, + subfolder=subfolder, + user_agent=user_agent, + allow_pickle=allow_pickle, + ) + + is_dora_scale_present = any("dora_scale" in k for k in state_dict) + if is_dora_scale_present: + warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new." + logger.warning(warn_msg) + state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k} + + final_state_dict = state_dict + is_non_diffusers_format = any(k.startswith("diffusion_model.") for k in state_dict) + has_connector = any(k.startswith("text_embedding_projection.") for k in state_dict) + if is_non_diffusers_format: + final_state_dict = _convert_non_diffusers_ltx2_lora_to_diffusers(state_dict) + if has_connector: + connectors_state_dict = _convert_non_diffusers_ltx2_lora_to_diffusers( + state_dict, "text_embedding_projection" + ) + final_state_dict.update(connectors_state_dict) + out = (final_state_dict, metadata) if return_lora_metadata else final_state_dict + return out + + def load_lora_weights( + self, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, + hotswap: bool = False, + **kwargs, + ): + """ + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details. + """ + if not USE_PEFT_BACKEND: + raise ValueError("PEFT backend is required for this method.") + + low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA) + if low_cpu_mem_usage and is_peft_version("<", "0.13.0"): + raise ValueError( + "`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`." + ) + + # if a dict is passed, copy it instead of modifying it inplace + if isinstance(pretrained_model_name_or_path_or_dict, dict): + pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy() + + # First, ensure that the checkpoint is a compatible one and can be successfully loaded. + kwargs["return_lora_metadata"] = True + state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) + + is_correct_format = all("lora" in key for key in state_dict.keys()) + if not is_correct_format: + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + + transformer_peft_state_dict = { + k: v for k, v in state_dict.items() if k.startswith(f"{self.transformer_name}.") + } + connectors_peft_state_dict = {k: v for k, v in state_dict.items() if k.startswith(f"{self.connectors_name}.")} + self.load_lora_into_transformer( + transformer_peft_state_dict, + transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer, + adapter_name=adapter_name, + metadata=metadata, + _pipeline=self, + low_cpu_mem_usage=low_cpu_mem_usage, + hotswap=hotswap, + ) + if connectors_peft_state_dict: + self.load_lora_into_transformer( + connectors_peft_state_dict, + transformer=getattr(self, self.connectors_name) + if not hasattr(self, "connectors") + else self.connectors, + adapter_name=adapter_name, + metadata=metadata, + _pipeline=self, + low_cpu_mem_usage=low_cpu_mem_usage, + hotswap=hotswap, + prefix=self.connectors_name, + ) + + @classmethod + def load_lora_into_transformer( + cls, + state_dict, + transformer, + adapter_name=None, + _pipeline=None, + low_cpu_mem_usage=False, + hotswap: bool = False, + metadata=None, + prefix: str = "transformer", + ): + """ + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details. + """ + if low_cpu_mem_usage and is_peft_version("<", "0.13.0"): + raise ValueError( + "`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`." + ) + + # Load the layers corresponding to transformer. + logger.info(f"Loading {prefix}.") + transformer.load_lora_adapter( + state_dict, + network_alphas=None, + adapter_name=adapter_name, + metadata=metadata, + _pipeline=_pipeline, + low_cpu_mem_usage=low_cpu_mem_usage, + hotswap=hotswap, + prefix=prefix, + ) + + @classmethod + # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights + def save_lora_weights( + cls, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, + is_main_process: bool = True, + weight_name: str = None, + save_function: Callable = None, + safe_serialization: bool = True, + transformer_lora_adapter_metadata: dict | None = None, + ): + r""" + See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. + """ + lora_layers = {} + lora_metadata = {} + + if transformer_lora_layers: + lora_layers[cls.transformer_name] = transformer_lora_layers + lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata + + if not lora_layers: + raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.") + + cls._save_lora_weights( + save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, + is_main_process=is_main_process, + weight_name=weight_name, + save_function=save_function, + safe_serialization=safe_serialization, + ) + + # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora + def fuse_lora( + self, + components: list[str] = ["transformer"], + lora_scale: float = 1.0, + safe_fusing: bool = False, + adapter_names: list[str] | None = None, + **kwargs, + ): + r""" + See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details. + """ + super().fuse_lora( + components=components, + lora_scale=lora_scale, + safe_fusing=safe_fusing, + adapter_names=adapter_names, + **kwargs, + ) + + # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora + def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ @@ -3024,7 +3253,7 @@ class SanaLoraLoaderMixin(LoraBaseMixin): # Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.lora_state_dict def lora_state_dict( cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs, ): r""" @@ -3077,8 +3306,8 @@ def lora_state_dict( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights def load_lora_weights( self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, hotswap: bool = False, **kwargs, ): @@ -3104,7 +3333,7 @@ def load_lora_weights( is_correct_format = all("lora" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint.") + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") self.load_lora_into_transformer( state_dict, @@ -3152,13 +3381,13 @@ def load_lora_into_transformer( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights def save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, safe_serialization: bool = True, - transformer_lora_adapter_metadata: Optional[dict] = None, + transformer_lora_adapter_metadata: dict | None = None, ): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. @@ -3186,10 +3415,10 @@ def save_lora_weights( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora def fuse_lora( self, - components: List[str] = ["transformer"], + components: list[str] = ["transformer"], lora_scale: float = 1.0, safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, + adapter_names: list[str] | None = None, **kwargs, ): r""" @@ -3204,7 +3433,208 @@ def fuse_lora( ) # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora - def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): + def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs): + r""" + See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. + """ + super().unfuse_lora(components=components, **kwargs) + + +class HeliosLoraLoaderMixin(LoraBaseMixin): + r""" + Load LoRA layers into [`HeliosTransformer3DModel`]. Specific to [`HeliosPipeline`] and [`HeliosPyramidPipeline`]. + """ + + _lora_loadable_modules = ["transformer"] + transformer_name = TRANSFORMER_NAME + + @classmethod + @validate_hf_hub_args + def lora_state_dict( + cls, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + **kwargs, + ): + r""" + See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details. + """ + # Load the main state dict first which has the LoRA layers for either of + # transformer and text encoder or both. + cache_dir = kwargs.pop("cache_dir", None) + force_download = kwargs.pop("force_download", False) + proxies = kwargs.pop("proxies", None) + local_files_only = kwargs.pop("local_files_only", None) + token = kwargs.pop("token", None) + revision = kwargs.pop("revision", None) + subfolder = kwargs.pop("subfolder", None) + weight_name = kwargs.pop("weight_name", None) + use_safetensors = kwargs.pop("use_safetensors", None) + return_lora_metadata = kwargs.pop("return_lora_metadata", False) + + allow_pickle = False + if use_safetensors is None: + use_safetensors = True + allow_pickle = True + + user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"} + + state_dict, metadata = _fetch_state_dict( + pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict, + weight_name=weight_name, + use_safetensors=use_safetensors, + local_files_only=local_files_only, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + token=token, + revision=revision, + subfolder=subfolder, + user_agent=user_agent, + allow_pickle=allow_pickle, + ) + if any(k.startswith("diffusion_model.") for k in state_dict): + state_dict = _convert_non_diffusers_wan_lora_to_diffusers(state_dict) + elif any(k.startswith("lora_unet_") for k in state_dict): + state_dict = _convert_musubi_wan_lora_to_diffusers(state_dict) + + is_dora_scale_present = any("dora_scale" in k for k in state_dict) + if is_dora_scale_present: + warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new." + logger.warning(warn_msg) + state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k} + + out = (state_dict, metadata) if return_lora_metadata else state_dict + return out + + def load_lora_weights( + self, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, + hotswap: bool = False, + **kwargs, + ): + """ + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details. + """ + if not USE_PEFT_BACKEND: + raise ValueError("PEFT backend is required for this method.") + + low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA) + if low_cpu_mem_usage and is_peft_version("<", "0.13.0"): + raise ValueError( + "`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`." + ) + + # if a dict is passed, copy it instead of modifying it inplace + if isinstance(pretrained_model_name_or_path_or_dict, dict): + pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy() + + # First, ensure that the checkpoint is a compatible one and can be successfully loaded. + kwargs["return_lora_metadata"] = True + state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) + is_correct_format = all("lora" in key for key in state_dict.keys()) + if not is_correct_format: + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + + self.load_lora_into_transformer( + state_dict, + transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer, + adapter_name=adapter_name, + metadata=metadata, + _pipeline=self, + low_cpu_mem_usage=low_cpu_mem_usage, + hotswap=hotswap, + ) + + @classmethod + # Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.load_lora_into_transformer with SD3Transformer2DModel->WanTransformer3DModel + def load_lora_into_transformer( + cls, + state_dict, + transformer, + adapter_name=None, + _pipeline=None, + low_cpu_mem_usage=False, + hotswap: bool = False, + metadata=None, + ): + """ + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details. + """ + if low_cpu_mem_usage and is_peft_version("<", "0.13.0"): + raise ValueError( + "`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`." + ) + + # Load the layers corresponding to transformer. + logger.info(f"Loading {cls.transformer_name}.") + transformer.load_lora_adapter( + state_dict, + network_alphas=None, + adapter_name=adapter_name, + metadata=metadata, + _pipeline=_pipeline, + low_cpu_mem_usage=low_cpu_mem_usage, + hotswap=hotswap, + ) + + @classmethod + # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights + def save_lora_weights( + cls, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, + is_main_process: bool = True, + weight_name: str = None, + save_function: Callable = None, + safe_serialization: bool = True, + transformer_lora_adapter_metadata: dict | None = None, + ): + r""" + See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. + """ + lora_layers = {} + lora_metadata = {} + + if transformer_lora_layers: + lora_layers[cls.transformer_name] = transformer_lora_layers + lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata + + if not lora_layers: + raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.") + + cls._save_lora_weights( + save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, + is_main_process=is_main_process, + weight_name=weight_name, + save_function=save_function, + safe_serialization=safe_serialization, + ) + + # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora + def fuse_lora( + self, + components: list[str] = ["transformer"], + lora_scale: float = 1.0, + safe_fusing: bool = False, + adapter_names: list[str] | None = None, + **kwargs, + ): + r""" + See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details. + """ + super().fuse_lora( + components=components, + lora_scale=lora_scale, + safe_fusing=safe_fusing, + adapter_names=adapter_names, + **kwargs, + ) + + # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora + def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ @@ -3223,7 +3653,7 @@ class HunyuanVideoLoraLoaderMixin(LoraBaseMixin): @validate_hf_hub_args def lora_state_dict( cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs, ): r""" @@ -3280,8 +3710,8 @@ def lora_state_dict( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights def load_lora_weights( self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, hotswap: bool = False, **kwargs, ): @@ -3307,7 +3737,7 @@ def load_lora_weights( is_correct_format = all("lora" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint.") + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") self.load_lora_into_transformer( state_dict, @@ -3355,13 +3785,13 @@ def load_lora_into_transformer( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights def save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, safe_serialization: bool = True, - transformer_lora_adapter_metadata: Optional[dict] = None, + transformer_lora_adapter_metadata: dict | None = None, ): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. @@ -3389,10 +3819,10 @@ def save_lora_weights( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora def fuse_lora( self, - components: List[str] = ["transformer"], + components: list[str] = ["transformer"], lora_scale: float = 1.0, safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, + adapter_names: list[str] | None = None, **kwargs, ): r""" @@ -3407,7 +3837,7 @@ def fuse_lora( ) # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora - def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): + def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ @@ -3426,7 +3856,7 @@ class Lumina2LoraLoaderMixin(LoraBaseMixin): @validate_hf_hub_args def lora_state_dict( cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs, ): r""" @@ -3484,8 +3914,8 @@ def lora_state_dict( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights def load_lora_weights( self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, hotswap: bool = False, **kwargs, ): @@ -3511,7 +3941,7 @@ def load_lora_weights( is_correct_format = all("lora" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint.") + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") self.load_lora_into_transformer( state_dict, @@ -3559,13 +3989,13 @@ def load_lora_into_transformer( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights def save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, safe_serialization: bool = True, - transformer_lora_adapter_metadata: Optional[dict] = None, + transformer_lora_adapter_metadata: dict | None = None, ): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. @@ -3593,10 +4023,10 @@ def save_lora_weights( # Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.fuse_lora def fuse_lora( self, - components: List[str] = ["transformer"], + components: list[str] = ["transformer"], lora_scale: float = 1.0, safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, + adapter_names: list[str] | None = None, **kwargs, ): r""" @@ -3611,7 +4041,7 @@ def fuse_lora( ) # Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.unfuse_lora - def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): + def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ @@ -3631,7 +4061,7 @@ class KandinskyLoraLoaderMixin(LoraBaseMixin): # Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.lora_state_dict def lora_state_dict( cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs, ): r""" @@ -3684,8 +4114,8 @@ def lora_state_dict( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights def load_lora_weights( self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, hotswap: bool = False, **kwargs, ): @@ -3711,7 +4141,7 @@ def load_lora_weights( is_correct_format = all("lora" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint.") + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") self.load_lora_into_transformer( state_dict, @@ -3759,13 +4189,13 @@ def load_lora_into_transformer( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights def save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, safe_serialization: bool = True, - transformer_lora_adapter_metadata: Optional[dict] = None, + transformer_lora_adapter_metadata: dict | None = None, ): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. @@ -3793,10 +4223,10 @@ def save_lora_weights( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora def fuse_lora( self, - components: List[str] = ["transformer"], + components: list[str] = ["transformer"], lora_scale: float = 1.0, safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, + adapter_names: list[str] | None = None, **kwargs, ): r""" @@ -3811,7 +4241,7 @@ def fuse_lora( ) # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora - def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): + def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ @@ -3830,7 +4260,7 @@ class WanLoraLoaderMixin(LoraBaseMixin): @validate_hf_hub_args def lora_state_dict( cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs, ): r""" @@ -3934,8 +4364,8 @@ def _maybe_expand_t2v_lora_for_i2v( def load_lora_weights( self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, hotswap: bool = False, **kwargs, ): @@ -3965,7 +4395,7 @@ def load_lora_weights( ) is_correct_format = all("lora" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint.") + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") load_into_transformer_2 = kwargs.pop("load_into_transformer_2", False) if load_into_transformer_2: @@ -4033,13 +4463,13 @@ def load_lora_into_transformer( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights def save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, safe_serialization: bool = True, - transformer_lora_adapter_metadata: Optional[dict] = None, + transformer_lora_adapter_metadata: dict | None = None, ): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. @@ -4067,10 +4497,10 @@ def save_lora_weights( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora def fuse_lora( self, - components: List[str] = ["transformer"], + components: list[str] = ["transformer"], lora_scale: float = 1.0, safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, + adapter_names: list[str] | None = None, **kwargs, ): r""" @@ -4085,7 +4515,7 @@ def fuse_lora( ) # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora - def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): + def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ @@ -4105,7 +4535,7 @@ class SkyReelsV2LoraLoaderMixin(LoraBaseMixin): # Copied from diffusers.loaders.lora_pipeline.WanLoraLoaderMixin.lora_state_dict def lora_state_dict( cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs, ): r""" @@ -4211,8 +4641,8 @@ def _maybe_expand_t2v_lora_for_i2v( # Copied from diffusers.loaders.lora_pipeline.WanLoraLoaderMixin.load_lora_weights def load_lora_weights( self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, hotswap: bool = False, **kwargs, ): @@ -4242,7 +4672,7 @@ def load_lora_weights( ) is_correct_format = all("lora" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint.") + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") load_into_transformer_2 = kwargs.pop("load_into_transformer_2", False) if load_into_transformer_2: @@ -4310,13 +4740,13 @@ def load_lora_into_transformer( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights def save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, safe_serialization: bool = True, - transformer_lora_adapter_metadata: Optional[dict] = None, + transformer_lora_adapter_metadata: dict | None = None, ): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. @@ -4344,10 +4774,10 @@ def save_lora_weights( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora def fuse_lora( self, - components: List[str] = ["transformer"], + components: list[str] = ["transformer"], lora_scale: float = 1.0, safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, + adapter_names: list[str] | None = None, **kwargs, ): r""" @@ -4362,7 +4792,7 @@ def fuse_lora( ) # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora - def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): + def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ @@ -4382,7 +4812,7 @@ class CogView4LoraLoaderMixin(LoraBaseMixin): # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.lora_state_dict def lora_state_dict( cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs, ): r""" @@ -4435,8 +4865,8 @@ def lora_state_dict( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights def load_lora_weights( self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, hotswap: bool = False, **kwargs, ): @@ -4462,7 +4892,7 @@ def load_lora_weights( is_correct_format = all("lora" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint.") + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") self.load_lora_into_transformer( state_dict, @@ -4510,13 +4940,13 @@ def load_lora_into_transformer( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights def save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, safe_serialization: bool = True, - transformer_lora_adapter_metadata: Optional[dict] = None, + transformer_lora_adapter_metadata: dict | None = None, ): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. @@ -4544,10 +4974,10 @@ def save_lora_weights( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora def fuse_lora( self, - components: List[str] = ["transformer"], + components: list[str] = ["transformer"], lora_scale: float = 1.0, safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, + adapter_names: list[str] | None = None, **kwargs, ): r""" @@ -4562,7 +4992,7 @@ def fuse_lora( ) # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora - def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): + def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ @@ -4581,7 +5011,7 @@ class HiDreamImageLoraLoaderMixin(LoraBaseMixin): @validate_hf_hub_args def lora_state_dict( cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs, ): r""" @@ -4638,8 +5068,8 @@ def lora_state_dict( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights def load_lora_weights( self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, hotswap: bool = False, **kwargs, ): @@ -4665,7 +5095,7 @@ def load_lora_weights( is_correct_format = all("lora" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint.") + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") self.load_lora_into_transformer( state_dict, @@ -4713,13 +5143,13 @@ def load_lora_into_transformer( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights def save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, safe_serialization: bool = True, - transformer_lora_adapter_metadata: Optional[dict] = None, + transformer_lora_adapter_metadata: dict | None = None, ): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. @@ -4747,10 +5177,10 @@ def save_lora_weights( # Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.fuse_lora def fuse_lora( self, - components: List[str] = ["transformer"], + components: list[str] = ["transformer"], lora_scale: float = 1.0, safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, + adapter_names: list[str] | None = None, **kwargs, ): r""" @@ -4765,7 +5195,7 @@ def fuse_lora( ) # Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.unfuse_lora - def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): + def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ @@ -4784,7 +5214,7 @@ class QwenImageLoraLoaderMixin(LoraBaseMixin): @validate_hf_hub_args def lora_state_dict( cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs, ): r""" @@ -4844,8 +5274,8 @@ def lora_state_dict( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights def load_lora_weights( self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, hotswap: bool = False, **kwargs, ): @@ -4871,7 +5301,7 @@ def load_lora_weights( is_correct_format = all("lora" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint.") + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") self.load_lora_into_transformer( state_dict, @@ -4919,13 +5349,13 @@ def load_lora_into_transformer( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights def save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, safe_serialization: bool = True, - transformer_lora_adapter_metadata: Optional[dict] = None, + transformer_lora_adapter_metadata: dict | None = None, ): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. @@ -4953,10 +5383,10 @@ def save_lora_weights( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora def fuse_lora( self, - components: List[str] = ["transformer"], + components: list[str] = ["transformer"], lora_scale: float = 1.0, safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, + adapter_names: list[str] | None = None, **kwargs, ): r""" @@ -4971,7 +5401,7 @@ def fuse_lora( ) # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora - def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): + def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ @@ -4990,7 +5420,7 @@ class ZImageLoraLoaderMixin(LoraBaseMixin): @validate_hf_hub_args def lora_state_dict( cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs, ): r""" @@ -5050,8 +5480,8 @@ def lora_state_dict( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights def load_lora_weights( self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, hotswap: bool = False, **kwargs, ): @@ -5077,7 +5507,7 @@ def load_lora_weights( is_correct_format = all("lora" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint.") + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") self.load_lora_into_transformer( state_dict, @@ -5125,13 +5555,13 @@ def load_lora_into_transformer( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights def save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, safe_serialization: bool = True, - transformer_lora_adapter_metadata: Optional[dict] = None, + transformer_lora_adapter_metadata: dict | None = None, ): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. @@ -5159,10 +5589,10 @@ def save_lora_weights( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora def fuse_lora( self, - components: List[str] = ["transformer"], + components: list[str] = ["transformer"], lora_scale: float = 1.0, safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, + adapter_names: list[str] | None = None, **kwargs, ): r""" @@ -5177,7 +5607,7 @@ def fuse_lora( ) # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora - def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): + def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ @@ -5196,7 +5626,7 @@ class Flux2LoraLoaderMixin(LoraBaseMixin): @validate_hf_hub_args def lora_state_dict( cls, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs, ): r""" @@ -5243,6 +5673,10 @@ def lora_state_dict( logger.warning(warn_msg) state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k} + is_peft_format = any(k.startswith("base_model.model.") for k in state_dict) + if is_peft_format: + state_dict = {k.replace("base_model.model.", "diffusion_model."): v for k, v in state_dict.items()} + is_ai_toolkit = any(k.startswith("diffusion_model.") for k in state_dict) if is_ai_toolkit: state_dict = _convert_non_diffusers_flux2_lora_to_diffusers(state_dict) @@ -5253,8 +5687,8 @@ def lora_state_dict( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights def load_lora_weights( self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - adapter_name: Optional[str] = None, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, hotswap: bool = False, **kwargs, ): @@ -5280,7 +5714,7 @@ def load_lora_weights( is_correct_format = all("lora" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint.") + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") self.load_lora_into_transformer( state_dict, @@ -5328,13 +5762,13 @@ def load_lora_into_transformer( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights def save_lora_weights( cls, - save_directory: Union[str, os.PathLike], - transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, safe_serialization: bool = True, - transformer_lora_adapter_metadata: Optional[dict] = None, + transformer_lora_adapter_metadata: dict | None = None, ): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. @@ -5362,10 +5796,10 @@ def save_lora_weights( # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora def fuse_lora( self, - components: List[str] = ["transformer"], + components: list[str] = ["transformer"], lora_scale: float = 1.0, safe_fusing: bool = False, - adapter_names: Optional[List[str]] = None, + adapter_names: list[str] | None = None, **kwargs, ): r""" @@ -5380,7 +5814,7 @@ def fuse_lora( ) # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora - def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs): + def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs): r""" See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. """ diff --git a/src/diffusers/loaders/peft.py b/src/diffusers/loaders/peft.py index 30a78f00b3f2..339d10359f04 100644 --- a/src/diffusers/loaders/peft.py +++ b/src/diffusers/loaders/peft.py @@ -17,7 +17,7 @@ import os from functools import partial from pathlib import Path -from typing import Dict, List, Literal, Optional, Union +from typing import Literal import safetensors import torch @@ -51,6 +51,7 @@ "FluxTransformer2DModel": lambda model_cls, weights: weights, "CogVideoXTransformer3DModel": lambda model_cls, weights: weights, "ConsisIDTransformer3DModel": lambda model_cls, weights: weights, + "HeliosTransformer3DModel": lambda model_cls, weights: weights, "MochiTransformer3DModel": lambda model_cls, weights: weights, "HunyuanVideoTransformer3DModel": lambda model_cls, weights: weights, "LTXVideoTransformer3DModel": lambda model_cls, weights: weights, @@ -63,9 +64,13 @@ "HunyuanVideoFramepackTransformer3DModel": lambda model_cls, weights: weights, "WanVACETransformer3DModel": lambda model_cls, weights: weights, "ChromaTransformer2DModel": lambda model_cls, weights: weights, + "ChromaRadianceTransformer2DModel": lambda model_cls, weights: weights, + "ChronoEditTransformer3DModel": lambda model_cls, weights: weights, "QwenImageTransformer2DModel": lambda model_cls, weights: weights, "Flux2Transformer2DModel": lambda model_cls, weights: weights, "ZImageTransformer2DModel": lambda model_cls, weights: weights, + "LTX2VideoTransformer3DModel": lambda model_cls, weights: weights, + "LTX2TextConnectors": lambda model_cls, weights: weights, } @@ -85,7 +90,7 @@ class PeftAdapterMixin: _hf_peft_config_loaded = False # kwargs for prepare_model_for_compiled_hotswap, if required - _prepare_lora_hotswap_kwargs: Optional[dict] = None + _prepare_lora_hotswap_kwargs: dict | None = None @classmethod # Copied from diffusers.loaders.lora_base.LoraBaseMixin._optionally_disable_offloading @@ -111,13 +116,13 @@ def load_lora_adapter( prefix (`str`, *optional*): Prefix to filter the state dict. - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. local_files_only (`bool`, *optional*, defaults to `False`): @@ -131,7 +136,7 @@ def load_lora_adapter( allowed by Git. subfolder (`str`, *optional*, defaults to `""`): The subfolder location of a model file within a larger model repository on the Hub or locally. - network_alphas (`Dict[str, float]`): + network_alphas (`dict[str, float]`): The value of the network alpha used for stable learning and preventing underflow. This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning). @@ -402,7 +407,7 @@ def save_lora_adapter( adapter_name: str = "default", upcast_before_saving: bool = False, safe_serialization: bool = True, - weight_name: Optional[str] = None, + weight_name: str | None = None, ): """ Save the LoRA parameters corresponding to the underlying model. @@ -466,16 +471,16 @@ def save_function(weights, filename): def set_adapters( self, - adapter_names: Union[List[str], str], - weights: Optional[Union[float, Dict, List[float], List[Dict], List[None]]] = None, + adapter_names: list[str] | str, + weights: float | dict | list[float] | list[dict] | list[None] | None = None, ): """ Set the currently active adapters for use in the diffusion network (e.g. unet, transformer, etc.). Args: - adapter_names (`List[str]` or `str`): + adapter_names (`list[str]` or `str`): The names of the adapters to use. - adapter_weights (`Union[List[float], float]`, *optional*): + weights (`Union[List[float], float]`, *optional*): The adapter(s) weights to use with the UNet. If `None`, the weights are set to `1.0` for all the adapters. @@ -492,7 +497,7 @@ def set_adapters( "jbilcke-hf/sdxl-cinematic-1", weight_name="pytorch_lora_weights.safetensors", adapter_name="cinematic" ) pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") - pipeline.unet.set_adapters(["cinematic", "pixel"], adapter_weights=[0.5, 0.5]) + pipeline.unet.set_adapters(["cinematic", "pixel"], weights=[0.5, 0.5]) ``` """ if not USE_PEFT_BACKEND: @@ -558,7 +563,7 @@ def add_adapter(self, adapter_config, adapter_name: str = "default") -> None: inject_adapter_in_model(adapter_config, self, adapter_name) self.set_adapter(adapter_name) - def set_adapter(self, adapter_name: Union[str, List[str]]) -> None: + def set_adapter(self, adapter_name: str | list[str]) -> None: """ Sets a specific adapter by forcing the model to only use that adapter and disables the other adapters. @@ -566,7 +571,7 @@ def set_adapter(self, adapter_name: Union[str, List[str]]) -> None: [documentation](https://huggingface.co/docs/peft). Args: - adapter_name (Union[str, List[str]])): + adapter_name (str | list[str])): The list of adapters to set or the adapter name in the case of a single adapter. """ check_peft_version(min_version=MIN_PEFT_VERSION) @@ -652,7 +657,7 @@ def enable_adapters(self) -> None: # support for older PEFT versions module.disable_adapters = False - def active_adapters(self) -> List[str]: + def active_adapters(self) -> list[str]: """ Gets the current list of active adapters of the model. @@ -775,12 +780,12 @@ def enable_lora(self): raise ValueError("PEFT backend is required for this method.") set_adapter_layers(self, enabled=True) - def delete_adapters(self, adapter_names: Union[List[str], str]): + def delete_adapters(self, adapter_names: list[str] | str): """ Delete an adapter's LoRA layers from the underlying model. Args: - adapter_names (`Union[List[str], str]`): + adapter_names (`list[str, str]`): The names (single string or list of strings) of the adapter to delete. Example: diff --git a/src/diffusers/loaders/single_file.py b/src/diffusers/loaders/single_file.py index 667f79437985..0b0d52d4a412 100644 --- a/src/diffusers/loaders/single_file.py +++ b/src/diffusers/loaders/single_file.py @@ -286,11 +286,11 @@ def from_single_file(cls, pretrained_model_link_or_path, **kwargs) -> Self: force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. local_files_only (`bool`, *optional*, defaults to `False`): diff --git a/src/diffusers/loaders/single_file_model.py b/src/diffusers/loaders/single_file_model.py index 0e4ebab7fe07..8927b3188a6c 100644 --- a/src/diffusers/loaders/single_file_model.py +++ b/src/diffusers/loaders/single_file_model.py @@ -15,7 +15,6 @@ import inspect import re from contextlib import nullcontext -from typing import Optional import torch from huggingface_hub.utils import validate_hf_hub_args @@ -31,6 +30,7 @@ convert_animatediff_checkpoint_to_diffusers, convert_auraflow_transformer_checkpoint_to_diffusers, convert_autoencoder_dc_checkpoint_to_diffusers, + convert_chroma_radiance_transformer_checkpoint_to_diffusers, convert_chroma_transformer_checkpoint_to_diffusers, convert_controlnet_checkpoint, convert_cosmos_transformer_checkpoint_to_diffusers, @@ -40,6 +40,9 @@ convert_hunyuan_video_transformer_to_diffusers, convert_ldm_unet_checkpoint, convert_ldm_vae_checkpoint, + convert_ltx2_audio_vae_to_diffusers, + convert_ltx2_transformer_to_diffusers, + convert_ltx2_vae_to_diffusers, convert_ltx_transformer_checkpoint_to_diffusers, convert_ltx_vae_checkpoint_to_diffusers, convert_lumina2_to_diffusers, @@ -112,6 +115,10 @@ "checkpoint_mapping_fn": convert_chroma_transformer_checkpoint_to_diffusers, "default_subfolder": "transformer", }, + "ChromaRadianceTransformer2DModel": { + "checkpoint_mapping_fn": convert_chroma_radiance_transformer_checkpoint_to_diffusers, + "default_subfolder": "transformer", + }, "LTXVideoTransformer3DModel": { "checkpoint_mapping_fn": convert_ltx_transformer_checkpoint_to_diffusers, "default_subfolder": "transformer", @@ -149,6 +156,10 @@ "checkpoint_mapping_fn": convert_wan_transformer_to_diffusers, "default_subfolder": "transformer", }, + "WanAnimateTransformer3DModel": { + "checkpoint_mapping_fn": convert_wan_transformer_to_diffusers, + "default_subfolder": "transformer", + }, "AutoencoderKLWan": { "checkpoint_mapping_fn": convert_wan_vae_to_diffusers, "default_subfolder": "vae", @@ -162,7 +173,7 @@ "default_subfolder": "transformer", }, "QwenImageTransformer2DModel": { - "checkpoint_mapping_fn": lambda x: x, + "checkpoint_mapping_fn": lambda checkpoint, **kwargs: checkpoint, "default_subfolder": "transformer", }, "Flux2Transformer2DModel": { @@ -176,6 +187,18 @@ "ZImageControlNetModel": { "checkpoint_mapping_fn": convert_z_image_controlnet_checkpoint_to_diffusers, }, + "LTX2VideoTransformer3DModel": { + "checkpoint_mapping_fn": convert_ltx2_transformer_to_diffusers, + "default_subfolder": "transformer", + }, + "AutoencoderKLLTX2Video": { + "checkpoint_mapping_fn": convert_ltx2_vae_to_diffusers, + "default_subfolder": "vae", + }, + "AutoencoderKLLTX2Audio": { + "checkpoint_mapping_fn": convert_ltx2_audio_vae_to_diffusers, + "default_subfolder": "audio_vae", + }, } @@ -216,7 +239,7 @@ class FromOriginalModelMixin: @classmethod @validate_hf_hub_args - def from_single_file(cls, pretrained_model_link_or_path_or_dict: Optional[str] = None, **kwargs) -> Self: + def from_single_file(cls, pretrained_model_link_or_path_or_dict: str | None = None, **kwargs) -> Self: r""" Instantiate a model from pretrained weights saved in the original `.ckpt` or `.safetensors` format. The model is set in evaluation mode (`model.eval()`) by default. @@ -243,11 +266,11 @@ def from_single_file(cls, pretrained_model_link_or_path_or_dict: Optional[str] = force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. local_files_only (`bool`, *optional*, defaults to `False`): diff --git a/src/diffusers/loaders/single_file_utils.py b/src/diffusers/loaders/single_file_utils.py index aac4835fe849..fd6c15639881 100644 --- a/src/diffusers/loaders/single_file_utils.py +++ b/src/diffusers/loaders/single_file_utils.py @@ -112,7 +112,8 @@ "model.diffusion_model.transformer_blocks.27.scale_shift_table", "patchify_proj.weight", "transformer_blocks.27.scale_shift_table", - "vae.per_channel_statistics.mean-of-means", + "vae.decoder.last_scale_shift_table", # 0.9.1, 0.9.5, 0.9.7, 0.9.8 + "vae.decoder.up_blocks.9.res_blocks.0.conv1.conv.weight", # 0.9.0 ], "autoencoder-dc": "decoder.stages.1.op_list.0.main.conv.conv.bias", "autoencoder-dc-sana": "encoder.project_in.conv.bias", @@ -120,7 +121,10 @@ "hunyuan-video": "txt_in.individual_token_refiner.blocks.0.adaLN_modulation.1.bias", "instruct-pix2pix": "model.diffusion_model.input_blocks.0.0.weight", "lumina2": ["model.diffusion_model.cap_embedder.0.weight", "cap_embedder.0.weight"], - "z-image-turbo": "cap_embedder.0.weight", + "z-image-turbo": [ + "model.diffusion_model.layers.0.adaLN_modulation.0.weight", + "layers.0.adaLN_modulation.0.weight", + ], "z-image-turbo-controlnet": "control_all_x_embedder.2-1.weight", "z-image-turbo-controlnet-2.x": "control_layers.14.adaLN_modulation.0.weight", "sana": [ @@ -132,6 +136,7 @@ "wan": ["model.diffusion_model.head.modulation", "head.modulation"], "wan_vae": "decoder.middle.0.residual.0.gamma", "wan_vace": "vace_blocks.0.after_proj.bias", + "wan_animate": "motion_encoder.dec.direction.weight", "hidream": "double_stream_blocks.0.block.adaLN_modulation.1.bias", "cosmos-1.0": [ "net.x_embedder.proj.1.weight", @@ -144,6 +149,11 @@ "net.pos_embedder.dim_spatial_range", ], "flux2": ["model.diffusion_model.single_stream_modulation.lin.weight", "single_stream_modulation.lin.weight"], + "ltx2": [ + "model.diffusion_model.av_ca_a2v_gate_adaln_single.emb.timestep_embedder.linear_1.weight", + "vae.per_channel_statistics.mean-of-means", + "audio_vae.per_channel_statistics.mean-of-means", + ], } DIFFUSERS_DEFAULT_PIPELINE_PATHS = { @@ -194,6 +204,7 @@ "flux-depth": {"pretrained_model_name_or_path": "black-forest-labs/FLUX.1-Depth-dev"}, "flux-schnell": {"pretrained_model_name_or_path": "black-forest-labs/FLUX.1-schnell"}, "flux-2-dev": {"pretrained_model_name_or_path": "black-forest-labs/FLUX.2-dev"}, + "chroma": {"pretrained_model_name_or_path": "lodestones/Chroma1-HD"}, "ltx-video": {"pretrained_model_name_or_path": "diffusers/LTX-Video-0.9.0"}, "ltx-video-0.9.1": {"pretrained_model_name_or_path": "diffusers/LTX-Video-0.9.1"}, "ltx-video-0.9.5": {"pretrained_model_name_or_path": "Lightricks/LTX-Video-0.9.5"}, @@ -210,6 +221,7 @@ "wan-t2v-1.3B": {"pretrained_model_name_or_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"}, "wan-t2v-14B": {"pretrained_model_name_or_path": "Wan-AI/Wan2.1-T2V-14B-Diffusers"}, "wan-i2v-14B": {"pretrained_model_name_or_path": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"}, + "wan-animate-14B": {"pretrained_model_name_or_path": "Wan-AI/Wan2.2-Animate-14B-Diffusers"}, "wan-vace-1.3B": {"pretrained_model_name_or_path": "Wan-AI/Wan2.1-VACE-1.3B-diffusers"}, "wan-vace-14B": {"pretrained_model_name_or_path": "Wan-AI/Wan2.1-VACE-14B-diffusers"}, "hidream": {"pretrained_model_name_or_path": "HiDream-ai/HiDream-I1-Dev"}, @@ -223,7 +235,9 @@ "cosmos-2.0-v2w-14B": {"pretrained_model_name_or_path": "nvidia/Cosmos-Predict2-14B-Video2World"}, "z-image-turbo": {"pretrained_model_name_or_path": "Tongyi-MAI/Z-Image-Turbo"}, "z-image-turbo-controlnet": {"pretrained_model_name_or_path": "hlky/Z-Image-Turbo-Fun-Controlnet-Union"}, - "z-image-turbo-controlnet-2.x": {"pretrained_model_name_or_path": "hlky/Z-Image-Turbo-Fun-Controlnet-Union-2.1"}, + "z-image-turbo-controlnet-2.0": {"pretrained_model_name_or_path": "hlky/Z-Image-Turbo-Fun-Controlnet-Union-2.0"}, + "z-image-turbo-controlnet-2.1": {"pretrained_model_name_or_path": "hlky/Z-Image-Turbo-Fun-Controlnet-Union-2.1"}, + "ltx2-dev": {"pretrained_model_name_or_path": "Lightricks/LTX-2"}, } # Use to configure model sample size when original config is provided @@ -668,7 +682,13 @@ def infer_diffusers_model_type(checkpoint): model_type = "flux-2-dev" elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["flux"]): - if any( + if any(c in checkpoint for c in ["distilled_guidance_layer.in_proj.bias"]): + # Should be updated once a repo exists + # if any(h in checkpoint for h in ["nerf_blocks.0.param_generator.bias"]): + # model_type = "chroma-radiance" + # else: + model_type = "chroma" + elif any( g in checkpoint for g in ["guidance_in.in_layer.bias", "model.diffusion_model.guidance_in.in_layer.bias"] ): if "model.diffusion_model.img_in.weight" in checkpoint: @@ -727,10 +747,7 @@ def infer_diffusers_model_type(checkpoint): ): model_type = "instruct-pix2pix" - elif ( - CHECKPOINT_KEY_NAMES["z-image-turbo"] in checkpoint - and checkpoint[CHECKPOINT_KEY_NAMES["z-image-turbo"]].shape[0] == 2560 - ): + elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["z-image-turbo"]): model_type = "z-image-turbo" elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["lumina2"]): @@ -751,6 +768,9 @@ def infer_diffusers_model_type(checkpoint): elif checkpoint[target_key].shape[0] == 5120: model_type = "wan-vace-14B" + if CHECKPOINT_KEY_NAMES["wan_animate"] in checkpoint: + model_type = "wan-animate-14B" + elif checkpoint[target_key].shape[0] == 1536: model_type = "wan-t2v-1.3B" elif checkpoint[target_key].shape[0] == 5120 and checkpoint[target_key].shape[1] == 16: @@ -784,11 +804,20 @@ def infer_diffusers_model_type(checkpoint): raise ValueError(f"Unexpected x_embedder shape: {x_embedder_shape} when loading Cosmos 2.0 model.") elif CHECKPOINT_KEY_NAMES["z-image-turbo-controlnet-2.x"] in checkpoint: - model_type = "z-image-turbo-controlnet-2.x" + before_proj_weight = checkpoint.get("control_noise_refiner.0.before_proj.weight", None) + if before_proj_weight is None: + model_type = "z-image-turbo-controlnet-2.0" + elif before_proj_weight is not None and torch.all(before_proj_weight == 0.0): + model_type = "z-image-turbo-controlnet-2.0" + else: + model_type = "z-image-turbo-controlnet-2.1" elif CHECKPOINT_KEY_NAMES["z-image-turbo-controlnet"] in checkpoint: model_type = "z-image-turbo-controlnet" + elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["ltx2"]): + model_type = "ltx2-dev" + else: model_type = "v1" @@ -3137,13 +3166,64 @@ def convert_sana_transformer_to_diffusers(checkpoint, **kwargs): def convert_wan_transformer_to_diffusers(checkpoint, **kwargs): + def generate_motion_encoder_mappings(): + mappings = { + "motion_encoder.dec.direction.weight": "motion_encoder.motion_synthesis_weight", + "motion_encoder.enc.net_app.convs.0.0.weight": "motion_encoder.conv_in.weight", + "motion_encoder.enc.net_app.convs.0.1.bias": "motion_encoder.conv_in.act_fn.bias", + "motion_encoder.enc.net_app.convs.8.weight": "motion_encoder.conv_out.weight", + "motion_encoder.enc.fc": "motion_encoder.motion_network", + } + + for i in range(7): + conv_idx = i + 1 + mappings.update( + { + f"motion_encoder.enc.net_app.convs.{conv_idx}.conv1.0.weight": f"motion_encoder.res_blocks.{i}.conv1.weight", + f"motion_encoder.enc.net_app.convs.{conv_idx}.conv1.1.bias": f"motion_encoder.res_blocks.{i}.conv1.act_fn.bias", + f"motion_encoder.enc.net_app.convs.{conv_idx}.conv2.1.weight": f"motion_encoder.res_blocks.{i}.conv2.weight", + f"motion_encoder.enc.net_app.convs.{conv_idx}.conv2.2.bias": f"motion_encoder.res_blocks.{i}.conv2.act_fn.bias", + f"motion_encoder.enc.net_app.convs.{conv_idx}.skip.1.weight": f"motion_encoder.res_blocks.{i}.conv_skip.weight", + } + ) + + return mappings + + def generate_face_adapter_mappings(): + return { + "face_adapter.fuser_blocks": "face_adapter", + ".k_norm.": ".norm_k.", + ".q_norm.": ".norm_q.", + ".linear1_q.": ".to_q.", + ".linear2.": ".to_out.", + "conv1_local.conv": "conv1_local", + "conv2.conv": "conv2", + "conv3.conv": "conv3", + } + + def split_tensor_handler(key, state_dict, split_pattern, target_keys): + tensor = state_dict.pop(key) + split_idx = tensor.shape[0] // 2 + + new_key_1 = key.replace(split_pattern, target_keys[0]) + new_key_2 = key.replace(split_pattern, target_keys[1]) + + state_dict[new_key_1] = tensor[:split_idx] + state_dict[new_key_2] = tensor[split_idx:] + + def reshape_bias_handler(key, state_dict): + if "motion_encoder.enc.net_app.convs." in key and ".bias" in key: + state_dict[key] = state_dict[key][0, :, 0, 0] + converted_state_dict = {} + # Strip model.diffusion_model prefix keys = list(checkpoint.keys()) for k in keys: if "model.diffusion_model." in k: checkpoint[k.replace("model.diffusion_model.", "")] = checkpoint.pop(k) + # Base transformer mappings TRANSFORMER_KEYS_RENAME_DICT = { "time_embedding.0": "condition_embedder.time_embedder.linear_1", "time_embedding.2": "condition_embedder.time_embedder.linear_2", @@ -3165,28 +3245,43 @@ def convert_wan_transformer_to_diffusers(checkpoint, **kwargs): "ffn.0": "ffn.net.0.proj", "ffn.2": "ffn.net.2", # Hack to swap the layer names - # The original model calls the norms in following order: norm1, norm3, norm2 - # We convert it to: norm1, norm2, norm3 "norm2": "norm__placeholder", "norm3": "norm2", "norm__placeholder": "norm3", - # For the I2V model + # I2V model "img_emb.proj.0": "condition_embedder.image_embedder.norm1", "img_emb.proj.1": "condition_embedder.image_embedder.ff.net.0.proj", "img_emb.proj.3": "condition_embedder.image_embedder.ff.net.2", "img_emb.proj.4": "condition_embedder.image_embedder.norm2", - # For the VACE model + # VACE model "before_proj": "proj_in", "after_proj": "proj_out", } + SPECIAL_KEYS_HANDLERS = {} + if any("face_adapter" in k for k in checkpoint.keys()): + TRANSFORMER_KEYS_RENAME_DICT.update(generate_face_adapter_mappings()) + SPECIAL_KEYS_HANDLERS[".linear1_kv."] = (split_tensor_handler, [".to_k.", ".to_v."]) + + if any("motion_encoder" in k for k in checkpoint.keys()): + TRANSFORMER_KEYS_RENAME_DICT.update(generate_motion_encoder_mappings()) + for key in list(checkpoint.keys()): - new_key = key[:] + reshape_bias_handler(key, checkpoint) + + for key in list(checkpoint.keys()): + new_key = key for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items(): new_key = new_key.replace(replace_key, rename_key) - converted_state_dict[new_key] = checkpoint.pop(key) + for key in list(converted_state_dict.keys()): + for pattern, (handler_fn, target_keys) in SPECIAL_KEYS_HANDLERS.items(): + if pattern not in key: + continue + handler_fn(key, converted_state_dict, pattern, target_keys) + break + return converted_state_dict @@ -3566,6 +3661,194 @@ def swap_scale_shift(weight): return converted_state_dict +def convert_chroma_radiance_transformer_checkpoint_to_diffusers(checkpoint, **kwargs): + converted_state_dict = {} + keys = list(checkpoint.keys()) + + for k in keys: + if "model.diffusion_model." in k: + checkpoint[k.replace("model.diffusion_model.", "")] = checkpoint.pop(k) + + num_layers = list(set(int(k.split(".", 2)[1]) for k in checkpoint if "double_blocks." in k))[-1] + 1 # noqa: C401 + num_single_layers = list(set(int(k.split(".", 2)[1]) for k in checkpoint if "single_blocks." in k))[-1] + 1 # noqa: C401 + num_guidance_layers = ( + list(set(int(k.split(".", 3)[2]) for k in checkpoint if "distilled_guidance_layer.layers." in k))[-1] + 1 # noqa: C401 + ) + num_nerf_layers = list(set(int(k.split(".", 2)[1]) for k in checkpoint if "nerf_blocks." in k))[-1] + 1 # noqa: C401 + mlp_ratio = 4.0 + inner_dim = 3072 + + # in SD3 original implementation of AdaLayerNormContinuous, it split linear projection output into shift, scale; + # while in diffusers it split into scale, shift. Here we swap the linear projection weights in order to be able to use diffusers implementation + def swap_scale_shift(weight): + shift, scale = weight.chunk(2, dim=0) + new_weight = torch.cat([scale, shift], dim=0) + return new_weight + + # guidance + converted_state_dict["distilled_guidance_layer.in_proj.bias"] = checkpoint.pop( + "distilled_guidance_layer.in_proj.bias" + ) + converted_state_dict["distilled_guidance_layer.in_proj.weight"] = checkpoint.pop( + "distilled_guidance_layer.in_proj.weight" + ) + converted_state_dict["distilled_guidance_layer.out_proj.bias"] = checkpoint.pop( + "distilled_guidance_layer.out_proj.bias" + ) + converted_state_dict["distilled_guidance_layer.out_proj.weight"] = checkpoint.pop( + "distilled_guidance_layer.out_proj.weight" + ) + for i in range(num_guidance_layers): + block_prefix = f"distilled_guidance_layer.layers.{i}." + converted_state_dict[f"{block_prefix}linear_1.bias"] = checkpoint.pop( + f"distilled_guidance_layer.layers.{i}.in_layer.bias" + ) + converted_state_dict[f"{block_prefix}linear_1.weight"] = checkpoint.pop( + f"distilled_guidance_layer.layers.{i}.in_layer.weight" + ) + converted_state_dict[f"{block_prefix}linear_2.bias"] = checkpoint.pop( + f"distilled_guidance_layer.layers.{i}.out_layer.bias" + ) + converted_state_dict[f"{block_prefix}linear_2.weight"] = checkpoint.pop( + f"distilled_guidance_layer.layers.{i}.out_layer.weight" + ) + converted_state_dict[f"distilled_guidance_layer.norms.{i}.weight"] = checkpoint.pop( + f"distilled_guidance_layer.norms.{i}.scale" + ) + + # context_embedder + converted_state_dict["context_embedder.weight"] = checkpoint.pop("txt_in.weight") + converted_state_dict["context_embedder.bias"] = checkpoint.pop("txt_in.bias") + + # double transformer blocks + for i in range(num_layers): + block_prefix = f"transformer_blocks.{i}." + # Q, K, V + sample_q, sample_k, sample_v = torch.chunk(checkpoint.pop(f"double_blocks.{i}.img_attn.qkv.weight"), 3, dim=0) + context_q, context_k, context_v = torch.chunk( + checkpoint.pop(f"double_blocks.{i}.txt_attn.qkv.weight"), 3, dim=0 + ) + sample_q_bias, sample_k_bias, sample_v_bias = torch.chunk( + checkpoint.pop(f"double_blocks.{i}.img_attn.qkv.bias"), 3, dim=0 + ) + context_q_bias, context_k_bias, context_v_bias = torch.chunk( + checkpoint.pop(f"double_blocks.{i}.txt_attn.qkv.bias"), 3, dim=0 + ) + converted_state_dict[f"{block_prefix}attn.to_q.weight"] = torch.cat([sample_q]) + converted_state_dict[f"{block_prefix}attn.to_q.bias"] = torch.cat([sample_q_bias]) + converted_state_dict[f"{block_prefix}attn.to_k.weight"] = torch.cat([sample_k]) + converted_state_dict[f"{block_prefix}attn.to_k.bias"] = torch.cat([sample_k_bias]) + converted_state_dict[f"{block_prefix}attn.to_v.weight"] = torch.cat([sample_v]) + converted_state_dict[f"{block_prefix}attn.to_v.bias"] = torch.cat([sample_v_bias]) + converted_state_dict[f"{block_prefix}attn.add_q_proj.weight"] = torch.cat([context_q]) + converted_state_dict[f"{block_prefix}attn.add_q_proj.bias"] = torch.cat([context_q_bias]) + converted_state_dict[f"{block_prefix}attn.add_k_proj.weight"] = torch.cat([context_k]) + converted_state_dict[f"{block_prefix}attn.add_k_proj.bias"] = torch.cat([context_k_bias]) + converted_state_dict[f"{block_prefix}attn.add_v_proj.weight"] = torch.cat([context_v]) + converted_state_dict[f"{block_prefix}attn.add_v_proj.bias"] = torch.cat([context_v_bias]) + # qk_norm + converted_state_dict[f"{block_prefix}attn.norm_q.weight"] = checkpoint.pop( + f"double_blocks.{i}.img_attn.norm.query_norm.scale" + ) + converted_state_dict[f"{block_prefix}attn.norm_k.weight"] = checkpoint.pop( + f"double_blocks.{i}.img_attn.norm.key_norm.scale" + ) + converted_state_dict[f"{block_prefix}attn.norm_added_q.weight"] = checkpoint.pop( + f"double_blocks.{i}.txt_attn.norm.query_norm.scale" + ) + converted_state_dict[f"{block_prefix}attn.norm_added_k.weight"] = checkpoint.pop( + f"double_blocks.{i}.txt_attn.norm.key_norm.scale" + ) + # ff img_mlp + converted_state_dict[f"{block_prefix}ff.net.0.proj.weight"] = checkpoint.pop( + f"double_blocks.{i}.img_mlp.0.weight" + ) + converted_state_dict[f"{block_prefix}ff.net.0.proj.bias"] = checkpoint.pop(f"double_blocks.{i}.img_mlp.0.bias") + converted_state_dict[f"{block_prefix}ff.net.2.weight"] = checkpoint.pop(f"double_blocks.{i}.img_mlp.2.weight") + converted_state_dict[f"{block_prefix}ff.net.2.bias"] = checkpoint.pop(f"double_blocks.{i}.img_mlp.2.bias") + converted_state_dict[f"{block_prefix}ff_context.net.0.proj.weight"] = checkpoint.pop( + f"double_blocks.{i}.txt_mlp.0.weight" + ) + converted_state_dict[f"{block_prefix}ff_context.net.0.proj.bias"] = checkpoint.pop( + f"double_blocks.{i}.txt_mlp.0.bias" + ) + converted_state_dict[f"{block_prefix}ff_context.net.2.weight"] = checkpoint.pop( + f"double_blocks.{i}.txt_mlp.2.weight" + ) + converted_state_dict[f"{block_prefix}ff_context.net.2.bias"] = checkpoint.pop( + f"double_blocks.{i}.txt_mlp.2.bias" + ) + # output projections. + converted_state_dict[f"{block_prefix}attn.to_out.0.weight"] = checkpoint.pop( + f"double_blocks.{i}.img_attn.proj.weight" + ) + converted_state_dict[f"{block_prefix}attn.to_out.0.bias"] = checkpoint.pop( + f"double_blocks.{i}.img_attn.proj.bias" + ) + converted_state_dict[f"{block_prefix}attn.to_add_out.weight"] = checkpoint.pop( + f"double_blocks.{i}.txt_attn.proj.weight" + ) + converted_state_dict[f"{block_prefix}attn.to_add_out.bias"] = checkpoint.pop( + f"double_blocks.{i}.txt_attn.proj.bias" + ) + + # single transformer blocks + for i in range(num_single_layers): + block_prefix = f"single_transformer_blocks.{i}." + # Q, K, V, mlp + mlp_hidden_dim = int(inner_dim * mlp_ratio) + split_size = (inner_dim, inner_dim, inner_dim, mlp_hidden_dim) + q, k, v, mlp = torch.split(checkpoint.pop(f"single_blocks.{i}.linear1.weight"), split_size, dim=0) + q_bias, k_bias, v_bias, mlp_bias = torch.split( + checkpoint.pop(f"single_blocks.{i}.linear1.bias"), split_size, dim=0 + ) + converted_state_dict[f"{block_prefix}attn.to_q.weight"] = torch.cat([q]) + converted_state_dict[f"{block_prefix}attn.to_q.bias"] = torch.cat([q_bias]) + converted_state_dict[f"{block_prefix}attn.to_k.weight"] = torch.cat([k]) + converted_state_dict[f"{block_prefix}attn.to_k.bias"] = torch.cat([k_bias]) + converted_state_dict[f"{block_prefix}attn.to_v.weight"] = torch.cat([v]) + converted_state_dict[f"{block_prefix}attn.to_v.bias"] = torch.cat([v_bias]) + converted_state_dict[f"{block_prefix}proj_mlp.weight"] = torch.cat([mlp]) + converted_state_dict[f"{block_prefix}proj_mlp.bias"] = torch.cat([mlp_bias]) + # qk norm + converted_state_dict[f"{block_prefix}attn.norm_q.weight"] = checkpoint.pop( + f"single_blocks.{i}.norm.query_norm.scale" + ) + converted_state_dict[f"{block_prefix}attn.norm_k.weight"] = checkpoint.pop( + f"single_blocks.{i}.norm.key_norm.scale" + ) + # output projections. + converted_state_dict[f"{block_prefix}proj_out.weight"] = checkpoint.pop(f"single_blocks.{i}.linear2.weight") + converted_state_dict[f"{block_prefix}proj_out.bias"] = checkpoint.pop(f"single_blocks.{i}.linear2.bias") + + # nerf + + converted_state_dict["nerf.nerf_embedder.embedder.0.bias"] = checkpoint.pop("nerf_image_embedder.embedder.0.bias") + converted_state_dict["nerf.nerf_embedder.embedder.0.weight"] = checkpoint.pop( + "nerf_image_embedder.embedder.0.weight" + ) + converted_state_dict["nerf.final_layer.conv.bias"] = checkpoint.pop("nerf_final_layer_conv.conv.bias") + converted_state_dict["nerf.final_layer.conv.weight"] = checkpoint.pop("nerf_final_layer_conv.conv.weight") + converted_state_dict["nerf.final_layer.norm.weight"] = checkpoint.pop("nerf_final_layer_conv.norm.scale") + + for i in range(num_nerf_layers): + block_prefix = f"nerf.blocks.{i}." + converted_state_dict[f"{block_prefix}norm.weight"] = checkpoint.pop(f"nerf_blocks.{i}.norm.scale") + converted_state_dict[f"{block_prefix}param_generator.bias"] = checkpoint.pop( + f"nerf_blocks.{i}.param_generator.bias" + ) + converted_state_dict[f"{block_prefix}param_generator.weight"] = checkpoint.pop( + f"nerf_blocks.{i}.param_generator.weight" + ) + + # patch + + converted_state_dict["x_embedder_patch.bias"] = checkpoint.pop("img_in_patch.bias") + converted_state_dict["x_embedder_patch.weight"] = checkpoint.pop("img_in_patch.weight") + + return converted_state_dict + + def convert_cosmos_transformer_checkpoint_to_diffusers(checkpoint, **kwargs): converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())} @@ -3852,6 +4135,7 @@ def convert_z_image_transformer_checkpoint_to_diffusers(checkpoint, **kwargs): ".attention.k_norm.weight": ".attention.norm_k.weight", ".attention.q_norm.weight": ".attention.norm_q.weight", ".attention.out.weight": ".attention.to_out.0.weight", + "model.diffusion_model.": "", } def convert_z_image_fused_attention(key: str, state_dict: dict[str, object]) -> None: @@ -3886,6 +4170,9 @@ def update_state_dict(state_dict: dict[str, object], old_key: str, new_key: str) update_state_dict(converted_state_dict, key, new_key) + if "norm_final.weight" in converted_state_dict.keys(): + _ = converted_state_dict.pop("norm_final.weight") + # Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in # special_keys_remap for key in list(converted_state_dict.keys()): @@ -3909,3 +4196,161 @@ def convert_z_image_controlnet_checkpoint_to_diffusers(checkpoint, config, **kwa return converted_state_dict else: raise ValueError("Unknown Z-Image Turbo ControlNet type.") + + +def convert_ltx2_transformer_to_diffusers(checkpoint, **kwargs): + LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT = { + # Transformer prefix + "model.diffusion_model.": "", + # Input Patchify Projections + "patchify_proj": "proj_in", + "audio_patchify_proj": "audio_proj_in", + # Modulation Parameters + # Handle adaln_single --> time_embed, audioln_single --> audio_time_embed separately as the original keys are + # substrings of the other modulation parameters below + "av_ca_video_scale_shift_adaln_single": "av_cross_attn_video_scale_shift", + "av_ca_a2v_gate_adaln_single": "av_cross_attn_video_a2v_gate", + "av_ca_audio_scale_shift_adaln_single": "av_cross_attn_audio_scale_shift", + "av_ca_v2a_gate_adaln_single": "av_cross_attn_audio_v2a_gate", + # Transformer Blocks + # Per-Block Cross Attention Modulation Parameters + "scale_shift_table_a2v_ca_video": "video_a2v_cross_attn_scale_shift_table", + "scale_shift_table_a2v_ca_audio": "audio_a2v_cross_attn_scale_shift_table", + # Attention QK Norms + "q_norm": "norm_q", + "k_norm": "norm_k", + } + + def update_state_dict_inplace(state_dict, old_key: str, new_key: str) -> None: + state_dict[new_key] = state_dict.pop(old_key) + + def remove_keys_inplace(key: str, state_dict) -> None: + state_dict.pop(key) + + def convert_ltx2_transformer_adaln_single(key: str, state_dict) -> None: + # Skip if not a weight, bias + if ".weight" not in key and ".bias" not in key: + return + + if key.startswith("adaln_single."): + new_key = key.replace("adaln_single.", "time_embed.") + param = state_dict.pop(key) + state_dict[new_key] = param + + if key.startswith("audio_adaln_single."): + new_key = key.replace("audio_adaln_single.", "audio_time_embed.") + param = state_dict.pop(key) + state_dict[new_key] = param + + return + + LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP = { + "video_embeddings_connector": remove_keys_inplace, + "audio_embeddings_connector": remove_keys_inplace, + "adaln_single": convert_ltx2_transformer_adaln_single, + } + + converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())} + + # Handle official code --> diffusers key remapping via the remap dict + for key in list(converted_state_dict.keys()): + new_key = key[:] + for replace_key, rename_key in LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT.items(): + new_key = new_key.replace(replace_key, rename_key) + + update_state_dict_inplace(converted_state_dict, key, new_key) + + # Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in + # special_keys_remap + for key in list(converted_state_dict.keys()): + for special_key, handler_fn_inplace in LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP.items(): + if special_key not in key: + continue + handler_fn_inplace(key, converted_state_dict) + + return converted_state_dict + + +def convert_ltx2_vae_to_diffusers(checkpoint, **kwargs): + LTX_2_0_VIDEO_VAE_RENAME_DICT = { + # Video VAE prefix + "vae.": "", + # Encoder + "down_blocks.0": "down_blocks.0", + "down_blocks.1": "down_blocks.0.downsamplers.0", + "down_blocks.2": "down_blocks.1", + "down_blocks.3": "down_blocks.1.downsamplers.0", + "down_blocks.4": "down_blocks.2", + "down_blocks.5": "down_blocks.2.downsamplers.0", + "down_blocks.6": "down_blocks.3", + "down_blocks.7": "down_blocks.3.downsamplers.0", + "down_blocks.8": "mid_block", + # Decoder + "up_blocks.0": "mid_block", + "up_blocks.1": "up_blocks.0.upsamplers.0", + "up_blocks.2": "up_blocks.0", + "up_blocks.3": "up_blocks.1.upsamplers.0", + "up_blocks.4": "up_blocks.1", + "up_blocks.5": "up_blocks.2.upsamplers.0", + "up_blocks.6": "up_blocks.2", + # Common + # For all 3D ResNets + "res_blocks": "resnets", + "per_channel_statistics.mean-of-means": "latents_mean", + "per_channel_statistics.std-of-means": "latents_std", + } + + def update_state_dict_inplace(state_dict, old_key: str, new_key: str) -> None: + state_dict[new_key] = state_dict.pop(old_key) + + def remove_keys_inplace(key: str, state_dict) -> None: + state_dict.pop(key) + + LTX_2_0_VAE_SPECIAL_KEYS_REMAP = { + "per_channel_statistics.channel": remove_keys_inplace, + "per_channel_statistics.mean-of-stds": remove_keys_inplace, + } + + converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())} + + # Handle official code --> diffusers key remapping via the remap dict + for key in list(converted_state_dict.keys()): + new_key = key[:] + for replace_key, rename_key in LTX_2_0_VIDEO_VAE_RENAME_DICT.items(): + new_key = new_key.replace(replace_key, rename_key) + + update_state_dict_inplace(converted_state_dict, key, new_key) + + # Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in + # special_keys_remap + for key in list(converted_state_dict.keys()): + for special_key, handler_fn_inplace in LTX_2_0_VAE_SPECIAL_KEYS_REMAP.items(): + if special_key not in key: + continue + handler_fn_inplace(key, converted_state_dict) + + return converted_state_dict + + +def convert_ltx2_audio_vae_to_diffusers(checkpoint, **kwargs): + LTX_2_0_AUDIO_VAE_RENAME_DICT = { + # Audio VAE prefix + "audio_vae.": "", + "per_channel_statistics.mean-of-means": "latents_mean", + "per_channel_statistics.std-of-means": "latents_std", + } + + def update_state_dict_inplace(state_dict, old_key: str, new_key: str) -> None: + state_dict[new_key] = state_dict.pop(old_key) + + converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())} + + # Handle official code --> diffusers key remapping via the remap dict + for key in list(converted_state_dict.keys()): + new_key = key[:] + for replace_key, rename_key in LTX_2_0_AUDIO_VAE_RENAME_DICT.items(): + new_key = new_key.replace(replace_key, rename_key) + + update_state_dict_inplace(converted_state_dict, key, new_key) + + return converted_state_dict diff --git a/src/diffusers/loaders/textual_inversion.py b/src/diffusers/loaders/textual_inversion.py index 63fc97ed431f..711de5f81343 100644 --- a/src/diffusers/loaders/textual_inversion.py +++ b/src/diffusers/loaders/textual_inversion.py @@ -11,15 +11,23 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Dict, List, Optional, Union +from __future__ import annotations + +import json import safetensors import torch from huggingface_hub.utils import validate_hf_hub_args +from tokenizers import Tokenizer as TokenizerFast from torch import nn from ..models.modeling_utils import load_state_dict -from ..utils import _get_model_file, is_accelerate_available, is_transformers_available, logging +from ..utils import ( + _get_model_file, + is_accelerate_available, + is_transformers_available, + logging, +) if is_transformers_available(): @@ -112,7 +120,7 @@ class TextualInversionLoaderMixin: Load Textual Inversion tokens and embeddings to the tokenizer and text encoder. """ - def maybe_convert_prompt(self, prompt: Union[str, List[str]], tokenizer: "PreTrainedTokenizer"): # noqa: F821 + def maybe_convert_prompt(self, prompt: str | list[str], tokenizer: "PreTrainedTokenizer"): # noqa: F821 r""" Processes prompts that include a special token corresponding to a multi-vector textual inversion embedding to be replaced with multiple special tokens each corresponding to one of the vectors. If the prompt has no textual @@ -127,14 +135,14 @@ def maybe_convert_prompt(self, prompt: Union[str, List[str]], tokenizer: "PreTra Returns: `str` or list of `str`: The converted prompt """ - if not isinstance(prompt, List): + if not isinstance(prompt, list): prompts = [prompt] else: prompts = prompt prompts = [self._maybe_convert_prompt(p, tokenizer) for p in prompts] - if not isinstance(prompt, List): + if not isinstance(prompt, list): return prompts[0] return prompts @@ -263,10 +271,10 @@ def _extend_tokens_and_embeddings(tokens, embeddings, tokenizer): @validate_hf_hub_args def load_textual_inversion( self, - pretrained_model_name_or_path: Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]], - token: Optional[Union[str, List[str]]] = None, - tokenizer: Optional["PreTrainedTokenizer"] = None, # noqa: F821 - text_encoder: Optional["PreTrainedModel"] = None, # noqa: F821 + pretrained_model_name_or_path: str | list[str] | dict[str, torch.Tensor] | list[dict[str, torch.Tensor]], + token: str | list[str] | None = None, + tokenizer: "PreTrainedTokenizer" | None = None, # noqa: F821 + text_encoder: "PreTrainedModel" | None = None, # noqa: F821 **kwargs, ): r""" @@ -274,7 +282,7 @@ def load_textual_inversion( Automatic1111 formats are supported). Parameters: - pretrained_model_name_or_path (`str` or `os.PathLike` or `List[str or os.PathLike]` or `Dict` or `List[Dict]`): + pretrained_model_name_or_path (`str` or `os.PathLike` or `list[str or os.PathLike]` or `Dict` or `list[Dict]`): Can be either one of the following or a list of them: - A string, the *model id* (for example `sd-concepts-library/low-poly-hd-logos-icons`) of a @@ -285,7 +293,7 @@ def load_textual_inversion( - A [torch state dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - token (`str` or `List[str]`, *optional*): + token (`str` or `list[str]`, *optional*): Override the token to use for the textual inversion weights. If `pretrained_model_name_or_path` is a list, then `token` must also be a list of equal length. text_encoder ([`~transformers.CLIPTextModel`], *optional*): @@ -299,14 +307,14 @@ def load_textual_inversion( - The saved textual inversion file is in 🤗 Diffusers format, but was saved under a specific weight name such as `text_inv.bin`. - The saved textual inversion file is in the Automatic1111 format. - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. local_files_only (`bool`, *optional*, defaults to `False`): @@ -458,9 +466,9 @@ def load_textual_inversion( def unload_textual_inversion( self, - tokens: Optional[Union[str, List[str]]] = None, - tokenizer: Optional["PreTrainedTokenizer"] = None, - text_encoder: Optional["PreTrainedModel"] = None, + tokens: str | list[str] | None = None, + tokenizer: "PreTrainedTokenizer" | None = None, + text_encoder: "PreTrainedModel" | None = None, ): r""" Unload Textual Inversion embeddings from the text encoder of [`StableDiffusionPipeline`] @@ -547,23 +555,39 @@ def unload_textual_inversion( else: last_special_token_id = added_token_id - # Delete from tokenizer - for token_id, token_to_remove in zip(token_ids, tokens): - del tokenizer._added_tokens_decoder[token_id] - del tokenizer._added_tokens_encoder[token_to_remove] - - # Make all token ids sequential in tokenizer - key_id = 1 - for token_id in tokenizer.added_tokens_decoder: - if token_id > last_special_token_id and token_id > last_special_token_id + key_id: - token = tokenizer._added_tokens_decoder[token_id] - tokenizer._added_tokens_decoder[last_special_token_id + key_id] = token + # Fast tokenizers (v5+) + if hasattr(tokenizer, "_tokenizer"): + # Fast tokenizers: serialize, filter tokens, reload + tokenizer_json = json.loads(tokenizer._tokenizer.to_str()) + new_id = last_special_token_id + 1 + filtered = [] + for tok in tokenizer_json.get("added_tokens", []): + if tok.get("content") in set(tokens): + continue + if not tok.get("special", False): + tok["id"] = new_id + new_id += 1 + filtered.append(tok) + tokenizer_json["added_tokens"] = filtered + tokenizer._tokenizer = TokenizerFast.from_str(json.dumps(tokenizer_json)) + else: + # Slow tokenizers + for token_id, token_to_remove in zip(token_ids, tokens): del tokenizer._added_tokens_decoder[token_id] - tokenizer._added_tokens_encoder[token.content] = last_special_token_id + key_id - key_id += 1 - tokenizer._update_trie() - # set correct total vocab size after removing tokens - tokenizer._update_total_vocab_size() + del tokenizer._added_tokens_encoder[token_to_remove] + + key_id = 1 + for token_id in list(tokenizer.added_tokens_decoder.keys()): + if token_id > last_special_token_id and token_id > last_special_token_id + key_id: + token = tokenizer._added_tokens_decoder[token_id] + tokenizer._added_tokens_decoder[last_special_token_id + key_id] = token + del tokenizer._added_tokens_decoder[token_id] + tokenizer._added_tokens_encoder[token.content] = last_special_token_id + key_id + key_id += 1 + if hasattr(tokenizer, "_update_trie"): + tokenizer._update_trie() + if hasattr(tokenizer, "_update_total_vocab_size"): + tokenizer._update_total_vocab_size() # Delete from text encoder text_embedding_dim = text_encoder.get_input_embeddings().embedding_dim diff --git a/src/diffusers/loaders/transformer_sd3.py b/src/diffusers/loaders/transformer_sd3.py index e3728082efdd..d1c3fff14e60 100644 --- a/src/diffusers/loaders/transformer_sd3.py +++ b/src/diffusers/loaders/transformer_sd3.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. from contextlib import nullcontext -from typing import Dict from ..models.attention_processor import SD3IPAdapterJointAttnProcessor2_0 from ..models.embeddings import IPAdapterTimeImageProjection @@ -29,8 +28,8 @@ class SD3Transformer2DLoadersMixin: """Load IP-Adapters and LoRA layers into a `[SD3Transformer2DModel]`.""" def _convert_ip_adapter_attn_to_diffusers( - self, state_dict: Dict, low_cpu_mem_usage: bool = _LOW_CPU_MEM_USAGE_DEFAULT - ) -> Dict: + self, state_dict: dict, low_cpu_mem_usage: bool = _LOW_CPU_MEM_USAGE_DEFAULT + ) -> dict: if low_cpu_mem_usage: if is_accelerate_available(): from accelerate import init_empty_weights @@ -87,7 +86,7 @@ def _convert_ip_adapter_attn_to_diffusers( return attn_procs def _convert_ip_adapter_image_proj_to_diffusers( - self, state_dict: Dict, low_cpu_mem_usage: bool = _LOW_CPU_MEM_USAGE_DEFAULT + self, state_dict: dict, low_cpu_mem_usage: bool = _LOW_CPU_MEM_USAGE_DEFAULT ) -> IPAdapterTimeImageProjection: if low_cpu_mem_usage: if is_accelerate_available(): @@ -155,7 +154,7 @@ def _convert_ip_adapter_image_proj_to_diffusers( return image_proj - def _load_ip_adapter_weights(self, state_dict: Dict, low_cpu_mem_usage: bool = _LOW_CPU_MEM_USAGE_DEFAULT) -> None: + def _load_ip_adapter_weights(self, state_dict: dict, low_cpu_mem_usage: bool = _LOW_CPU_MEM_USAGE_DEFAULT) -> None: """Sets IP-Adapter attention processors, image projection, and loads state_dict. Args: diff --git a/src/diffusers/loaders/unet.py b/src/diffusers/loaders/unet.py index c5e56af156fc..9dab3bc667ea 100644 --- a/src/diffusers/loaders/unet.py +++ b/src/diffusers/loaders/unet.py @@ -15,7 +15,7 @@ from collections import defaultdict from contextlib import nullcontext from pathlib import Path -from typing import Callable, Dict, Union +from typing import Callable import safetensors import torch @@ -66,7 +66,7 @@ class UNet2DConditionLoadersMixin: unet_name = UNET_NAME @validate_hf_hub_args - def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], **kwargs): + def load_attn_procs(self, pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs): r""" Load pretrained attention processor layers into [`UNet2DConditionModel`]. Attention processor layers have to be defined in @@ -85,14 +85,14 @@ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict - A [torch state dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. local_files_only (`bool`, *optional*, defaults to `False`): @@ -106,7 +106,7 @@ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict allowed by Git. subfolder (`str`, *optional*, defaults to `""`): The subfolder location of a model file within a larger model repository on the Hub or locally. - network_alphas (`Dict[str, float]`): + network_alphas (`dict[str, float]`): The value of the network alpha used for stable learning and preventing underflow. This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning). @@ -412,7 +412,7 @@ def _optionally_disable_offloading(cls, _pipeline): def save_attn_procs( self, - save_directory: Union[str, os.PathLike], + save_directory: str | os.PathLike, is_main_process: bool = True, weight_name: str = None, save_function: Callable = None, diff --git a/src/diffusers/loaders/unet_loader_utils.py b/src/diffusers/loaders/unet_loader_utils.py index d5b0e83cbd9e..d6aff55d6d22 100644 --- a/src/diffusers/loaders/unet_loader_utils.py +++ b/src/diffusers/loaders/unet_loader_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import copy -from typing import TYPE_CHECKING, Dict, List, Union +from typing import TYPE_CHECKING from torch import nn @@ -40,9 +40,7 @@ def _translate_into_actual_layer_name(name): return ".".join((updown, block, attn)) -def _maybe_expand_lora_scales( - unet: "UNet2DConditionModel", weight_scales: List[Union[float, Dict]], default_scale=1.0 -): +def _maybe_expand_lora_scales(unet: "UNet2DConditionModel", weight_scales: list[float | dict], default_scale=1.0): blocks_with_transformer = { "down": [i for i, block in enumerate(unet.down_blocks) if hasattr(block, "attentions")], "up": [i for i, block in enumerate(unet.up_blocks) if hasattr(block, "attentions")], @@ -64,9 +62,9 @@ def _maybe_expand_lora_scales( def _maybe_expand_lora_scales_for_one_adapter( - scales: Union[float, Dict], - blocks_with_transformer: Dict[str, int], - transformer_per_block: Dict[str, int], + scales: float | dict, + blocks_with_transformer: dict[str, int], + transformer_per_block: dict[str, int], model: nn.Module, default_scale: float = 1.0, ): @@ -74,11 +72,11 @@ def _maybe_expand_lora_scales_for_one_adapter( Expands the inputs into a more granular dictionary. See the example below for more details. Parameters: - scales (`Union[float, Dict]`): + scales (`float | Dict`): Scales dict to expand. - blocks_with_transformer (`Dict[str, int]`): + blocks_with_transformer (`dict[str, int]`): Dict with keys 'up' and 'down', showing which blocks have transformer layers - transformer_per_block (`Dict[str, int]`): + transformer_per_block (`dict[str, int]`): Dict with keys 'up' and 'down', showing how many transformer layers each block has E.g. turns diff --git a/src/diffusers/loaders/utils.py b/src/diffusers/loaders/utils.py index 2d39e7bfb7d2..72d6453a8bbf 100644 --- a/src/diffusers/loaders/utils.py +++ b/src/diffusers/loaders/utils.py @@ -12,13 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Dict import torch class AttnProcsLayers(torch.nn.Module): - def __init__(self, state_dict: Dict[str, torch.Tensor]): + def __init__(self, state_dict: dict[str, torch.Tensor]): super().__init__() self.layers = torch.nn.ModuleList(state_dict.values()) self.mapping = dict(enumerate(state_dict.keys())) diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index c4664f00cad2..244f2963104b 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -41,17 +41,21 @@ _import_structure["autoencoders.autoencoder_kl_hunyuanimage_refiner"] = ["AutoencoderKLHunyuanImageRefiner"] _import_structure["autoencoders.autoencoder_kl_hunyuanvideo15"] = ["AutoencoderKLHunyuanVideo15"] _import_structure["autoencoders.autoencoder_kl_ltx"] = ["AutoencoderKLLTXVideo"] + _import_structure["autoencoders.autoencoder_kl_ltx2"] = ["AutoencoderKLLTX2Video"] + _import_structure["autoencoders.autoencoder_kl_ltx2_audio"] = ["AutoencoderKLLTX2Audio"] _import_structure["autoencoders.autoencoder_kl_magvit"] = ["AutoencoderKLMagvit"] _import_structure["autoencoders.autoencoder_kl_mochi"] = ["AutoencoderKLMochi"] _import_structure["autoencoders.autoencoder_kl_qwenimage"] = ["AutoencoderKLQwenImage"] _import_structure["autoencoders.autoencoder_kl_temporal_decoder"] = ["AutoencoderKLTemporalDecoder"] _import_structure["autoencoders.autoencoder_kl_wan"] = ["AutoencoderKLWan"] _import_structure["autoencoders.autoencoder_oobleck"] = ["AutoencoderOobleck"] + _import_structure["autoencoders.autoencoder_rae"] = ["AutoencoderRAE"] _import_structure["autoencoders.autoencoder_tiny"] = ["AutoencoderTiny"] _import_structure["autoencoders.consistency_decoder_vae"] = ["ConsistencyDecoderVAE"] _import_structure["autoencoders.vq_model"] = ["VQModel"] _import_structure["cache_utils"] = ["CacheMixin"] _import_structure["controlnets.controlnet"] = ["ControlNetModel"] + _import_structure["controlnets.controlnet_cosmos"] = ["CosmosControlNetModel"] _import_structure["controlnets.controlnet_flux"] = ["FluxControlNetModel", "FluxMultiControlNetModel"] _import_structure["controlnets.controlnet_hunyuan"] = [ "HunyuanDiT2DControlNetModel", @@ -88,7 +92,10 @@ _import_structure["transformers.transformer_allegro"] = ["AllegroTransformer3DModel"] _import_structure["transformers.transformer_bria"] = ["BriaTransformer2DModel"] _import_structure["transformers.transformer_bria_fibo"] = ["BriaFiboTransformer2DModel"] - _import_structure["transformers.transformer_chroma"] = ["ChromaTransformer2DModel"] + _import_structure["transformers.transformer_chroma"] = [ + "ChromaTransformer2DModel", + "ChromaRadianceTransformer2DModel", + ] _import_structure["transformers.transformer_chronoedit"] = ["ChronoEditTransformer3DModel"] _import_structure["transformers.transformer_cogview3plus"] = ["CogView3PlusTransformer2DModel"] _import_structure["transformers.transformer_cogview4"] = ["CogView4Transformer2DModel"] @@ -96,6 +103,8 @@ _import_structure["transformers.transformer_easyanimate"] = ["EasyAnimateTransformer3DModel"] _import_structure["transformers.transformer_flux"] = ["FluxTransformer2DModel"] _import_structure["transformers.transformer_flux2"] = ["Flux2Transformer2DModel"] + _import_structure["transformers.transformer_glm_image"] = ["GlmImageTransformer2DModel"] + _import_structure["transformers.transformer_helios"] = ["HeliosTransformer3DModel"] _import_structure["transformers.transformer_hidream_image"] = ["HiDreamImageTransformer2DModel"] _import_structure["transformers.transformer_hunyuan_video"] = ["HunyuanVideoTransformer3DModel"] _import_structure["transformers.transformer_hunyuan_video15"] = ["HunyuanVideo15Transformer3DModel"] @@ -104,6 +113,7 @@ _import_structure["transformers.transformer_kandinsky"] = ["Kandinsky5Transformer3DModel"] _import_structure["transformers.transformer_longcat_image"] = ["LongCatImageTransformer2DModel"] _import_structure["transformers.transformer_ltx"] = ["LTXVideoTransformer3DModel"] + _import_structure["transformers.transformer_ltx2"] = ["LTX2VideoTransformer3DModel"] _import_structure["transformers.transformer_lumina2"] = ["Lumina2Transformer2DModel"] _import_structure["transformers.transformer_mochi"] = ["MochiTransformer3DModel"] _import_structure["transformers.transformer_omnigen"] = ["OmniGenTransformer2DModel"] @@ -153,6 +163,8 @@ AutoencoderKLHunyuanImageRefiner, AutoencoderKLHunyuanVideo, AutoencoderKLHunyuanVideo15, + AutoencoderKLLTX2Audio, + AutoencoderKLLTX2Video, AutoencoderKLLTXVideo, AutoencoderKLMagvit, AutoencoderKLMochi, @@ -160,6 +172,7 @@ AutoencoderKLTemporalDecoder, AutoencoderKLWan, AutoencoderOobleck, + AutoencoderRAE, AutoencoderTiny, ConsistencyDecoderVAE, VQModel, @@ -169,6 +182,7 @@ ControlNetModel, ControlNetUnionModel, ControlNetXSAdapter, + CosmosControlNetModel, FluxControlNetModel, FluxMultiControlNetModel, HunyuanDiT2DControlNetModel, @@ -191,6 +205,7 @@ AuraFlowTransformer2DModel, BriaFiboTransformer2DModel, BriaTransformer2DModel, + ChromaRadianceTransformer2DModel, ChromaTransformer2DModel, ChronoEditTransformer3DModel, CogVideoXTransformer3DModel, @@ -203,6 +218,8 @@ EasyAnimateTransformer3DModel, Flux2Transformer2DModel, FluxTransformer2DModel, + GlmImageTransformer2DModel, + HeliosTransformer3DModel, HiDreamImageTransformer2DModel, HunyuanDiT2DModel, HunyuanImageTransformer2DModel, @@ -212,6 +229,7 @@ Kandinsky5Transformer3DModel, LatteTransformer3DModel, LongCatImageTransformer2DModel, + LTX2VideoTransformer3DModel, LTXVideoTransformer3DModel, Lumina2Transformer2DModel, LuminaNextDiT2DModel, diff --git a/src/diffusers/models/_modeling_parallel.py b/src/diffusers/models/_modeling_parallel.py index 2a4eb520c796..ed966dc8fe98 100644 --- a/src/diffusers/models/_modeling_parallel.py +++ b/src/diffusers/models/_modeling_parallel.py @@ -16,9 +16,10 @@ # limitations under the License. from dataclasses import dataclass -from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union +from typing import TYPE_CHECKING, Literal import torch +import torch.distributed as dist from ..utils import get_logger @@ -62,11 +63,14 @@ class ContextParallelConfig: """ - ring_degree: Optional[int] = None - ulysses_degree: Optional[int] = None + ring_degree: int | None = None + ulysses_degree: int | None = None convert_to_fp32: bool = True # TODO: support alltoall rotate_method: Literal["allgather", "alltoall"] = "allgather" + # Whether to enable ulysses anything attention to support + # any sequence lengths and any head numbers. + ulysses_anything: bool = False _rank: int = None _world_size: int = None @@ -90,21 +94,22 @@ def __post_init__(self): ) if self.ring_degree < 1 or self.ulysses_degree < 1: raise ValueError("`ring_degree` and `ulysses_degree` must be greater than or equal to 1.") - if self.ring_degree > 1 and self.ulysses_degree > 1: - raise ValueError( - "Unified Ulysses-Ring attention is not yet supported. Please set either `ring_degree` or `ulysses_degree` to 1." - ) if self.rotate_method != "allgather": raise NotImplementedError( f"Only rotate_method='allgather' is supported for now, but got {self.rotate_method}." ) + if self.ulysses_anything: + if self.ulysses_degree == 1: + raise ValueError("ulysses_degree must be greater than 1 for ulysses_anything to be enabled.") + if self.ring_degree > 1: + raise ValueError("ulysses_anything cannot be enabled when ring_degree > 1.") @property - def mesh_shape(self) -> Tuple[int, int]: + def mesh_shape(self) -> tuple[int, int]: return (self.ring_degree, self.ulysses_degree) @property - def mesh_dim_names(self) -> Tuple[str, str]: + def mesh_dim_names(self) -> tuple[str, str]: """Dimension names for the device mesh.""" return ("ring", "ulysses") @@ -136,7 +141,7 @@ class ParallelConfig: Configuration for context parallelism. """ - context_parallel_config: Optional[ContextParallelConfig] = None + context_parallel_config: ContextParallelConfig | None = None _rank: int = None _world_size: int = None @@ -149,7 +154,7 @@ def setup( world_size: int, device: torch.device, *, - mesh: Optional[torch.distributed.device_mesh.DeviceMesh] = None, + mesh: torch.distributed.device_mesh.DeviceMesh | None = None, ): self._rank = rank self._world_size = world_size @@ -177,7 +182,7 @@ class ContextParallelInput: """ split_dim: int - expected_dims: Optional[int] = None + expected_dims: int | None = None split_output: bool = False def __repr__(self): @@ -198,7 +203,7 @@ class ContextParallelOutput: """ gather_dim: int - expected_dims: Optional[int] = None + expected_dims: int | None = None def __repr__(self): return f"ContextParallelOutput(gather_dim={self.gather_dim}, expected_dims={self.expected_dims})" @@ -209,19 +214,17 @@ def __repr__(self): # If the key is a string, it denotes the name of the parameter in the forward function. # If the key is an integer, split_output must be set to True, and it denotes the index of the output # to be split across context parallel region. -ContextParallelInputType = Dict[ - Union[str, int], Union[ContextParallelInput, List[ContextParallelInput], Tuple[ContextParallelInput, ...]] +ContextParallelInputType = dict[ + str | int, ContextParallelInput | list[ContextParallelInput] | tuple[ContextParallelInput, ...] ] # A dictionary where keys denote the output to be gathered across context parallel region, and the # value denotes the gathering configuration. -ContextParallelOutputType = Union[ - ContextParallelOutput, List[ContextParallelOutput], Tuple[ContextParallelOutput, ...] -] +ContextParallelOutputType = ContextParallelOutput | list[ContextParallelOutput] | tuple[ContextParallelOutput, ...] # A dictionary where keys denote the module id, and the value denotes how the inputs/outputs of # the module should be split/gathered across context parallel region. -ContextParallelModelPlan = Dict[str, Union[ContextParallelInputType, ContextParallelOutputType]] +ContextParallelModelPlan = dict[str, ContextParallelInputType | ContextParallelOutputType] # Example of a ContextParallelModelPlan (QwenImageTransformer2DModel): @@ -261,3 +264,39 @@ def __repr__(self): # # ContextParallelOutput: # specifies how to gather the input tensor in the post-forward hook in the layer it is attached to + + +# Below are utility functions for distributed communication in context parallelism. +def gather_size_by_comm(size: int, group: dist.ProcessGroup) -> list[int]: + r"""Gather the local size from all ranks. + size: int, local size return: list[int], list of size from all ranks + """ + # NOTE(Serving/CP Safety): + # Do NOT cache this collective result. + # + # In "Ulysses Anything" mode, `size` (e.g. per-rank local seq_len / S_LOCAL) + # may legitimately differ across ranks. If we cache based on the *local* `size`, + # different ranks can have different cache hit/miss patterns across time. + # + # That can lead to a catastrophic distributed hang: + # - some ranks hit cache and *skip* dist.all_gather() + # - other ranks miss cache and *enter* dist.all_gather() + # This mismatched collective participation will stall the process group and + # eventually trigger NCCL watchdog timeouts (often surfacing later as ALLTOALL + # timeouts in Ulysses attention). + world_size = dist.get_world_size(group=group) + # HACK: Use Gloo backend for all_gather to avoid H2D and D2H overhead + comm_backends = str(dist.get_backend(group=group)) + # NOTE: e.g., dist.init_process_group(backend="cpu:gloo,cuda:nccl") + gather_device = "cpu" if "cpu" in comm_backends else torch.accelerator.current_accelerator() + gathered_sizes = [torch.empty((1,), device=gather_device, dtype=torch.int64) for _ in range(world_size)] + dist.all_gather( + gathered_sizes, + torch.tensor([size], device=gather_device, dtype=torch.int64), + group=group, + ) + + gathered_sizes = [s[0].item() for s in gathered_sizes] + # NOTE: DON'T use tolist here due to graph break - Explanation: + # Backend compiler `inductor` failed with aten._local_scalar_dense.default + return gathered_sizes diff --git a/src/diffusers/models/adapter.py b/src/diffusers/models/adapter.py index e475fe6bee88..f0652c581a3e 100644 --- a/src/diffusers/models/adapter.py +++ b/src/diffusers/models/adapter.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import os -from typing import Callable, List, Optional, Union +from typing import Callable import torch import torch.nn as nn @@ -34,11 +34,11 @@ class MultiAdapter(ModelMixin): or saving. Args: - adapters (`List[T2IAdapter]`, *optional*, defaults to None): + adapters (`list[T2IAdapter]`, *optional*, defaults to None): A list of `T2IAdapter` model instances. """ - def __init__(self, adapters: List["T2IAdapter"]): + def __init__(self, adapters: list["T2IAdapter"]): super(MultiAdapter, self).__init__() self.num_adapter = len(adapters) @@ -73,7 +73,7 @@ def __init__(self, adapters: List["T2IAdapter"]): self.total_downscale_factor = first_adapter_total_downscale_factor self.downscale_factor = first_adapter_downscale_factor - def forward(self, xs: torch.Tensor, adapter_weights: Optional[List[float]] = None) -> List[torch.Tensor]: + def forward(self, xs: torch.Tensor, adapter_weights: list[float] | None = None) -> list[torch.Tensor]: r""" Args: xs (`torch.Tensor`): @@ -81,7 +81,7 @@ def forward(self, xs: torch.Tensor, adapter_weights: Optional[List[float]] = Non models, concatenated along dimension 1(channel dimension). The `channel` dimension should be equal to `num_adapter` * number of channel per image. - adapter_weights (`List[float]`, *optional*, defaults to None): + adapter_weights (`list[float]`, *optional*, defaults to None): A list of floats representing the weights which will be multiplied by each adapter's output before summing them together. If `None`, equal weights will be used for all adapters. """ @@ -104,11 +104,11 @@ def forward(self, xs: torch.Tensor, adapter_weights: Optional[List[float]] = Non def save_pretrained( self, - save_directory: Union[str, os.PathLike], + save_directory: str | os.PathLike, is_main_process: bool = True, save_function: Callable = None, safe_serialization: bool = True, - variant: Optional[str] = None, + variant: str | None = None, ): """ Save a model and its configuration file to a specified directory, allowing it to be re-loaded with the @@ -145,7 +145,7 @@ def save_pretrained( model_path_to_save = model_path_to_save + f"_{idx}" @classmethod - def from_pretrained(cls, pretrained_model_path: Optional[Union[str, os.PathLike]], **kwargs): + def from_pretrained(cls, pretrained_model_path: str | os.PathLike | None, **kwargs): r""" Instantiate a pretrained `MultiAdapter` model from multiple pre-trained adapter models. @@ -165,7 +165,7 @@ def from_pretrained(cls, pretrained_model_path: Optional[Union[str, os.PathLike] Override the default `torch.dtype` and load the model under this dtype. output_loading_info(`bool`, *optional*, defaults to `False`): Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. - device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*): + device_map (`str` or `dict[str, int | str | torch.device]`, *optional*): A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer name, once a given module name is inside, every submodule of it will be sent to the same device. @@ -229,7 +229,7 @@ class T2IAdapter(ModelMixin, ConfigMixin): in_channels (`int`, *optional*, defaults to `3`): The number of channels in the adapter's input (*control image*). Set it to 1 if you're using a gray scale image. - channels (`List[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): + channels (`list[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): The number of channels in each downsample block's output hidden state. The `len(block_out_channels)` determines the number of downsample blocks in the adapter. num_res_blocks (`int`, *optional*, defaults to `2`): @@ -244,7 +244,7 @@ class T2IAdapter(ModelMixin, ConfigMixin): def __init__( self, in_channels: int = 3, - channels: List[int] = [320, 640, 1280, 1280], + channels: list[int] = [320, 640, 1280, 1280], num_res_blocks: int = 2, downscale_factor: int = 8, adapter_type: str = "full_adapter", @@ -263,7 +263,7 @@ def __init__( "'full_adapter_xl' or 'light_adapter'." ) - def forward(self, x: torch.Tensor) -> List[torch.Tensor]: + def forward(self, x: torch.Tensor) -> list[torch.Tensor]: r""" This function processes the input tensor `x` through the adapter model and returns a list of feature tensors, each representing information extracted at a different scale from the input. The length of the list is @@ -295,7 +295,7 @@ class FullAdapter(nn.Module): def __init__( self, in_channels: int = 3, - channels: List[int] = [320, 640, 1280, 1280], + channels: list[int] = [320, 640, 1280, 1280], num_res_blocks: int = 2, downscale_factor: int = 8, ): @@ -318,7 +318,7 @@ def __init__( self.total_downscale_factor = downscale_factor * 2 ** (len(channels) - 1) - def forward(self, x: torch.Tensor) -> List[torch.Tensor]: + def forward(self, x: torch.Tensor) -> list[torch.Tensor]: r""" This method processes the input tensor `x` through the FullAdapter model and performs operations including pixel unshuffling, convolution, and a stack of AdapterBlocks. It returns a list of feature tensors, each @@ -345,7 +345,7 @@ class FullAdapterXL(nn.Module): def __init__( self, in_channels: int = 3, - channels: List[int] = [320, 640, 1280, 1280], + channels: list[int] = [320, 640, 1280, 1280], num_res_blocks: int = 2, downscale_factor: int = 16, ): @@ -370,7 +370,7 @@ def __init__( # XL has only one downsampling AdapterBlock. self.total_downscale_factor = downscale_factor * 2 - def forward(self, x: torch.Tensor) -> List[torch.Tensor]: + def forward(self, x: torch.Tensor) -> list[torch.Tensor]: r""" This method takes the tensor x as input and processes it through FullAdapterXL model. It consists of operations including unshuffling pixels, applying convolution layer and appending each block into list of feature tensors. @@ -473,7 +473,7 @@ class LightAdapter(nn.Module): def __init__( self, in_channels: int = 3, - channels: List[int] = [320, 640, 1280], + channels: list[int] = [320, 640, 1280], num_res_blocks: int = 4, downscale_factor: int = 8, ): @@ -496,7 +496,7 @@ def __init__( self.total_downscale_factor = downscale_factor * (2 ** len(channels)) - def forward(self, x: torch.Tensor) -> List[torch.Tensor]: + def forward(self, x: torch.Tensor) -> list[torch.Tensor]: r""" This method takes the input tensor x and performs downscaling and appends it in list of feature tensors. Each feature tensor corresponds to a different level of processing within the LightAdapter. diff --git a/src/diffusers/models/attention.py b/src/diffusers/models/attention.py index 8b583d1a1cce..36d0893734c7 100644 --- a/src/diffusers/models/attention.py +++ b/src/diffusers/models/attention.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch import torch.nn as nn @@ -38,7 +38,7 @@ class AttentionMixin: @property - def attn_processors(self) -> Dict[str, AttentionProcessor]: + def attn_processors(self) -> dict[str, AttentionProcessor]: r""" Returns: `dict` of attention processors: A dictionary containing all attention processors used in the model with @@ -47,7 +47,7 @@ def attn_processors(self) -> Dict[str, AttentionProcessor]: # set recursively processors = {} - def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]): + def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: dict[str, AttentionProcessor]): if hasattr(module, "get_processor"): processors[f"{name}.processor"] = module.get_processor() @@ -61,7 +61,7 @@ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: return processors - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): + def set_attn_processor(self, processor: AttentionProcessor | dict[str, AttentionProcessor]): r""" Sets the attention processor to use to compute attention. @@ -185,7 +185,7 @@ def set_use_npu_flash_attention(self, use_npu_flash_attention: bool) -> None: def set_use_xla_flash_attention( self, use_xla_flash_attention: bool, - partition_spec: Optional[Tuple[Optional[str], ...]] = None, + partition_spec: tuple[str | None, ...] | None = None, is_flux=False, ) -> None: """ @@ -194,7 +194,7 @@ def set_use_xla_flash_attention( Args: use_xla_flash_attention (`bool`): Whether to use pallas flash attention kernel from `torch_xla` or not. - partition_spec (`Tuple[]`, *optional*): + partition_spec (`tuple[]`, *optional*): Specify the partition specification if using SPMD. Otherwise None. is_flux (`bool`, *optional*, defaults to `False`): Whether the model is a Flux model. @@ -206,7 +206,7 @@ def set_use_xla_flash_attention( self.set_attention_backend("_native_xla") def set_use_memory_efficient_attention_xformers( - self, use_memory_efficient_attention_xformers: bool, attention_op: Optional[Callable] = None + self, use_memory_efficient_attention_xformers: bool, attention_op: Callable | None = None ) -> None: """ Set whether to use memory efficient attention from `xformers` or not. @@ -402,7 +402,7 @@ def head_to_batch_dim(self, tensor: torch.Tensor, out_dim: int = 3) -> torch.Ten return tensor def get_attention_scores( - self, query: torch.Tensor, key: torch.Tensor, attention_mask: Optional[torch.Tensor] = None + self, query: torch.Tensor, key: torch.Tensor, attention_mask: torch.Tensor | None = None ) -> torch.Tensor: """ Compute the attention scores. @@ -597,7 +597,7 @@ def __init__( num_attention_heads: int, attention_head_dim: int, context_pre_only: bool = False, - qk_norm: Optional[str] = None, + qk_norm: str | None = None, use_dual_attention: bool = False, ): super().__init__() @@ -673,7 +673,7 @@ def __init__( self._chunk_dim = 0 # Copied from diffusers.models.attention.BasicTransformerBlock.set_chunk_feed_forward - def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0): + def set_chunk_feed_forward(self, chunk_size: int | None, dim: int = 0): # Sets chunk feed-forward self._chunk_size = chunk_size self._chunk_dim = dim @@ -683,8 +683,8 @@ def forward( hidden_states: torch.FloatTensor, encoder_hidden_states: torch.FloatTensor, temb: torch.FloatTensor, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + joint_attention_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: joint_attention_kwargs = joint_attention_kwargs or {} if self.use_dual_attention: norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_hidden_states2, gate_msa2 = self.norm1( @@ -790,9 +790,9 @@ def __init__( num_attention_heads: int, attention_head_dim: int, dropout=0.0, - cross_attention_dim: Optional[int] = None, + cross_attention_dim: int | None = None, activation_fn: str = "geglu", - num_embeds_ada_norm: Optional[int] = None, + num_embeds_ada_norm: int | None = None, attention_bias: bool = False, only_cross_attention: bool = False, double_self_attention: bool = False, @@ -802,11 +802,11 @@ def __init__( norm_eps: float = 1e-5, final_dropout: bool = False, attention_type: str = "default", - positional_embeddings: Optional[str] = None, - num_positional_embeddings: Optional[int] = None, - ada_norm_continous_conditioning_embedding_dim: Optional[int] = None, - ada_norm_bias: Optional[int] = None, - ff_inner_dim: Optional[int] = None, + positional_embeddings: str | None = None, + num_positional_embeddings: int | None = None, + ada_norm_continous_conditioning_embedding_dim: int | None = None, + ada_norm_bias: int | None = None, + ff_inner_dim: int | None = None, ff_bias: bool = True, attention_out_bias: bool = True, ): @@ -952,7 +952,7 @@ def __init__( self._chunk_size = None self._chunk_dim = 0 - def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0): + def set_chunk_feed_forward(self, chunk_size: int | None, dim: int = 0): # Sets chunk feed-forward self._chunk_size = chunk_size self._chunk_dim = dim @@ -960,13 +960,13 @@ def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0): def forward( self, hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - timestep: Optional[torch.LongTensor] = None, - cross_attention_kwargs: Dict[str, Any] = None, - class_labels: Optional[torch.LongTensor] = None, - added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, + attention_mask: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + timestep: torch.LongTensor | None = None, + cross_attention_kwargs: dict[str, Any] = None, + class_labels: torch.LongTensor | None = None, + added_cond_kwargs: dict[str, torch.Tensor] | None = None, ) -> torch.Tensor: if cross_attention_kwargs is not None: if cross_attention_kwargs.get("scale", None) is not None: @@ -1099,8 +1099,8 @@ def __init__( self, dim: int, inner_dim: int, - multiple_of: Optional[int] = 256, - ffn_dim_multiplier: Optional[float] = None, + multiple_of: int | None = 256, + ffn_dim_multiplier: float | None = None, ): super().__init__() # custom hidden_size factor multiplier @@ -1148,7 +1148,7 @@ def __init__( time_mix_inner_dim: int, num_attention_heads: int, attention_head_dim: int, - cross_attention_dim: Optional[int] = None, + cross_attention_dim: int | None = None, ): super().__init__() self.is_res = dim == time_mix_inner_dim @@ -1195,7 +1195,7 @@ def __init__( self._chunk_size = None self._chunk_dim = None - def set_chunk_feed_forward(self, chunk_size: Optional[int], **kwargs): + def set_chunk_feed_forward(self, chunk_size: int | None, **kwargs): # Sets chunk feed-forward self._chunk_size = chunk_size # chunk dim should be hardcoded to 1 to have better speed vs. memory trade-off @@ -1205,7 +1205,7 @@ def forward( self, hidden_states: torch.Tensor, num_frames: int, - encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, ) -> torch.Tensor: # Notice that normalization is always applied before the real computation in the following blocks. # 0. Self-Attention @@ -1268,7 +1268,7 @@ def __init__( kv_input_dim: int, kv_input_dim_proj_use_bias: bool, dropout=0.0, - cross_attention_dim: Optional[int] = None, + cross_attention_dim: int | None = None, attention_bias: bool = False, attention_out_bias: bool = True, ): @@ -1393,9 +1393,9 @@ def __init__( num_attention_heads: int, attention_head_dim: int, dropout: float = 0.0, - cross_attention_dim: Optional[int] = None, + cross_attention_dim: int | None = None, activation_fn: str = "geglu", - num_embeds_ada_norm: Optional[int] = None, + num_embeds_ada_norm: int | None = None, attention_bias: bool = False, only_cross_attention: bool = False, double_self_attention: bool = False, @@ -1404,9 +1404,9 @@ def __init__( norm_type: str = "layer_norm", norm_eps: float = 1e-5, final_dropout: bool = False, - positional_embeddings: Optional[str] = None, - num_positional_embeddings: Optional[int] = None, - ff_inner_dim: Optional[int] = None, + positional_embeddings: str | None = None, + num_positional_embeddings: int | None = None, + ff_inner_dim: int | None = None, ff_bias: bool = True, attention_out_bias: bool = True, context_length: int = 16, @@ -1501,7 +1501,7 @@ def __init__( self._chunk_size = None self._chunk_dim = 0 - def _get_frame_indices(self, num_frames: int) -> List[Tuple[int, int]]: + def _get_frame_indices(self, num_frames: int) -> list[tuple[int, int]]: frame_indices = [] for i in range(0, num_frames - self.context_length + 1, self.context_stride): window_start = i @@ -1509,7 +1509,7 @@ def _get_frame_indices(self, num_frames: int) -> List[Tuple[int, int]]: frame_indices.append((window_start, window_end)) return frame_indices - def _get_frame_weights(self, num_frames: int, weighting_scheme: str = "pyramid") -> List[float]: + def _get_frame_weights(self, num_frames: int, weighting_scheme: str = "pyramid") -> list[float]: if weighting_scheme == "flat": weights = [1.0] * num_frames @@ -1548,7 +1548,7 @@ def set_free_noise_properties( self.context_stride = context_stride self.weighting_scheme = weighting_scheme - def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0) -> None: + def set_chunk_feed_forward(self, chunk_size: int | None, dim: int = 0) -> None: # Sets chunk feed-forward self._chunk_size = chunk_size self._chunk_dim = dim @@ -1556,10 +1556,10 @@ def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0) -> Non def forward( self, hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Dict[str, Any] = None, + attention_mask: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] = None, *args, **kwargs, ) -> torch.Tensor: @@ -1696,7 +1696,7 @@ class FeedForward(nn.Module): def __init__( self, dim: int, - dim_out: Optional[int] = None, + dim_out: int | None = None, mult: int = 4, dropout: float = 0.0, activation_fn: str = "geglu", diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index 310c44457c27..5b1f831ed060 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -12,15 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import contextlib import functools import inspect import math from dataclasses import dataclass from enum import Enum -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Callable import torch +import torch.distributed as dist +import torch.nn.functional as F if torch.distributed.is_available(): @@ -34,6 +38,7 @@ is_flash_attn_available, is_flash_attn_version, is_kernels_available, + is_kernels_version, is_sageattention_available, is_sageattention_version, is_torch_npu_available, @@ -44,6 +49,8 @@ is_xformers_version, ) from ..utils.constants import DIFFUSERS_ATTN_BACKEND, DIFFUSERS_ATTN_CHECKS +from ..utils.torch_utils import maybe_allow_in_graph +from ._modeling_parallel import gather_size_by_comm if TYPE_CHECKING: @@ -56,6 +63,8 @@ _REQUIRED_XLA_VERSION = "2.2" _REQUIRED_XFORMERS_VERSION = "0.0.29" +logger = get_logger(__name__) # pylint: disable=invalid-name + _CAN_USE_FLASH_ATTN = is_flash_attn_available() and is_flash_attn_version(">=", _REQUIRED_FLASH_VERSION) _CAN_USE_FLASH_ATTN_3 = is_flash_attn_3_available() _CAN_USE_AITER_ATTN = is_aiter_available() and is_aiter_version(">=", _REQUIRED_AITER_VERSION) @@ -67,8 +76,18 @@ if _CAN_USE_FLASH_ATTN: - from flash_attn import flash_attn_func, flash_attn_varlen_func - from flash_attn.flash_attn_interface import _wrapped_flash_attn_backward, _wrapped_flash_attn_forward + try: + from flash_attn import flash_attn_func, flash_attn_varlen_func + from flash_attn.flash_attn_interface import _wrapped_flash_attn_backward, _wrapped_flash_attn_forward + except (ImportError, OSError, RuntimeError) as e: + # Handle ABI mismatch or other import failures gracefully. + # This can happen when flash_attn was compiled against a different PyTorch version. + logger.warning(f"flash_attn is installed but failed to import: {e}. Falling back to native PyTorch attention.") + _CAN_USE_FLASH_ATTN = False + flash_attn_func = None + flash_attn_varlen_func = None + _wrapped_flash_attn_backward = None + _wrapped_flash_attn_forward = None else: flash_attn_func = None flash_attn_varlen_func = None @@ -77,26 +96,47 @@ if _CAN_USE_FLASH_ATTN_3: - from flash_attn_interface import flash_attn_func as flash_attn_3_func - from flash_attn_interface import flash_attn_varlen_func as flash_attn_3_varlen_func + try: + from flash_attn_interface import flash_attn_func as flash_attn_3_func + from flash_attn_interface import flash_attn_varlen_func as flash_attn_3_varlen_func + except (ImportError, OSError, RuntimeError) as e: + logger.warning(f"flash_attn_3 failed to import: {e}. Falling back to native attention.") + _CAN_USE_FLASH_ATTN_3 = False + flash_attn_3_func = None + flash_attn_3_varlen_func = None else: flash_attn_3_func = None flash_attn_3_varlen_func = None if _CAN_USE_AITER_ATTN: - from aiter import flash_attn_func as aiter_flash_attn_func + try: + from aiter import flash_attn_func as aiter_flash_attn_func + except (ImportError, OSError, RuntimeError) as e: + logger.warning(f"aiter failed to import: {e}. Falling back to native attention.") + _CAN_USE_AITER_ATTN = False + aiter_flash_attn_func = None else: aiter_flash_attn_func = None if _CAN_USE_SAGE_ATTN: - from sageattention import ( - sageattn, - sageattn_qk_int8_pv_fp8_cuda, - sageattn_qk_int8_pv_fp8_cuda_sm90, - sageattn_qk_int8_pv_fp16_cuda, - sageattn_qk_int8_pv_fp16_triton, - sageattn_varlen, - ) + try: + from sageattention import ( + sageattn, + sageattn_qk_int8_pv_fp8_cuda, + sageattn_qk_int8_pv_fp8_cuda_sm90, + sageattn_qk_int8_pv_fp16_cuda, + sageattn_qk_int8_pv_fp16_triton, + sageattn_varlen, + ) + except (ImportError, OSError, RuntimeError) as e: + logger.warning(f"sageattention failed to import: {e}. Falling back to native attention.") + _CAN_USE_SAGE_ATTN = False + sageattn = None + sageattn_qk_int8_pv_fp8_cuda = None + sageattn_qk_int8_pv_fp8_cuda_sm90 = None + sageattn_qk_int8_pv_fp16_cuda = None + sageattn_qk_int8_pv_fp16_triton = None + sageattn_varlen = None else: sageattn = None sageattn_qk_int8_pv_fp16_cuda = None @@ -107,26 +147,48 @@ if _CAN_USE_FLEX_ATTN: - # We cannot import the flex_attention function from the package directly because it is expected (from the - # pytorch documentation) that the user may compile it. If we import directly, we will not have access to the - # compiled function. - import torch.nn.attention.flex_attention as flex_attention + try: + # We cannot import the flex_attention function from the package directly because it is expected (from the + # pytorch documentation) that the user may compile it. If we import directly, we will not have access to the + # compiled function. + import torch.nn.attention.flex_attention as flex_attention + except (ImportError, OSError, RuntimeError) as e: + logger.warning(f"flex_attention failed to import: {e}. Falling back to native attention.") + _CAN_USE_FLEX_ATTN = False + flex_attention = None +else: + flex_attention = None if _CAN_USE_NPU_ATTN: - from torch_npu import npu_fusion_attention + try: + from torch_npu import npu_fusion_attention + except (ImportError, OSError, RuntimeError) as e: + logger.warning(f"torch_npu failed to import: {e}. Falling back to native attention.") + _CAN_USE_NPU_ATTN = False + npu_fusion_attention = None else: npu_fusion_attention = None if _CAN_USE_XLA_ATTN: - from torch_xla.experimental.custom_kernel import flash_attention as xla_flash_attention + try: + from torch_xla.experimental.custom_kernel import flash_attention as xla_flash_attention + except (ImportError, OSError, RuntimeError) as e: + logger.warning(f"torch_xla failed to import: {e}. Falling back to native attention.") + _CAN_USE_XLA_ATTN = False + xla_flash_attention = None else: xla_flash_attention = None if _CAN_USE_XFORMERS_ATTN: - import xformers.ops as xops + try: + import xformers.ops as xops + except (ImportError, OSError, RuntimeError) as e: + logger.warning(f"xformers failed to import: {e}. Falling back to native attention.") + _CAN_USE_XFORMERS_ATTN = False + xops = None else: xops = None @@ -152,8 +214,6 @@ def wrap(func): _register_fake = register_fake_no_op -logger = get_logger(__name__) # pylint: disable=invalid-name - # TODO(aryan): Add support for the following: # - Sage Attention++ # - block sparse, radial and other attention methods @@ -215,7 +275,7 @@ class _AttentionBackendRegistry: def register( cls, backend: AttentionBackendName, - constraints: Optional[List[Callable]] = None, + constraints: list[Callable] | None = None, supports_context_parallel: bool = False, ): logger.debug(f"Registering attention backend: {backend} with constraints: {constraints}") @@ -235,6 +295,10 @@ def decorator(func): def get_active_backend(cls): return cls._active_backend, cls._backends[cls._active_backend] + @classmethod + def set_active_backend(cls, backend: str): + cls._active_backend = backend + @classmethod def list_backends(cls): return list(cls._backends.keys()) @@ -254,35 +318,51 @@ class _HubKernelConfig: repo_id: str function_attr: str - revision: Optional[str] = None - kernel_fn: Optional[Callable] = None + revision: str | None = None + version: int | None = None + kernel_fn: Callable | None = None + wrapped_forward_attr: str | None = None + wrapped_backward_attr: str | None = None + wrapped_forward_fn: Callable | None = None + wrapped_backward_fn: Callable | None = None # Registry for hub-based attention kernels -_HUB_KERNELS_REGISTRY: Dict["AttentionBackendName", _HubKernelConfig] = { - # TODO: temporary revision for now. Remove when merged upstream into `main`. +_HUB_KERNELS_REGISTRY: dict["AttentionBackendName", _HubKernelConfig] = { AttentionBackendName._FLASH_3_HUB: _HubKernelConfig( - repo_id="kernels-community/flash-attn3", function_attr="flash_attn_func", revision="fake-ops-return-probs" + repo_id="kernels-community/flash-attn3", + function_attr="flash_attn_func", + wrapped_forward_attr="flash_attn_interface._flash_attn_forward", + wrapped_backward_attr="flash_attn_interface._flash_attn_backward", + version=1, ), AttentionBackendName._FLASH_3_VARLEN_HUB: _HubKernelConfig( repo_id="kernels-community/flash-attn3", function_attr="flash_attn_varlen_func", - # revision="fake-ops-return-probs", + version=1, ), AttentionBackendName.FLASH_HUB: _HubKernelConfig( - repo_id="kernels-community/flash-attn2", function_attr="flash_attn_func", revision=None + repo_id="kernels-community/flash-attn2", + function_attr="flash_attn_func", + wrapped_forward_attr="flash_attn_interface._wrapped_flash_attn_forward", + wrapped_backward_attr="flash_attn_interface._wrapped_flash_attn_backward", + version=1, ), AttentionBackendName.FLASH_VARLEN_HUB: _HubKernelConfig( - repo_id="kernels-community/flash-attn2", function_attr="flash_attn_varlen_func", revision=None + repo_id="kernels-community/flash-attn2", + function_attr="flash_attn_varlen_func", + version=1, ), AttentionBackendName.SAGE_HUB: _HubKernelConfig( - repo_id="kernels-community/sage_attention", function_attr="sageattn", revision=None + repo_id="kernels-community/sage-attention", + function_attr="sageattn", + version=1, ), } @contextlib.contextmanager -def attention_backend(backend: Union[str, AttentionBackendName] = AttentionBackendName.NATIVE): +def attention_backend(backend: str | AttentionBackendName = AttentionBackendName.NATIVE): """ Context manager to set the active attention backend. """ @@ -294,27 +374,27 @@ def attention_backend(backend: Union[str, AttentionBackendName] = AttentionBacke _maybe_download_kernel_for_backend(backend) old_backend = _AttentionBackendRegistry._active_backend - _AttentionBackendRegistry._active_backend = backend + _AttentionBackendRegistry.set_active_backend(backend) try: yield finally: - _AttentionBackendRegistry._active_backend = old_backend + _AttentionBackendRegistry.set_active_backend(old_backend) def dispatch_attention_fn( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[torch.Tensor] = None, + attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, enable_gqa: bool = False, - attention_kwargs: Optional[Dict[str, Any]] = None, + attention_kwargs: dict[str, Any] | None = None, *, - backend: Optional[AttentionBackendName] = None, - parallel_config: Optional["ParallelConfig"] = None, + backend: AttentionBackendName | None = None, + parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: attention_kwargs = attention_kwargs or {} @@ -348,6 +428,7 @@ def dispatch_attention_fn( check(**kwargs) kwargs = {k: v for k, v in kwargs.items() if k in _AttentionBackendRegistry._supported_arg_names[backend_name]} + return backend_fn(**kwargs) @@ -355,7 +436,7 @@ def dispatch_attention_fn( # A list of very simple functions to catch common errors quickly when debugging. -def _check_attn_mask_or_causal(attn_mask: Optional[torch.Tensor], is_causal: bool, **kwargs) -> None: +def _check_attn_mask_or_causal(attn_mask: torch.Tensor | None, is_causal: bool, **kwargs) -> None: if attn_mask is not None and is_causal: raise ValueError("`is_causal` cannot be True when `attn_mask` is not None.") @@ -401,7 +482,7 @@ def _check_shape( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[torch.Tensor] = None, + attn_mask: torch.Tensor | None = None, **kwargs, ) -> None: # Expected shapes: @@ -445,6 +526,10 @@ def _check_attention_backend_requirements(backend: AttentionBackendName) -> None raise RuntimeError( f"Backend '{backend.value}' is not usable because the `kernels` package isn't available. Please install it with `pip install kernels`." ) + if not is_kernels_version(">=", "0.12"): + raise RuntimeError( + f"Backend '{backend.value}' needs to be used with a `kernels` version of at least 0.12. Please update with `pip install -U kernels`." + ) elif backend == AttentionBackendName.AITER: if not _CAN_USE_AITER_ATTN: @@ -495,7 +580,7 @@ def _prepare_for_flash_attn_or_sage_varlen_without_mask( batch_size: int, seq_len_q: int, seq_len_kv: int, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): seqlens_q = torch.full((batch_size,), seq_len_q, dtype=torch.int32, device=device) seqlens_k = torch.full((batch_size,), seq_len_kv, dtype=torch.int32, device=device) @@ -512,7 +597,7 @@ def _prepare_for_flash_attn_or_sage_varlen_with_mask( batch_size: int, seq_len_q: int, attn_mask: torch.Tensor, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): seqlens_q = torch.full((batch_size,), seq_len_q, dtype=torch.int32, device=device) seqlens_k = attn_mask.sum(dim=1, dtype=torch.int32) @@ -529,8 +614,8 @@ def _prepare_for_flash_attn_or_sage_varlen( batch_size: int, seq_len_q: int, seq_len_kv: int, - attn_mask: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, + attn_mask: torch.Tensor | None = None, + device: torch.device | None = None, ) -> None: if attn_mask is None: return _prepare_for_flash_attn_or_sage_varlen_without_mask(batch_size, seq_len_q, seq_len_kv, device) @@ -594,22 +679,39 @@ def _flex_attention_causal_mask_mod(batch_idx, head_idx, q_idx, kv_idx): # ===== Helpers for downloading kernels ===== +def _resolve_kernel_attr(module, attr_path: str): + target = module + for attr in attr_path.split("."): + if not hasattr(target, attr): + raise AttributeError(f"Kernel module '{module.__name__}' does not define attribute path '{attr_path}'.") + target = getattr(target, attr) + return target + + def _maybe_download_kernel_for_backend(backend: AttentionBackendName) -> None: if backend not in _HUB_KERNELS_REGISTRY: return config = _HUB_KERNELS_REGISTRY[backend] - if config.kernel_fn is not None: + needs_kernel = config.kernel_fn is None + needs_wrapped_forward = config.wrapped_forward_attr is not None and config.wrapped_forward_fn is None + needs_wrapped_backward = config.wrapped_backward_attr is not None and config.wrapped_backward_fn is None + + if not (needs_kernel or needs_wrapped_forward or needs_wrapped_backward): return try: from kernels import get_kernel - kernel_module = get_kernel(config.repo_id, revision=config.revision) - kernel_func = getattr(kernel_module, config.function_attr) + kernel_module = get_kernel(config.repo_id, revision=config.revision, version=config.version) + if needs_kernel: + config.kernel_fn = _resolve_kernel_attr(kernel_module, config.function_attr) - # Cache the downloaded kernel function in the config object - config.kernel_fn = kernel_func + if needs_wrapped_forward: + config.wrapped_forward_fn = _resolve_kernel_attr(kernel_module, config.wrapped_forward_attr) + + if needs_wrapped_backward: + config.wrapped_backward_fn = _resolve_kernel_attr(kernel_module, config.wrapped_backward_attr) except Exception as e: logger.error(f"An error occurred while fetching kernel '{config.repo_id}' from the Hub: {e}") @@ -625,22 +727,22 @@ def _wrapped_flash_attn_3( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - softmax_scale: Optional[float] = None, + softmax_scale: float | None = None, causal: bool = False, - qv: Optional[torch.Tensor] = None, - q_descale: Optional[torch.Tensor] = None, - k_descale: Optional[torch.Tensor] = None, - v_descale: Optional[torch.Tensor] = None, + qv: torch.Tensor | None = None, + q_descale: torch.Tensor | None = None, + k_descale: torch.Tensor | None = None, + v_descale: torch.Tensor | None = None, attention_chunk: int = 0, softcap: float = 0.0, num_splits: int = 1, - pack_gqa: Optional[bool] = None, + pack_gqa: bool | None = None, deterministic: bool = False, sm_margin: int = 0, -) -> Tuple[torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor]: # Hardcoded for now because pytorch does not support tuple/int type hints window_size = (-1, -1) - out, lse, *_ = flash_attn_3_func( + result = flash_attn_3_func( q=q, k=k, v=v, @@ -657,7 +759,9 @@ def _wrapped_flash_attn_3( pack_gqa=pack_gqa, deterministic=deterministic, sm_margin=sm_margin, + return_attn_probs=True, ) + out, lse, *_ = result lse = lse.permute(0, 2, 1) return out, lse @@ -667,19 +771,19 @@ def _( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - softmax_scale: Optional[float] = None, + softmax_scale: float | None = None, causal: bool = False, - qv: Optional[torch.Tensor] = None, - q_descale: Optional[torch.Tensor] = None, - k_descale: Optional[torch.Tensor] = None, - v_descale: Optional[torch.Tensor] = None, + qv: torch.Tensor | None = None, + q_descale: torch.Tensor | None = None, + k_descale: torch.Tensor | None = None, + v_descale: torch.Tensor | None = None, attention_chunk: int = 0, softcap: float = 0.0, num_splits: int = 1, - pack_gqa: Optional[bool] = None, + pack_gqa: bool | None = None, deterministic: bool = False, sm_margin: int = 0, -) -> Tuple[torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor]: window_size = (-1, -1) # noqa: F841 # A lot of the parameters here are not yet used in any way within diffusers. # We can safely ignore for now and keep the fake op shape propagation simple. @@ -696,14 +800,14 @@ def _native_attention_forward_op( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[torch.Tensor] = None, + attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, enable_gqa: bool = False, return_lse: bool = False, _save_ctx: bool = True, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ): # Native attention does not return_lse if return_lse: @@ -779,14 +883,14 @@ def _cudnn_attention_forward_op( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[torch.Tensor] = None, + attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, enable_gqa: bool = False, return_lse: bool = False, _save_ctx: bool = True, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ): if enable_gqa: raise ValueError("`enable_gqa` is not yet supported for cuDNN attention.") @@ -876,14 +980,14 @@ def _native_flash_attention_forward_op( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[torch.Tensor] = None, + attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, enable_gqa: bool = False, return_lse: bool = False, _save_ctx: bool = True, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ): if enable_gqa: raise ValueError("`enable_gqa` is not yet supported for native flash attention.") @@ -965,14 +1069,14 @@ def _flash_attention_forward_op( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[torch.Tensor] = None, + attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, enable_gqa: bool = False, return_lse: bool = False, _save_ctx: bool = True, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ): if attn_mask is not None: raise ValueError("`attn_mask` is not yet supported for flash-attn 2.") @@ -1060,19 +1164,271 @@ def _flash_attention_backward_op( return grad_query, grad_key, grad_value +def _flash_attention_hub_forward_op( + ctx: torch.autograd.function.FunctionCtx, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor | None = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: float | None = None, + enable_gqa: bool = False, + return_lse: bool = False, + _save_ctx: bool = True, + _parallel_config: "ParallelConfig" | None = None, +): + if attn_mask is not None: + raise ValueError("`attn_mask` is not yet supported for flash-attn hub kernels.") + if enable_gqa: + raise ValueError("`enable_gqa` is not yet supported for flash-attn hub kernels.") + + config = _HUB_KERNELS_REGISTRY[AttentionBackendName.FLASH_HUB] + wrapped_forward_fn = config.wrapped_forward_fn + wrapped_backward_fn = config.wrapped_backward_fn + if wrapped_forward_fn is None or wrapped_backward_fn is None: + raise RuntimeError( + "Flash attention hub kernels must expose `_wrapped_flash_attn_forward` and `_wrapped_flash_attn_backward` " + "for context parallel execution." + ) + + if scale is None: + scale = query.shape[-1] ** (-0.5) + + window_size = (-1, -1) + softcap = 0.0 + alibi_slopes = None + deterministic = False + grad_enabled = any(x.requires_grad for x in (query, key, value)) + + if grad_enabled or (_parallel_config is not None and _parallel_config.context_parallel_config._world_size > 1): + dropout_p = dropout_p if dropout_p > 0 else 1e-30 + + with torch.set_grad_enabled(grad_enabled): + out, lse, S_dmask, rng_state = wrapped_forward_fn( + query, + key, + value, + dropout_p, + scale, + is_causal, + window_size[0], + window_size[1], + softcap, + alibi_slopes, + return_lse, + ) + lse = lse.permute(0, 2, 1).contiguous() + + if _save_ctx: + ctx.save_for_backward(query, key, value, out, lse, rng_state) + ctx.dropout_p = dropout_p + ctx.scale = scale + ctx.is_causal = is_causal + ctx.window_size = window_size + ctx.softcap = softcap + ctx.alibi_slopes = alibi_slopes + ctx.deterministic = deterministic + + return (out, lse) if return_lse else out + + +def _flash_attention_hub_backward_op( + ctx: torch.autograd.function.FunctionCtx, + grad_out: torch.Tensor, + *args, + **kwargs, +): + config = _HUB_KERNELS_REGISTRY[AttentionBackendName.FLASH_HUB] + wrapped_backward_fn = config.wrapped_backward_fn + if wrapped_backward_fn is None: + raise RuntimeError( + "Flash attention hub kernels must expose `_wrapped_flash_attn_backward` for context parallel execution." + ) + + query, key, value, out, lse, rng_state = ctx.saved_tensors + grad_query, grad_key, grad_value = torch.empty_like(query), torch.empty_like(key), torch.empty_like(value) + + _ = wrapped_backward_fn( + grad_out, + query, + key, + value, + out, + lse, + grad_query, + grad_key, + grad_value, + ctx.dropout_p, + ctx.scale, + ctx.is_causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.alibi_slopes, + ctx.deterministic, + rng_state, + ) + + grad_query = grad_query[..., : grad_out.shape[-1]] + grad_key = grad_key[..., : grad_out.shape[-1]] + grad_value = grad_value[..., : grad_out.shape[-1]] + + return grad_query, grad_key, grad_value + + +def _flash_attention_3_hub_forward_op( + ctx: torch.autograd.function.FunctionCtx, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor | None = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: float | None = None, + enable_gqa: bool = False, + return_lse: bool = False, + _save_ctx: bool = True, + _parallel_config: "ParallelConfig" | None = None, + *, + window_size: tuple[int, int] = (-1, -1), + softcap: float = 0.0, + num_splits: int = 1, + pack_gqa: bool | None = None, + deterministic: bool = False, + sm_margin: int = 0, +): + if attn_mask is not None: + raise ValueError("`attn_mask` is not yet supported for flash-attn 3 hub kernels.") + if dropout_p != 0.0: + raise ValueError("`dropout_p` is not yet supported for flash-attn 3 hub kernels.") + if enable_gqa: + raise ValueError("`enable_gqa` is not yet supported for flash-attn 3 hub kernels.") + + config = _HUB_KERNELS_REGISTRY[AttentionBackendName._FLASH_3_HUB] + wrapped_forward_fn = config.wrapped_forward_fn + if wrapped_forward_fn is None: + raise RuntimeError( + "Flash attention 3 hub kernels must expose `flash_attn_interface._flash_attn_forward` " + "for context parallel execution." + ) + + if scale is None: + scale = query.shape[-1] ** (-0.5) + + out, softmax_lse, *_ = wrapped_forward_fn( + query, + key, + value, + None, + None, # k_new, v_new + None, # qv + None, # out + None, + None, + None, # cu_seqlens_q/k/k_new + None, + None, # seqused_q/k + None, + None, # max_seqlen_q/k + None, + None, + None, # page_table, kv_batch_idx, leftpad_k + None, + None, + None, # rotary_cos/sin, seqlens_rotary + None, + None, + None, # q_descale, k_descale, v_descale + scale, + causal=is_causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=0, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + + lse = softmax_lse.permute(0, 2, 1).contiguous() if return_lse else None + + if _save_ctx: + ctx.save_for_backward(query, key, value, out, softmax_lse) + ctx.scale = scale + ctx.is_causal = is_causal + ctx.window_size = window_size + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + + return (out, lse) if return_lse else out + + +def _flash_attention_3_hub_backward_op( + ctx: torch.autograd.function.FunctionCtx, + grad_out: torch.Tensor, + *args, + **kwargs, +): + config = _HUB_KERNELS_REGISTRY[AttentionBackendName._FLASH_3_HUB] + wrapped_backward_fn = config.wrapped_backward_fn + if wrapped_backward_fn is None: + raise RuntimeError( + "Flash attention 3 hub kernels must expose `flash_attn_interface._flash_attn_backward` " + "for context parallel execution." + ) + + query, key, value, out, softmax_lse = ctx.saved_tensors + grad_query = torch.empty_like(query) + grad_key = torch.empty_like(key) + grad_value = torch.empty_like(value) + + wrapped_backward_fn( + grad_out, + query, + key, + value, + out, + softmax_lse, + None, + None, # cu_seqlens_q, cu_seqlens_k + None, + None, # seqused_q, seqused_k + None, + None, # max_seqlen_q, max_seqlen_k + grad_query, + grad_key, + grad_value, + ctx.scale, + ctx.is_causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + + grad_query = grad_query[..., : grad_out.shape[-1]] + grad_key = grad_key[..., : grad_out.shape[-1]] + grad_value = grad_value[..., : grad_out.shape[-1]] + + return grad_query, grad_key, grad_value + + def _sage_attention_forward_op( ctx: torch.autograd.function.FunctionCtx, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[torch.Tensor] = None, + attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, enable_gqa: bool = False, return_lse: bool = False, _save_ctx: bool = True, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ): if attn_mask is not None: raise ValueError("`attn_mask` is not yet supported for Sage attention.") @@ -1098,6 +1454,46 @@ def _sage_attention_forward_op( return (out, lse) if return_lse else out +def _sage_attention_hub_forward_op( + ctx: torch.autograd.function.FunctionCtx, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor | None = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: float | None = None, + enable_gqa: bool = False, + return_lse: bool = False, + _save_ctx: bool = True, + _parallel_config: "ParallelConfig" | None = None, +): + if attn_mask is not None: + raise ValueError("`attn_mask` is not yet supported for Sage attention.") + if dropout_p > 0.0: + raise ValueError("`dropout_p` is not yet supported for Sage attention.") + if enable_gqa: + raise ValueError("`enable_gqa` is not yet supported for Sage attention.") + + func = _HUB_KERNELS_REGISTRY[AttentionBackendName.SAGE_HUB].kernel_fn + out = func( + q=query, + k=key, + v=value, + tensor_layout="NHD", + is_causal=is_causal, + sm_scale=scale, + return_lse=return_lse, + ) + + lse = None + if return_lse: + out, lse, *_ = out + lse = lse.permute(0, 2, 1).contiguous() + + return (out, lse) if return_lse else out + + def _sage_attention_backward_op( ctx: torch.autograd.function.FunctionCtx, grad_out: torch.Tensor, @@ -1106,48 +1502,361 @@ def _sage_attention_backward_op( raise NotImplementedError("Backward pass is not implemented for Sage attention.") -# ===== Context parallel ===== - - -# Reference: -# - https://github.com/pytorch/pytorch/blob/f58a680d09e13658a52c6ba05c63c15759846bcc/torch/distributed/_functional_collectives.py#L827 -# - https://github.com/pytorch/pytorch/blob/f58a680d09e13658a52c6ba05c63c15759846bcc/torch/distributed/_functional_collectives.py#L246 -# For fullgraph=True tracing compatibility (since FakeTensor does not have a `wait` method): -def _wait_tensor(tensor): - if isinstance(tensor, funcol.AsyncCollectiveTensor): - tensor = tensor.wait() - return tensor +def _maybe_modify_attn_mask_npu(query: torch.Tensor, key: torch.Tensor, attn_mask: torch.Tensor | None = None): + # Skip Attention Mask if all values are 1, `None` mask can speedup the computation + if attn_mask is not None and torch.all(attn_mask != 0): + attn_mask = None + # Reshape Attention Mask: [batch_size, seq_len_k] -> [batch_size, 1, sqe_len_q, seq_len_k] + # https://www.hiascend.com/document/detail/zh/Pytorch/730/apiref/torchnpuCustomsapi/docs/context/torch_npu-npu_fusion_attention.md + if ( + attn_mask is not None + and attn_mask.ndim == 2 + and attn_mask.shape[0] == query.shape[0] + and attn_mask.shape[1] == key.shape[1] + ): + B, Sq, Skv = attn_mask.shape[0], query.shape[1], key.shape[1] + attn_mask = ~attn_mask.to(torch.bool) + attn_mask = attn_mask.unsqueeze(1).expand(B, Sq, Skv).unsqueeze(1).contiguous() -def _all_to_all_single(x: torch.Tensor, group) -> torch.Tensor: - shape = x.shape - # HACK: We need to flatten because despite making tensors contiguous, torch single-file-ization - # to benchmark triton codegen fails somewhere: - # buf25 = torch.ops._c10d_functional.all_to_all_single.default(buf24, [1, 1], [1, 1], '3') - # ValueError: Tensors must be contiguous - x = x.flatten() - x = funcol.all_to_all_single(x, None, None, group) - x = x.reshape(shape) - x = _wait_tensor(x) - return x + return attn_mask -class TemplatedRingAttention(torch.autograd.Function): - @staticmethod - def forward( - ctx: torch.autograd.function.FunctionCtx, - query: torch.Tensor, - key: torch.Tensor, +def _npu_attention_forward_op( + ctx: torch.autograd.function.FunctionCtx, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor | None = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: float | None = None, + enable_gqa: bool = False, + return_lse: bool = False, + _save_ctx: bool = True, + _parallel_config: "ParallelConfig" | None = None, +): + if return_lse: + raise ValueError("NPU attention backend does not support setting `return_lse=True`.") + + attn_mask = _maybe_modify_attn_mask_npu(query, key, attn_mask) + + out = npu_fusion_attention( + query, + key, + value, + query.size(2), # num_heads + atten_mask=attn_mask, + input_layout="BSND", + pse=None, + scale=1.0 / math.sqrt(query.shape[-1]) if scale is None else scale, + pre_tockens=65536, + next_tockens=65536, + keep_prob=1.0 - dropout_p, + sync=False, + inner_precise=0, + )[0] + + return out + + +# Not implemented yet. +def _npu_attention_backward_op( + ctx: torch.autograd.function.FunctionCtx, + grad_out: torch.Tensor, + *args, + **kwargs, +): + raise NotImplementedError("Backward pass is not implemented for Npu Fusion Attention.") + + +# ===== Context parallel ===== + + +# Reference: +# - https://github.com/pytorch/pytorch/blob/f58a680d09e13658a52c6ba05c63c15759846bcc/torch/distributed/_functional_collectives.py#L827 +# - https://github.com/pytorch/pytorch/blob/f58a680d09e13658a52c6ba05c63c15759846bcc/torch/distributed/_functional_collectives.py#L246 +# For fullgraph=True tracing compatibility (since FakeTensor does not have a `wait` method): +def _wait_tensor(tensor): + if isinstance(tensor, funcol.AsyncCollectiveTensor): + tensor = tensor.wait() + return tensor + + +def _all_to_all_single(x: torch.Tensor, group) -> torch.Tensor: + shape = x.shape + # HACK: We need to flatten because despite making tensors contiguous, torch single-file-ization + # to benchmark triton codegen fails somewhere: + # buf25 = torch.ops._c10d_functional.all_to_all_single.default(buf24, [1, 1], [1, 1], '3') + # ValueError: Tensors must be contiguous + x = x.flatten() + x = funcol.all_to_all_single(x, None, None, group) + x = x.reshape(shape) + x = _wait_tensor(x) + return x + + +def _all_to_all_dim_exchange(x: torch.Tensor, scatter_idx: int = 2, gather_idx: int = 1, group=None) -> torch.Tensor: + """ + Perform dimension sharding / reassembly across processes using _all_to_all_single. + + This utility reshapes and redistributes tensor `x` across the given process group, across sequence dimension or + head dimension flexibly by accepting scatter_idx and gather_idx. + + Args: + x (torch.Tensor): + Input tensor. Expected shapes: + - When scatter_idx=2, gather_idx=1: (batch_size, seq_len_local, num_heads, head_dim) + - When scatter_idx=1, gather_idx=2: (batch_size, seq_len, num_heads_local, head_dim) + scatter_idx (int) : + Dimension along which the tensor is partitioned before all-to-all. + gather_idx (int): + Dimension along which the output is reassembled after all-to-all. + group : + Distributed process group for the Ulysses group. + + Returns: + torch.Tensor: Tensor with globally exchanged dimensions. + - For (scatter_idx=2 → gather_idx=1): (batch_size, seq_len, num_heads_local, head_dim) + - For (scatter_idx=1 → gather_idx=2): (batch_size, seq_len_local, num_heads, head_dim) + """ + group_world_size = torch.distributed.get_world_size(group) + + if scatter_idx == 2 and gather_idx == 1: + # Used before Ulysses sequence parallel (SP) attention. Scatters the gathers sequence + # dimension and scatters head dimension + batch_size, seq_len_local, num_heads, head_dim = x.shape + seq_len = seq_len_local * group_world_size + num_heads_local = num_heads // group_world_size + + # B, S_LOCAL, H, D -> group_world_size, S_LOCAL, B, H_LOCAL, D + x_temp = ( + x.reshape(batch_size, seq_len_local, group_world_size, num_heads_local, head_dim) + .transpose(0, 2) + .contiguous() + ) + + if group_world_size > 1: + out = _all_to_all_single(x_temp, group=group) + else: + out = x_temp + # group_world_size, S_LOCAL, B, H_LOCAL, D -> B, S, H_LOCAL, D + out = out.reshape(seq_len, batch_size, num_heads_local, head_dim).permute(1, 0, 2, 3).contiguous() + out = out.reshape(batch_size, seq_len, num_heads_local, head_dim) + return out + elif scatter_idx == 1 and gather_idx == 2: + # Used after ulysses sequence parallel in unified SP. gathers the head dimension + # scatters back the sequence dimension. + batch_size, seq_len, num_heads_local, head_dim = x.shape + num_heads = num_heads_local * group_world_size + seq_len_local = seq_len // group_world_size + + # B, S, H_LOCAL, D -> group_world_size, H_LOCAL, S_LOCAL, B, D + x_temp = ( + x.reshape(batch_size, group_world_size, seq_len_local, num_heads_local, head_dim) + .permute(1, 3, 2, 0, 4) + .reshape(group_world_size, num_heads_local, seq_len_local, batch_size, head_dim) + ) + + if group_world_size > 1: + output = _all_to_all_single(x_temp, group) + else: + output = x_temp + output = output.reshape(num_heads, seq_len_local, batch_size, head_dim).transpose(0, 2).contiguous() + output = output.reshape(batch_size, seq_len_local, num_heads, head_dim) + return output + else: + raise ValueError("Invalid scatter/gather indices for _all_to_all_dim_exchange.") + + +class SeqAllToAllDim(torch.autograd.Function): + """ + all_to_all operation for unified sequence parallelism. uses _all_to_all_dim_exchange, see _all_to_all_dim_exchange + for more info. + """ + + @staticmethod + def forward(ctx, group, input, scatter_id=2, gather_id=1): + ctx.group = group + ctx.scatter_id = scatter_id + ctx.gather_id = gather_id + return _all_to_all_dim_exchange(input, scatter_id, gather_id, group) + + @staticmethod + def backward(ctx, grad_outputs): + grad_input = SeqAllToAllDim.apply( + ctx.group, + grad_outputs, + ctx.gather_id, # reversed + ctx.scatter_id, # reversed + ) + return (None, grad_input, None, None) + + +# Below are helper functions to handle abritrary head num and abritrary sequence length for Ulysses Anything Attention. +def _maybe_pad_qkv_head(x: torch.Tensor, H: int, group: dist.ProcessGroup) -> tuple[torch.Tensor, int]: + r"""Maybe pad the head dimension to be divisible by world_size. + x: torch.Tensor, shape (B, S_LOCAL, H, D) H: int, original global head num return: tuple[torch.Tensor, int], padded + tensor (B, S_LOCAL, H + H_PAD, D) and H_PAD + """ + world_size = dist.get_world_size(group=group) + H_PAD = 0 + if H % world_size != 0: + H_PAD = world_size - (H % world_size) + NEW_H_LOCAL = (H + H_PAD) // world_size + # e.g., Allow: H=30, world_size=8 -> NEW_H_LOCAL=4, H_PAD=2. + # NOT ALLOW: H=30, world_size=16 -> NEW_H_LOCAL=2, H_PAD=14. + assert H_PAD < NEW_H_LOCAL, f"Padding head num {H_PAD} should be less than new local head num {NEW_H_LOCAL}" + x = F.pad(x, (0, 0, 0, H_PAD)).contiguous() + return x, H_PAD + + +def _maybe_unpad_qkv_head(x: torch.Tensor, H_PAD: int, group: dist.ProcessGroup) -> torch.Tensor: + r"""Maybe unpad the head dimension. + x: torch.Tensor, shape (B, S_GLOBAL, H_LOCAL + H_PAD, D) H_PAD: int, head padding num return: torch.Tensor, + unpadded tensor (B, S_GLOBAL, H_LOCAL, D) + """ + rank = dist.get_rank(group=group) + world_size = dist.get_world_size(group=group) + # Only the last rank may have padding + if H_PAD > 0 and rank == world_size - 1: + x = x[:, :, :-H_PAD, :] + return x.contiguous() + + +def _maybe_pad_o_head(x: torch.Tensor, H: int, group: dist.ProcessGroup) -> tuple[torch.Tensor, int]: + r"""Maybe pad the head dimension to be divisible by world_size. + x: torch.Tensor, shape (B, S_GLOBAL, H_LOCAL, D) H: int, original global head num return: tuple[torch.Tensor, int], + padded tensor (B, S_GLOBAL, H_LOCAL + H_PAD, D) and H_PAD + """ + if H is None: + return x, 0 + + rank = dist.get_rank(group=group) + world_size = dist.get_world_size(group=group) + H_PAD = 0 + # Only the last rank may need padding + if H % world_size != 0: + # We need to broadcast H_PAD to all ranks to keep consistency + # in unpadding step later for all ranks. + H_PAD = world_size - (H % world_size) + NEW_H_LOCAL = (H + H_PAD) // world_size + assert H_PAD < NEW_H_LOCAL, f"Padding head num {H_PAD} should be less than new local head num {NEW_H_LOCAL}" + if rank == world_size - 1: + x = F.pad(x, (0, 0, 0, H_PAD)).contiguous() + return x, H_PAD + + +def _maybe_unpad_o_head(x: torch.Tensor, H_PAD: int, group: dist.ProcessGroup) -> torch.Tensor: + r"""Maybe unpad the head dimension. + x: torch.Tensor, shape (B, S_LOCAL, H_GLOBAL + H_PAD, D) H_PAD: int, head padding num return: torch.Tensor, + unpadded tensor (B, S_LOCAL, H_GLOBAL, D) + """ + if H_PAD > 0: + x = x[:, :, :-H_PAD, :] + return x.contiguous() + + +def ulysses_anything_metadata(query: torch.Tensor, **kwargs) -> dict: + # query: (B, S_LOCAL, H_GLOBAL, D) + assert len(query.shape) == 4, "Query tensor must be 4-dimensional of shape (B, S_LOCAL, H_GLOBAL, D)" + extra_kwargs = {} + extra_kwargs["NUM_QO_HEAD"] = query.shape[2] + extra_kwargs["Q_S_LOCAL"] = query.shape[1] + # Add other kwargs if needed in future + return extra_kwargs + + +@maybe_allow_in_graph +def all_to_all_single_any_qkv_async( + x: torch.Tensor, group: dist.ProcessGroup, **kwargs +) -> Callable[..., torch.Tensor]: + r""" + x: torch.Tensor, shape (B, S_LOCAL, H, D) return: Callable that returns (B, S_GLOBAL, H_LOCAL, D) + """ + world_size = dist.get_world_size(group=group) + B, S_LOCAL, H, D = x.shape + x, H_PAD = _maybe_pad_qkv_head(x, H, group) + H_LOCAL = (H + H_PAD) // world_size + # (world_size, S_LOCAL, B, H_LOCAL, D) + x = x.reshape(B, S_LOCAL, world_size, H_LOCAL, D).permute(2, 1, 0, 3, 4).contiguous() + + input_split_sizes = [S_LOCAL] * world_size + # S_LOCAL maybe not equal for all ranks in dynamic shape case, + # since we don't know the actual shape before this timing, thus, + # we have to use all gather to collect the S_LOCAL first. + output_split_sizes = gather_size_by_comm(S_LOCAL, group) + x = x.flatten(0, 1) # (world_size * S_LOCAL, B, H_LOCAL, D) + x = funcol.all_to_all_single(x, output_split_sizes, input_split_sizes, group) + + def wait() -> torch.Tensor: + nonlocal x, H_PAD + x = _wait_tensor(x) # (S_GLOBAL, B, H_LOCAL, D) + # (S_GLOBAL, B, H_LOCAL, D) + # -> (B, S_GLOBAL, H_LOCAL, D) + x = x.permute(1, 0, 2, 3).contiguous() + x = _maybe_unpad_qkv_head(x, H_PAD, group) + return x + + return wait + + +@maybe_allow_in_graph +def all_to_all_single_any_o_async(x: torch.Tensor, group: dist.ProcessGroup, **kwargs) -> Callable[..., torch.Tensor]: + r""" + x: torch.Tensor, shape (B, S_GLOBAL, H_LOCAL, D) return: Callable that returns (B, S_LOCAL, H_GLOBAL, D) + """ + # Assume H is provided in kwargs, since we can't infer H from x's shape. + # The padding logic needs H to determine if padding is necessary. + H = kwargs.get("NUM_QO_HEAD", None) + world_size = dist.get_world_size(group=group) + + x, H_PAD = _maybe_pad_o_head(x, H, group) + shape = x.shape # (B, S_GLOBAL, H_LOCAL, D) + (B, S_GLOBAL, H_LOCAL, D) = shape + + # input_split: e.g, S_GLOBAL=9 input splits across ranks [[5,4], [5,4],..] + # output_split: e.g, S_GLOBAL=9 output splits across ranks [[5,5], [4,4],..] + + # WARN: In some cases, e.g, joint attn in Qwen-Image, the S_LOCAL can not infer + # from tensor split due to: if c = torch.cat((a, b)), world_size=4, then, + # c.tensor_split(4)[0].shape[1] may != to (a.tensor_split(4)[0].shape[1] + + # b.tensor_split(4)[0].shape[1]) + + S_LOCAL = kwargs.get("Q_S_LOCAL") + input_split_sizes = gather_size_by_comm(S_LOCAL, group) + x = x.permute(1, 0, 2, 3).contiguous() # (S_GLOBAL, B, H_LOCAL, D) + output_split_sizes = [S_LOCAL] * world_size + x = funcol.all_to_all_single(x, output_split_sizes, input_split_sizes, group) + + def wait() -> torch.Tensor: + nonlocal x, H_PAD + x = _wait_tensor(x) # (S_GLOBAL, B, H_LOCAL, D) + x = x.reshape(world_size, S_LOCAL, B, H_LOCAL, D) + x = x.permute(2, 1, 0, 3, 4).contiguous() + x = x.reshape(B, S_LOCAL, world_size * H_LOCAL, D) + x = _maybe_unpad_o_head(x, H_PAD, group) + return x + + return wait + + +class TemplatedRingAttention(torch.autograd.Function): + @staticmethod + def forward( + ctx: torch.autograd.function.FunctionCtx, + query: torch.Tensor, + key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[torch.Tensor], + attn_mask: torch.Tensor | None, dropout_p: float, is_causal: bool, - scale: Optional[float], + scale: float | None, enable_gqa: bool, return_lse: bool, forward_op, backward_op, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ): ring_mesh = _parallel_config.context_parallel_config._ring_mesh rank = _parallel_config.context_parallel_config._ring_local_rank @@ -1192,7 +1901,10 @@ def forward( out = out.to(torch.float32) lse = lse.to(torch.float32) - lse = lse.unsqueeze(-1) + # Refer to: + # https://github.com/huggingface/diffusers/pull/12693#issuecomment-3627519544 + if is_torch_version("<", "2.9.0"): + lse = lse.unsqueeze(-1) if prev_out is not None: out = prev_out - torch.nn.functional.sigmoid(lse - prev_lse) * (prev_out - out) lse = prev_lse - torch.nn.functional.logsigmoid(prev_lse - lse) @@ -1253,7 +1965,7 @@ def backward( grad_query, grad_key, grad_value = (x.to(grad_out.dtype) for x in (grad_query, grad_key, grad_value)) - return grad_query, grad_key, grad_value, None, None, None, None, None, None, None, None + return grad_query, grad_key, grad_value, None, None, None, None, None, None, None, None, None class TemplatedUlyssesAttention(torch.autograd.Function): @@ -1263,15 +1975,15 @@ def forward( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[torch.Tensor], + attn_mask: torch.Tensor | None, dropout_p: float, is_causal: bool, - scale: Optional[float], + scale: float | None, enable_gqa: bool, return_lse: bool, forward_op, backward_op, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ): ulysses_mesh = _parallel_config.context_parallel_config._ulysses_mesh world_size = _parallel_config.context_parallel_config.ulysses_degree @@ -1348,34 +2060,173 @@ def backward( x.flatten(0, 1).permute(1, 2, 0, 3).contiguous() for x in (grad_query, grad_key, grad_value) ) - return grad_query, grad_key, grad_value, None, None, None, None, None, None, None, None + return grad_query, grad_key, grad_value, None, None, None, None, None, None, None, None, None + + +class TemplatedUlyssesAnythingAttention(torch.autograd.Function): + @staticmethod + def forward( + ctx: torch.autograd.function.FunctionCtx, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor, + dropout_p: float, + is_causal: bool, + scale: float, + enable_gqa: bool, + return_lse: bool, + forward_op, + backward_op, + _parallel_config: "ParallelConfig" | None = None, + **kwargs, + ): + ulysses_mesh = _parallel_config.context_parallel_config._ulysses_mesh + group = ulysses_mesh.get_group() + + ctx.forward_op = forward_op + ctx.backward_op = backward_op + ctx._parallel_config = _parallel_config + + metadata = ulysses_anything_metadata(query) + query_wait = all_to_all_single_any_qkv_async(query, group, **metadata) + key_wait = all_to_all_single_any_qkv_async(key, group, **metadata) + value_wait = all_to_all_single_any_qkv_async(value, group, **metadata) + + query = query_wait() # type: torch.Tensor + key = key_wait() # type: torch.Tensor + value = value_wait() # type: torch.Tensor + + out = forward_op( + ctx, + query, + key, + value, + attn_mask, + dropout_p, + is_causal, + scale, + enable_gqa, + return_lse, + _save_ctx=False, # ulysses anything only support forward pass now. + _parallel_config=_parallel_config, + ) + if return_lse: + out, lse, *_ = out + + # out: (B, S_Q_GLOBAL, H_LOCAL, D) -> (B, S_Q_LOCAL, H_GLOBAL, D) + out_wait = all_to_all_single_any_o_async(out, group, **metadata) + + if return_lse: + # lse: (B, S_Q_GLOBAL, H_LOCAL) + lse = lse.unsqueeze(-1) # (B, S_Q_GLOBAL, H_LOCAL, D=1) + lse_wait = all_to_all_single_any_o_async(lse, group, **metadata) + out = out_wait() # type: torch.Tensor + lse = lse_wait() # type: torch.Tensor + lse = lse.squeeze(-1).contiguous() # (B, S_Q_LOCAL, H_GLOBAL) + else: + out = out_wait() # type: torch.Tensor + lse = None + + return (out, lse) if return_lse else out + + @staticmethod + def backward( + ctx: torch.autograd.function.FunctionCtx, + grad_out: torch.Tensor, + *args, + ): + raise NotImplementedError("Backward pass for Ulysses Anything Attention in diffusers is not implemented yet.") + + +def _templated_unified_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor, + dropout_p: float, + is_causal: bool, + scale: float, + enable_gqa: bool, + return_lse: bool, + forward_op, + backward_op, + _parallel_config: "ParallelConfig" | None = None, + scatter_idx: int = 2, + gather_idx: int = 1, +): + """ + Unified Sequence Parallelism attention combining Ulysses and ring attention. See: https://arxiv.org/abs/2405.07719 + """ + ulysses_mesh = _parallel_config.context_parallel_config._ulysses_mesh + ulysses_group = ulysses_mesh.get_group() + + query = SeqAllToAllDim.apply(ulysses_group, query, scatter_idx, gather_idx) + key = SeqAllToAllDim.apply(ulysses_group, key, scatter_idx, gather_idx) + value = SeqAllToAllDim.apply(ulysses_group, value, scatter_idx, gather_idx) + out = TemplatedRingAttention.apply( + query, + key, + value, + attn_mask, + dropout_p, + is_causal, + scale, + enable_gqa, + return_lse, + forward_op, + backward_op, + _parallel_config, + ) + if return_lse: + context_layer, lse, *_ = out + else: + context_layer = out + # context_layer is of shape (B, S, H_LOCAL, D) + output = SeqAllToAllDim.apply( + ulysses_group, + context_layer, + gather_idx, + scatter_idx, + ) + if return_lse: + # lse is of shape (B, S, H_LOCAL, 1) + # Refer to: + # https://github.com/huggingface/diffusers/pull/12693#issuecomment-3627519544 + if is_torch_version("<", "2.9.0"): + lse = lse.unsqueeze(-1) # (B, S, H_LOCAL, 1) + lse = SeqAllToAllDim.apply(ulysses_group, lse, gather_idx, scatter_idx) + lse = lse.squeeze(-1) + return (output, lse) + return output def _templated_context_parallel_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[torch.Tensor] = None, + attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, enable_gqa: bool = False, return_lse: bool = False, *, forward_op, backward_op, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ): - if attn_mask is not None: - raise ValueError("Attention mask is not yet supported for templated attention.") if is_causal: raise ValueError("Causal attention is not yet supported for templated attention.") if enable_gqa: raise ValueError("GQA is not yet supported for templated attention.") # TODO: add support for unified attention with ring/ulysses degree both being > 1 - if _parallel_config.context_parallel_config.ring_degree > 1: - return TemplatedRingAttention.apply( + if ( + _parallel_config.context_parallel_config.ring_degree > 1 + and _parallel_config.context_parallel_config.ulysses_degree > 1 + ): + return _templated_unified_attention( query, key, value, @@ -1389,8 +2240,8 @@ def _templated_context_parallel_attention( backward_op, _parallel_config, ) - elif _parallel_config.context_parallel_config.ulysses_degree > 1: - return TemplatedUlyssesAttention.apply( + elif _parallel_config.context_parallel_config.ring_degree > 1: + return TemplatedRingAttention.apply( query, key, value, @@ -1404,6 +2255,38 @@ def _templated_context_parallel_attention( backward_op, _parallel_config, ) + elif _parallel_config.context_parallel_config.ulysses_degree > 1: + if _parallel_config.context_parallel_config.ulysses_anything: + # For Any sequence lengths and Any head num support + return TemplatedUlyssesAnythingAttention.apply( + query, + key, + value, + attn_mask, + dropout_p, + is_causal, + scale, + enable_gqa, + return_lse, + forward_op, + backward_op, + _parallel_config, + ) + else: + return TemplatedUlyssesAttention.apply( + query, + key, + value, + attn_mask, + dropout_p, + is_causal, + scale, + enable_gqa, + return_lse, + forward_op, + backward_op, + _parallel_config, + ) else: raise ValueError("Reaching this branch of code is unexpected. Please report a bug.") @@ -1420,13 +2303,17 @@ def _flash_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, + attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: lse = None + if attn_mask is not None: + raise ValueError("`attn_mask` is not supported for flash-attn 2.") + if _parallel_config is None: out = flash_attn_func( q=query, @@ -1463,31 +2350,53 @@ def _flash_attention( @_AttentionBackendRegistry.register( AttentionBackendName.FLASH_HUB, constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape], - supports_context_parallel=False, + supports_context_parallel=True, ) def _flash_attention_hub( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, + attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: lse = None + if attn_mask is not None: + raise ValueError("`attn_mask` is not supported for flash-attn 2.") + func = _HUB_KERNELS_REGISTRY[AttentionBackendName.FLASH_HUB].kernel_fn - out = func( - q=query, - k=key, - v=value, - dropout_p=dropout_p, - softmax_scale=scale, - causal=is_causal, - return_attn_probs=return_lse, - ) - if return_lse: - out, lse, *_ = out + if _parallel_config is None: + out = func( + q=query, + k=key, + v=value, + dropout_p=dropout_p, + softmax_scale=scale, + causal=is_causal, + return_attn_probs=return_lse, + ) + if return_lse: + out, lse, *_ = out + else: + out = _templated_context_parallel_attention( + query, + key, + value, + None, + dropout_p, + is_causal, + scale, + False, + return_lse, + forward_op=_flash_attention_hub_forward_op, + backward_op=_flash_attention_hub_backward_op, + _parallel_config=_parallel_config, + ) + if return_lse: + out, lse = out return (out, lse) if return_lse else out @@ -1501,12 +2410,12 @@ def _flash_varlen_attention_hub( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[torch.Tensor] = None, + attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, - scale: Optional[float] = None, + scale: float | None = None, is_causal: bool = False, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: batch_size, seq_len_q, _, _ = query.shape _, seq_len_kv, _, _ = key.shape @@ -1557,12 +2466,12 @@ def _flash_varlen_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[torch.Tensor] = None, + attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, - scale: Optional[float] = None, + scale: float | None = None, is_causal: bool = False, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: batch_size, seq_len_q, _, _ = query.shape _, seq_len_kv, _, _ = key.shape @@ -1612,11 +2521,15 @@ def _flash_attention_3( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - scale: Optional[float] = None, + attn_mask: torch.Tensor | None = None, + scale: float | None = None, is_causal: bool = False, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: + if attn_mask is not None: + raise ValueError("`attn_mask` is not supported for flash-attn 3.") + out, lse = _wrapped_flash_attn_3( q=query, k=key, @@ -1630,45 +2543,83 @@ def _flash_attention_3( @_AttentionBackendRegistry.register( AttentionBackendName._FLASH_3_HUB, constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape], - supports_context_parallel=False, + supports_context_parallel=True, ) def _flash_attention_3_hub( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - scale: Optional[float] = None, + attn_mask: torch.Tensor | None = None, + scale: float | None = None, is_causal: bool = False, - window_size: Tuple[int, int] = (-1, -1), + window_size: tuple[int, int] = (-1, -1), softcap: float = 0.0, deterministic: bool = False, return_attn_probs: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: - if _parallel_config: - raise NotImplementedError(f"{AttentionBackendName._FLASH_3_HUB.value} is not implemented for parallelism yet.") + if attn_mask is not None: + raise ValueError("`attn_mask` is not supported for flash-attn 3.") func = _HUB_KERNELS_REGISTRY[AttentionBackendName._FLASH_3_HUB].kernel_fn - out = func( - q=query, - k=key, - v=value, - softmax_scale=scale, - causal=is_causal, - qv=None, - q_descale=None, - k_descale=None, - v_descale=None, + if _parallel_config is None: + out = func( + q=query, + k=key, + v=value, + softmax_scale=scale, + causal=is_causal, + qv=None, + q_descale=None, + k_descale=None, + v_descale=None, + window_size=window_size, + softcap=softcap, + num_splits=1, + pack_gqa=None, + deterministic=deterministic, + sm_margin=0, + return_attn_probs=return_attn_probs, + ) + return (out[0], out[1]) if return_attn_probs else out + + forward_op = functools.partial( + _flash_attention_3_hub_forward_op, + window_size=window_size, + softcap=softcap, + num_splits=1, + pack_gqa=None, + deterministic=deterministic, + sm_margin=0, + ) + backward_op = functools.partial( + _flash_attention_3_hub_backward_op, window_size=window_size, softcap=softcap, num_splits=1, pack_gqa=None, deterministic=deterministic, sm_margin=0, - return_attn_probs=return_attn_probs, ) - # When `return_attn_probs` is True, the above returns a tuple of - # actual outputs and lse. - return (out[0], out[1]) if return_attn_probs else out + out = _templated_context_parallel_attention( + query, + key, + value, + None, + 0.0, + is_causal, + scale, + False, + return_attn_probs, + forward_op=forward_op, + backward_op=backward_op, + _parallel_config=_parallel_config, + ) + if return_attn_probs: + out, lse = out + return out, lse + + return out @_AttentionBackendRegistry.register( @@ -1680,11 +2631,11 @@ def _flash_attention_3_varlen_hub( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[torch.Tensor] = None, - scale: Optional[float] = None, + attn_mask: torch.Tensor | None = None, + scale: float | None = None, is_causal: bool = False, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: batch_size, seq_len_q, _, _ = query.shape _, seq_len_kv, _, _ = key.shape @@ -1733,11 +2684,11 @@ def _flash_varlen_attention_3( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[torch.Tensor] = None, - scale: Optional[float] = None, + attn_mask: torch.Tensor | None = None, + scale: float | None = None, is_causal: bool = False, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: batch_size, seq_len_q, _, _ = query.shape _, seq_len_kv, _, _ = key.shape @@ -1761,7 +2712,7 @@ def _flash_varlen_attention_3( key_packed = torch.cat(key_valid, dim=0) value_packed = torch.cat(value_valid, dim=0) - out, lse, *_ = flash_attn_3_varlen_func( + result = flash_attn_3_varlen_func( q=query_packed, k=key_packed, v=value_packed, @@ -1771,7 +2722,13 @@ def _flash_varlen_attention_3( max_seqlen_k=max_seqlen_k, softmax_scale=scale, causal=is_causal, + return_attn_probs=return_lse, ) + if isinstance(result, tuple): + out, lse, *_ = result + else: + out = result + lse = None out = out.unflatten(0, (batch_size, -1)) return (out, lse) if return_lse else out @@ -1785,12 +2742,16 @@ def _aiter_flash_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, + attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: + if attn_mask is not None: + raise ValueError("`attn_mask` is not supported for aiter attention") + if not return_lse and torch.is_grad_enabled(): # aiter requires return_lse=True by assertion when gradients are enabled. out, lse, *_ = aiter_flash_attn_func( @@ -1826,12 +2787,12 @@ def _native_flex_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[Union[torch.Tensor, "flex_attention.BlockMask"]] = None, + attn_mask: torch.Tensor | "flex_attention.BlockMask" | None = None, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, enable_gqa: bool = False, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: # TODO: should we LRU cache the block mask creation? score_mod = None @@ -1881,6 +2842,43 @@ def score_mod(score, batch_idx, head_idx, q_idx, kv_idx): return out +def _prepare_additive_attn_mask( + attn_mask: torch.Tensor, target_dtype: torch.dtype, reshape_4d: bool = True +) -> torch.Tensor: + """ + Convert a 2D attention mask to an additive mask, optionally reshaping to 4D for SDPA. + + This helper is used by both native SDPA and xformers backends to handle both boolean and additive masks. + + Args: + attn_mask: 2D tensor [batch_size, seq_len_k] + - Boolean: True means attend, False means mask out + - Additive: 0.0 means attend, -inf means mask out + target_dtype: The dtype to convert the mask to (usually query.dtype) + reshape_4d: If True, reshape from [batch_size, seq_len_k] to [batch_size, 1, 1, seq_len_k] for broadcasting + + Returns: + Additive mask tensor where 0.0 means attend and -inf means mask out. Shape is [batch_size, seq_len_k] if + reshape_4d=False, or [batch_size, 1, 1, seq_len_k] if reshape_4d=True. + """ + # Check if the mask is boolean or already additive + if attn_mask.dtype == torch.bool: + # Convert boolean to additive: True -> 0.0, False -> -inf + attn_mask = torch.where(attn_mask, 0.0, float("-inf")) + # Convert to target dtype + attn_mask = attn_mask.to(dtype=target_dtype) + else: + # Already additive mask - just ensure correct dtype + attn_mask = attn_mask.to(dtype=target_dtype) + + # Optionally reshape to 4D for broadcasting in attention mechanisms + if reshape_4d: + batch_size, seq_len_k = attn_mask.shape + attn_mask = attn_mask.view(batch_size, 1, 1, seq_len_k) + + return attn_mask + + @_AttentionBackendRegistry.register( AttentionBackendName.NATIVE, constraints=[_check_device, _check_shape], @@ -1890,16 +2888,29 @@ def _native_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[torch.Tensor] = None, + attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, enable_gqa: bool = False, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: if return_lse: raise ValueError("Native attention backend does not support setting `return_lse=True`.") + + # Reshape 2D mask to 4D for SDPA + # SDPA accepts both boolean masks (torch.bool) and additive masks (float) + if ( + attn_mask is not None + and attn_mask.ndim == 2 + and attn_mask.shape[0] == query.shape[0] + and attn_mask.shape[1] == key.shape[1] + ): + # Just reshape [batch_size, seq_len_k] -> [batch_size, 1, 1, seq_len_k] + # SDPA handles both boolean and additive masks correctly + attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) + if _parallel_config is None: query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) out = torch.nn.functional.scaled_dot_product_attention( @@ -1941,13 +2952,13 @@ def _native_cudnn_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[torch.Tensor] = None, + attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, enable_gqa: bool = False, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: lse = None if _parallel_config is None and not return_lse: @@ -1993,13 +3004,13 @@ def _native_efficient_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[torch.Tensor] = None, + attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, enable_gqa: bool = False, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: if return_lse: raise ValueError("Native efficient attention backend does not support setting `return_lse=True`.") @@ -2028,13 +3039,17 @@ def _native_flash_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, + attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, enable_gqa: bool = False, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: + if attn_mask is not None: + raise ValueError("`attn_mask` is not supported for aiter attention") + lse = None if _parallel_config is None and not return_lse: query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) @@ -2079,13 +3094,13 @@ def _native_math_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[torch.Tensor] = None, + attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, enable_gqa: bool = False, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: if return_lse: raise ValueError("Native math attention backend does not support setting `return_lse=True`.") @@ -2108,34 +3123,53 @@ def _native_math_attention( @_AttentionBackendRegistry.register( AttentionBackendName._NATIVE_NPU, constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape], + supports_context_parallel=True, ) def _native_npu_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, + attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, - scale: Optional[float] = None, + scale: float | None = None, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: if return_lse: raise ValueError("NPU attention backend does not support setting `return_lse=True`.") - query, key, value = (x.transpose(1, 2).contiguous() for x in (query, key, value)) - out = npu_fusion_attention( - query, - key, - value, - query.size(1), # num_heads - input_layout="BNSD", - pse=None, - scale=1.0 / math.sqrt(query.shape[-1]) if scale is None else scale, - pre_tockens=65536, - next_tockens=65536, - keep_prob=1.0 - dropout_p, - sync=False, - inner_precise=0, - )[0] - out = out.transpose(1, 2).contiguous() + if _parallel_config is None: + attn_mask = _maybe_modify_attn_mask_npu(query, key, attn_mask) + + out = npu_fusion_attention( + query, + key, + value, + query.size(2), # num_heads + atten_mask=attn_mask, + input_layout="BSND", + pse=None, + scale=1.0 / math.sqrt(query.shape[-1]) if scale is None else scale, + pre_tockens=65536, + next_tockens=65536, + keep_prob=1.0 - dropout_p, + sync=False, + inner_precise=0, + )[0] + else: + out = _templated_context_parallel_attention( + query, + key, + value, + attn_mask, + dropout_p, + None, + scale, + None, + return_lse, + forward_op=_npu_attention_forward_op, + backward_op=_npu_attention_backward_op, + _parallel_config=_parallel_config, + ) return out @@ -2148,10 +3182,13 @@ def _native_xla_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, + attn_mask: torch.Tensor | None = None, is_causal: bool = False, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: + if attn_mask is not None: + raise ValueError("`attn_mask` is not supported for XLA attention") if return_lse: raise ValueError("XLA attention backend does not support setting `return_lse=True`.") query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) @@ -2175,11 +3212,14 @@ def _sage_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, + attn_mask: torch.Tensor | None = None, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: + if attn_mask is not None: + raise ValueError("`attn_mask` is not supported for sage attention") lse = None if _parallel_config is None: out = sageattn( @@ -2217,17 +3257,20 @@ def _sage_attention( @_AttentionBackendRegistry.register( AttentionBackendName.SAGE_HUB, constraints=[_check_device_cuda, _check_qkv_dtype_bf16_or_fp16, _check_shape], - supports_context_parallel=False, + supports_context_parallel=True, ) def _sage_attention_hub( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, + attn_mask: torch.Tensor | None = None, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: + if attn_mask is not None: + raise ValueError("`attn_mask` is not supported for sage attention") lse = None func = _HUB_KERNELS_REGISTRY[AttentionBackendName.SAGE_HUB].kernel_fn if _parallel_config is None: @@ -2242,6 +3285,23 @@ def _sage_attention_hub( ) if return_lse: out, lse, *_ = out + else: + out = _templated_context_parallel_attention( + query, + key, + value, + None, + 0.0, + is_causal, + scale, + False, + return_lse, + forward_op=_sage_attention_hub_forward_op, + backward_op=_sage_attention_backward_op, + _parallel_config=_parallel_config, + ) + if return_lse: + out, lse = out return (out, lse) if return_lse else out @@ -2254,11 +3314,11 @@ def _sage_varlen_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[torch.Tensor] = None, + attn_mask: torch.Tensor | None = None, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: if return_lse: raise ValueError("Sage varlen backend does not support setting `return_lse=True`.") @@ -2309,11 +3369,14 @@ def _sage_qk_int8_pv_fp8_cuda_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, + attn_mask: torch.Tensor | None = None, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: + if attn_mask is not None: + raise ValueError("`attn_mask` is not supported for sage attention") return sageattn_qk_int8_pv_fp8_cuda( q=query, k=key, @@ -2333,11 +3396,14 @@ def _sage_qk_int8_pv_fp8_cuda_sm90_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, + attn_mask: torch.Tensor | None = None, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: + if attn_mask is not None: + raise ValueError("`attn_mask` is not supported for sage attention") return sageattn_qk_int8_pv_fp8_cuda_sm90( q=query, k=key, @@ -2357,11 +3423,14 @@ def _sage_qk_int8_pv_fp16_cuda_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, + attn_mask: torch.Tensor | None = None, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: + if attn_mask is not None: + raise ValueError("`attn_mask` is not supported for sage attention") return sageattn_qk_int8_pv_fp16_cuda( q=query, k=key, @@ -2381,11 +3450,14 @@ def _sage_qk_int8_pv_fp16_triton_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, + attn_mask: torch.Tensor | None = None, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: + if attn_mask is not None: + raise ValueError("`attn_mask` is not supported for sage attention") return sageattn_qk_int8_pv_fp16_triton( q=query, k=key, @@ -2405,13 +3477,13 @@ def _xformers_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_mask: Optional[torch.Tensor] = None, + attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, - scale: Optional[float] = None, + scale: float | None = None, enable_gqa: bool = False, return_lse: bool = False, - _parallel_config: Optional["ParallelConfig"] = None, + _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: if return_lse: raise ValueError("xformers attention backend does not support setting `return_lse=True`.") @@ -2423,10 +3495,34 @@ def _xformers_attention( attn_mask = xops.LowerTriangularMask() elif attn_mask is not None: if attn_mask.ndim == 2: - attn_mask = attn_mask.view(attn_mask.size(0), 1, attn_mask.size(1), 1) + # Convert 2D mask to 4D for xformers + # Mask can be boolean (True=attend, False=mask) or additive (0.0=attend, -inf=mask) + # xformers requires 4D additive masks [batch, heads, seq_q, seq_k] + # Need memory alignment - create larger tensor and slice for alignment + original_seq_len = attn_mask.size(1) + aligned_seq_len = ((original_seq_len + 7) // 8) * 8 # Round up to multiple of 8 + + # Create aligned 4D tensor and slice to ensure proper memory layout + aligned_mask = torch.zeros( + (batch_size, num_heads_q, seq_len_q, aligned_seq_len), + dtype=query.dtype, + device=query.device, + ) + # Convert to 4D additive mask (handles both boolean and additive inputs) + mask_additive = _prepare_additive_attn_mask( + attn_mask, target_dtype=query.dtype + ) # [batch, 1, 1, seq_len_k] + # Broadcast to [batch, heads, seq_q, seq_len_k] + aligned_mask[:, :, :, :original_seq_len] = mask_additive + # Mask out the padding (already -inf from zeros -> where with default) + aligned_mask[:, :, :, original_seq_len:] = float("-inf") + + # Slice to actual size with proper alignment + attn_mask = aligned_mask[:, :, :, :seq_len_kv] elif attn_mask.ndim != 4: raise ValueError("Only 2D and 4D attention masks are supported for xformers attention.") - attn_mask = attn_mask.expand(batch_size, num_heads_q, seq_len_q, seq_len_kv).type_as(query) + elif attn_mask.ndim == 4: + attn_mask = attn_mask.expand(batch_size, num_heads_q, seq_len_q, seq_len_kv).type_as(query) if enable_gqa: if num_heads_q % num_heads_kv != 0: diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py index 66455d733aee..f83a1753dac5 100755 --- a/src/diffusers/models/attention_processor.py +++ b/src/diffusers/models/attention_processor.py @@ -11,9 +11,11 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import inspect import math -from typing import Callable, List, Optional, Tuple, Union +from typing import Callable import torch import torch.nn.functional as F @@ -105,21 +107,21 @@ class Attention(nn.Module): def __init__( self, query_dim: int, - cross_attention_dim: Optional[int] = None, + cross_attention_dim: int | None = None, heads: int = 8, - kv_heads: Optional[int] = None, + kv_heads: int | None = None, dim_head: int = 64, dropout: float = 0.0, bias: bool = False, upcast_attention: bool = False, upcast_softmax: bool = False, - cross_attention_norm: Optional[str] = None, + cross_attention_norm: str | None = None, cross_attention_norm_num_groups: int = 32, - qk_norm: Optional[str] = None, - added_kv_proj_dim: Optional[int] = None, - added_proj_bias: Optional[bool] = True, - norm_num_groups: Optional[int] = None, - spatial_norm_dim: Optional[int] = None, + qk_norm: str | None = None, + added_kv_proj_dim: int | None = None, + added_proj_bias: bool | None = True, + norm_num_groups: int | None = None, + spatial_norm_dim: int | None = None, out_bias: bool = True, scale_qk: bool = True, only_cross_attention: bool = False, @@ -127,7 +129,7 @@ def __init__( rescale_output_factor: float = 1.0, residual_connection: bool = False, _from_deprecated_attn_block: bool = False, - processor: Optional["AttnProcessor"] = None, + processor: "AttnProcessor" | None = None, out_dim: int = None, out_context_dim: int = None, context_pre_only=None, @@ -309,7 +311,7 @@ def __init__( def set_use_xla_flash_attention( self, use_xla_flash_attention: bool, - partition_spec: Optional[Tuple[Optional[str], ...]] = None, + partition_spec: tuple[str | None, ...] | None = None, is_flux=False, ) -> None: r""" @@ -318,7 +320,7 @@ def set_use_xla_flash_attention( Args: use_xla_flash_attention (`bool`): Whether to use pallas flash attention kernel from `torch_xla` or not. - partition_spec (`Tuple[]`, *optional*): + partition_spec (`tuple[]`, *optional*): Specify the partition specification if using SPMD. Otherwise None. """ if use_xla_flash_attention: @@ -357,7 +359,7 @@ def set_use_npu_flash_attention(self, use_npu_flash_attention: bool) -> None: self.set_processor(processor) def set_use_memory_efficient_attention_xformers( - self, use_memory_efficient_attention_xformers: bool, attention_op: Optional[Callable] = None + self, use_memory_efficient_attention_xformers: bool, attention_op: Callable | None = None ) -> None: r""" Set whether to use memory efficient attention from `xformers` or not. @@ -567,8 +569,8 @@ def get_processor(self, return_deprecated_lora: bool = False) -> "AttentionProce def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, **cross_attention_kwargs, ) -> torch.Tensor: r""" @@ -655,7 +657,7 @@ def head_to_batch_dim(self, tensor: torch.Tensor, out_dim: int = 3) -> torch.Ten return tensor def get_attention_scores( - self, query: torch.Tensor, key: torch.Tensor, attention_mask: Optional[torch.Tensor] = None + self, query: torch.Tensor, key: torch.Tensor, attention_mask: torch.Tensor | None = None ) -> torch.Tensor: r""" Compute the attention scores. @@ -868,11 +870,11 @@ def __init__( self, in_channels: int, out_channels: int, - num_attention_heads: Optional[int] = None, + num_attention_heads: int | None = None, attention_head_dim: int = 8, mult: float = 1.0, norm_type: str = "batch_norm", - kernel_sizes: Tuple[int, ...] = (5,), + kernel_sizes: tuple[int, ...] = (5,), eps: float = 1e-15, residual_connection: bool = False, ): @@ -938,8 +940,8 @@ def __init__( dropout: float = 0.0, bias: bool = False, added_proj_bias: bool = True, - out_dim: Optional[int] = None, - out_context_dim: Optional[int] = None, + out_dim: int | None = None, + out_context_dim: int | None = None, out_bias: bool = True, context_pre_only: bool = False, eps: float = 1e-5, @@ -980,8 +982,8 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, **kwargs, ): return self.processor( @@ -1006,7 +1008,7 @@ def __call__( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, attention_mask: torch.Tensor, - image_rotary_emb: Optional[torch.Tensor] = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: query = attn.to_q(hidden_states) key = attn.to_k(hidden_states) @@ -1107,9 +1109,9 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - temb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + temb: torch.Tensor | None = None, *args, **kwargs, ) -> torch.Tensor: @@ -1193,8 +1195,8 @@ def __init__( self, train_kv: bool = True, train_q_out: bool = True, - hidden_size: Optional[int] = None, - cross_attention_dim: Optional[int] = None, + hidden_size: int | None = None, + cross_attention_dim: int | None = None, out_bias: bool = True, dropout: float = 0.0, ): @@ -1219,8 +1221,8 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: batch_size, sequence_length, _ = hidden_states.shape attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) @@ -1284,8 +1286,8 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, *args, **kwargs, ) -> torch.Tensor: @@ -1357,8 +1359,8 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, *args, **kwargs, ) -> torch.Tensor: @@ -1429,7 +1431,7 @@ def __call__( attn: Attention, hidden_states: torch.FloatTensor, encoder_hidden_states: torch.FloatTensor = None, - attention_mask: Optional[torch.FloatTensor] = None, + attention_mask: torch.FloatTensor | None = None, *args, **kwargs, ) -> torch.FloatTensor: @@ -1517,7 +1519,7 @@ def __call__( attn: Attention, hidden_states: torch.FloatTensor, encoder_hidden_states: torch.FloatTensor = None, - attention_mask: Optional[torch.FloatTensor] = None, + attention_mask: torch.FloatTensor | None = None, ) -> torch.FloatTensor: residual = hidden_states @@ -1673,7 +1675,7 @@ def __call__( attn: Attention, hidden_states: torch.FloatTensor, encoder_hidden_states: torch.FloatTensor = None, - attention_mask: Optional[torch.FloatTensor] = None, + attention_mask: torch.FloatTensor | None = None, *args, **kwargs, ) -> torch.FloatTensor: @@ -1836,7 +1838,7 @@ def __call__( attn: Attention, hidden_states: torch.FloatTensor, encoder_hidden_states: torch.FloatTensor = None, - attention_mask: Optional[torch.FloatTensor] = None, + attention_mask: torch.FloatTensor | None = None, *args, **kwargs, ) -> torch.FloatTensor: @@ -1915,7 +1917,7 @@ class XFormersJointAttnProcessor: operator. """ - def __init__(self, attention_op: Optional[Callable] = None): + def __init__(self, attention_op: Callable | None = None): self.attention_op = attention_op def __call__( @@ -1923,7 +1925,7 @@ def __call__( attn: Attention, hidden_states: torch.FloatTensor, encoder_hidden_states: torch.FloatTensor = None, - attention_mask: Optional[torch.FloatTensor] = None, + attention_mask: torch.FloatTensor | None = None, *args, **kwargs, ) -> torch.FloatTensor: @@ -2004,10 +2006,10 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - temb: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + temb: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: residual = hidden_states @@ -2287,8 +2289,8 @@ def __call__( attn: Attention, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: text_seq_length = encoder_hidden_states.size(1) @@ -2356,8 +2358,8 @@ def __call__( attn: Attention, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: text_seq_length = encoder_hidden_states.size(1) @@ -2424,15 +2426,15 @@ class XFormersAttnAddedKVProcessor: operator. """ - def __init__(self, attention_op: Optional[Callable] = None): + def __init__(self, attention_op: Callable | None = None): self.attention_op = attention_op def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: residual = hidden_states hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2) @@ -2495,16 +2497,16 @@ class XFormersAttnProcessor: operator. """ - def __init__(self, attention_op: Optional[Callable] = None): + def __init__(self, attention_op: Callable | None = None): self.attention_op = attention_op def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - temb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + temb: torch.Tensor | None = None, *args, **kwargs, ) -> torch.Tensor: @@ -2593,9 +2595,9 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - temb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + temb: torch.Tensor | None = None, *args, **kwargs, ) -> torch.Tensor: @@ -2704,9 +2706,9 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - temb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + temb: torch.Tensor | None = None, *args, **kwargs, ) -> torch.Tensor: @@ -2790,7 +2792,7 @@ class XLAFlashAttnProcessor2_0: Processor for implementing scaled dot-product attention with pallas flash attention kernel if using `torch_xla`. """ - def __init__(self, partition_spec: Optional[Tuple[Optional[str], ...]] = None): + def __init__(self, partition_spec: tuple[str | None, ...] | None = None): if not hasattr(F, "scaled_dot_product_attention"): raise ImportError( "XLAFlashAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0." @@ -2805,9 +2807,9 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - temb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + temb: torch.Tensor | None = None, *args, **kwargs, ) -> torch.Tensor: @@ -2914,8 +2916,8 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: residual = hidden_states is_single_frame = hidden_states.shape[1] == 1 @@ -3001,7 +3003,7 @@ def __init__(self): def apply_partial_rotary_emb( self, x: torch.Tensor, - freqs_cis: Tuple[torch.Tensor], + freqs_cis: tuple[torch.Tensor], ) -> torch.Tensor: from .embeddings import apply_rotary_emb @@ -3017,9 +3019,9 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: from .embeddings import apply_rotary_emb @@ -3133,10 +3135,10 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - temb: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + temb: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: from .embeddings import apply_rotary_emb @@ -3234,10 +3236,10 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - temb: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + temb: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: from .embeddings import apply_rotary_emb @@ -3337,10 +3339,10 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - temb: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + temb: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: from .embeddings import apply_rotary_emb @@ -3460,10 +3462,10 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - temb: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + temb: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: from .embeddings import apply_rotary_emb @@ -3582,10 +3584,10 @@ def __call__( attn: Attention, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - query_rotary_emb: Optional[torch.Tensor] = None, - key_rotary_emb: Optional[torch.Tensor] = None, - base_sequence_length: Optional[int] = None, + attention_mask: torch.Tensor | None = None, + query_rotary_emb: torch.Tensor | None = None, + key_rotary_emb: torch.Tensor | None = None, + base_sequence_length: int | None = None, ) -> torch.Tensor: from .embeddings import apply_rotary_emb @@ -3682,9 +3684,9 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - temb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + temb: torch.Tensor | None = None, *args, **kwargs, ) -> torch.Tensor: @@ -3792,11 +3794,11 @@ def __init__( self, train_kv: bool = True, train_q_out: bool = False, - hidden_size: Optional[int] = None, - cross_attention_dim: Optional[int] = None, + hidden_size: int | None = None, + cross_attention_dim: int | None = None, out_bias: bool = True, dropout: float = 0.0, - attention_op: Optional[Callable] = None, + attention_op: Callable | None = None, ): super().__init__() self.train_kv = train_kv @@ -3820,8 +3822,8 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: batch_size, sequence_length, _ = ( hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape @@ -3905,8 +3907,8 @@ def __init__( self, train_kv: bool = True, train_q_out: bool = True, - hidden_size: Optional[int] = None, - cross_attention_dim: Optional[int] = None, + hidden_size: int | None = None, + cross_attention_dim: int | None = None, out_bias: bool = True, dropout: float = 0.0, ): @@ -3931,8 +3933,8 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: batch_size, sequence_length, _ = hidden_states.shape attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) @@ -4012,8 +4014,8 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: residual = hidden_states @@ -4099,9 +4101,9 @@ def __call__( self, attn: "Attention", hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - temb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + temb: torch.Tensor | None = None, ) -> torch.Tensor: residual = hidden_states @@ -4212,9 +4214,9 @@ class IPAdapterAttnProcessor(nn.Module): The hidden size of the attention layer. cross_attention_dim (`int`): The number of channels in the `encoder_hidden_states`. - num_tokens (`int`, `Tuple[int]` or `List[int]`, defaults to `(4,)`): + num_tokens (`int`, `tuple[int]` or `list[int]`, defaults to `(4,)`): The context length of the image features. - scale (`float` or List[`float`], defaults to 1.0): + scale (`float` or list[`float`], defaults to 1.0): the weight scale of image prompt. """ @@ -4245,11 +4247,11 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - temb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + temb: torch.Tensor | None = None, scale: float = 1.0, - ip_adapter_masks: Optional[torch.Tensor] = None, + ip_adapter_masks: torch.Tensor | None = None, ): residual = hidden_states @@ -4305,7 +4307,7 @@ def __call__( hidden_states = attn.batch_to_head_dim(hidden_states) if ip_adapter_masks is not None: - if not isinstance(ip_adapter_masks, List): + if not isinstance(ip_adapter_masks, list): # for backward compatibility, we accept `ip_adapter_mask` as a tensor of shape [num_ip_adapter, 1, height, width] ip_adapter_masks = list(ip_adapter_masks.unsqueeze(1)) if not (len(ip_adapter_masks) == len(self.scale) == len(ip_hidden_states)): @@ -4412,9 +4414,9 @@ class IPAdapterAttnProcessor2_0(torch.nn.Module): The hidden size of the attention layer. cross_attention_dim (`int`): The number of channels in the `encoder_hidden_states`. - num_tokens (`int`, `Tuple[int]` or `List[int]`, defaults to `(4,)`): + num_tokens (`int`, `tuple[int]` or `list[int]`, defaults to `(4,)`): The context length of the image features. - scale (`float` or `List[float]`, defaults to 1.0): + scale (`float` or `list[float]`, defaults to 1.0): the weight scale of image prompt. """ @@ -4450,11 +4452,11 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - temb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + temb: torch.Tensor | None = None, scale: float = 1.0, - ip_adapter_masks: Optional[torch.Tensor] = None, + ip_adapter_masks: torch.Tensor | None = None, ): residual = hidden_states @@ -4524,7 +4526,7 @@ def __call__( hidden_states = hidden_states.to(query.dtype) if ip_adapter_masks is not None: - if not isinstance(ip_adapter_masks, List): + if not isinstance(ip_adapter_masks, list): # for backward compatibility, we accept `ip_adapter_mask` as a tensor of shape [num_ip_adapter, 1, height, width] ip_adapter_masks = list(ip_adapter_masks.unsqueeze(1)) if not (len(ip_adapter_masks) == len(self.scale) == len(ip_hidden_states)): @@ -4644,9 +4646,9 @@ class IPAdapterXFormersAttnProcessor(torch.nn.Module): The hidden size of the attention layer. cross_attention_dim (`int`): The number of channels in the `encoder_hidden_states`. - num_tokens (`int`, `Tuple[int]` or `List[int]`, defaults to `(4,)`): + num_tokens (`int`, `tuple[int]` or `list[int]`, defaults to `(4,)`): The context length of the image features. - scale (`float` or `List[float]`, defaults to 1.0): + scale (`float` or `list[float]`, defaults to 1.0): the weight scale of image prompt. attention_op (`Callable`, *optional*, defaults to `None`): The base @@ -4661,7 +4663,7 @@ def __init__( cross_attention_dim=None, num_tokens=(4,), scale=1.0, - attention_op: Optional[Callable] = None, + attention_op: Callable | None = None, ): super().__init__() @@ -4690,11 +4692,11 @@ def __call__( self, attn: Attention, hidden_states: torch.FloatTensor, - encoder_hidden_states: Optional[torch.FloatTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: torch.FloatTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + temb: torch.FloatTensor | None = None, scale: float = 1.0, - ip_adapter_masks: Optional[torch.FloatTensor] = None, + ip_adapter_masks: torch.FloatTensor | None = None, ): residual = hidden_states @@ -4763,7 +4765,7 @@ def __call__( if ip_hidden_states: if ip_adapter_masks is not None: - if not isinstance(ip_adapter_masks, List): + if not isinstance(ip_adapter_masks, list): # for backward compatibility, we accept `ip_adapter_mask` as a tensor of shape [num_ip_adapter, 1, height, width] ip_adapter_masks = list(ip_adapter_masks.unsqueeze(1)) if not (len(ip_adapter_masks) == len(self.scale) == len(ip_hidden_states)): @@ -4911,7 +4913,7 @@ def __call__( attn: Attention, hidden_states: torch.FloatTensor, encoder_hidden_states: torch.FloatTensor = None, - attention_mask: Optional[torch.FloatTensor] = None, + attention_mask: torch.FloatTensor | None = None, ip_hidden_states: torch.FloatTensor = None, temb: torch.FloatTensor = None, ) -> torch.FloatTensor: @@ -5054,9 +5056,9 @@ def __call__( self, attn: Attention, hidden_states: torch.FloatTensor, - encoder_hidden_states: Optional[torch.FloatTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: torch.FloatTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + temb: torch.FloatTensor | None = None, ) -> torch.Tensor: residual = hidden_states if attn.spatial_norm is not None: @@ -5153,9 +5155,9 @@ def __call__( self, attn: Attention, hidden_states: torch.FloatTensor, - encoder_hidden_states: Optional[torch.FloatTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: torch.FloatTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + temb: torch.FloatTensor | None = None, ) -> torch.Tensor: residual = hidden_states if attn.spatial_norm is not None: @@ -5345,8 +5347,8 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: original_dtype = hidden_states.dtype @@ -5397,8 +5399,8 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: original_dtype = hidden_states.dtype @@ -5452,8 +5454,8 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: original_dtype = hidden_states.dtype @@ -5622,56 +5624,56 @@ def __new__(cls, *args, **kwargs): FluxIPAdapterJointAttnProcessor2_0, ) -AttentionProcessor = Union[ - AttnProcessor, - CustomDiffusionAttnProcessor, - AttnAddedKVProcessor, - AttnAddedKVProcessor2_0, - JointAttnProcessor2_0, - PAGJointAttnProcessor2_0, - PAGCFGJointAttnProcessor2_0, - FusedJointAttnProcessor2_0, - AllegroAttnProcessor2_0, - AuraFlowAttnProcessor2_0, - FusedAuraFlowAttnProcessor2_0, - FluxAttnProcessor2_0, - FluxAttnProcessor2_0_NPU, - FusedFluxAttnProcessor2_0, - FusedFluxAttnProcessor2_0_NPU, - CogVideoXAttnProcessor2_0, - FusedCogVideoXAttnProcessor2_0, - XFormersAttnAddedKVProcessor, - XFormersAttnProcessor, - XLAFlashAttnProcessor2_0, - AttnProcessorNPU, - AttnProcessor2_0, - MochiVaeAttnProcessor2_0, - MochiAttnProcessor2_0, - StableAudioAttnProcessor2_0, - HunyuanAttnProcessor2_0, - FusedHunyuanAttnProcessor2_0, - PAGHunyuanAttnProcessor2_0, - PAGCFGHunyuanAttnProcessor2_0, - LuminaAttnProcessor2_0, - FusedAttnProcessor2_0, - CustomDiffusionXFormersAttnProcessor, - CustomDiffusionAttnProcessor2_0, - SlicedAttnProcessor, - SlicedAttnAddedKVProcessor, - SanaLinearAttnProcessor2_0, - PAGCFGSanaLinearAttnProcessor2_0, - PAGIdentitySanaLinearAttnProcessor2_0, - SanaMultiscaleLinearAttention, - SanaMultiscaleAttnProcessor2_0, - SanaMultiscaleAttentionProjection, - IPAdapterAttnProcessor, - IPAdapterAttnProcessor2_0, - IPAdapterXFormersAttnProcessor, - SD3IPAdapterJointAttnProcessor2_0, - PAGIdentitySelfAttnProcessor2_0, - PAGCFGIdentitySelfAttnProcessor2_0, - LoRAAttnProcessor, - LoRAAttnProcessor2_0, - LoRAXFormersAttnProcessor, - LoRAAttnAddedKVProcessor, -] +AttentionProcessor = ( + AttnProcessor + | CustomDiffusionAttnProcessor + | AttnAddedKVProcessor + | AttnAddedKVProcessor2_0 + | JointAttnProcessor2_0 + | PAGJointAttnProcessor2_0 + | PAGCFGJointAttnProcessor2_0 + | FusedJointAttnProcessor2_0 + | AllegroAttnProcessor2_0 + | AuraFlowAttnProcessor2_0 + | FusedAuraFlowAttnProcessor2_0 + | FluxAttnProcessor2_0 + | FluxAttnProcessor2_0_NPU + | FusedFluxAttnProcessor2_0 + | FusedFluxAttnProcessor2_0_NPU + | CogVideoXAttnProcessor2_0 + | FusedCogVideoXAttnProcessor2_0 + | XFormersAttnAddedKVProcessor + | XFormersAttnProcessor + | XLAFlashAttnProcessor2_0 + | AttnProcessorNPU + | AttnProcessor2_0 + | MochiVaeAttnProcessor2_0 + | MochiAttnProcessor2_0 + | StableAudioAttnProcessor2_0 + | HunyuanAttnProcessor2_0 + | FusedHunyuanAttnProcessor2_0 + | PAGHunyuanAttnProcessor2_0 + | PAGCFGHunyuanAttnProcessor2_0 + | LuminaAttnProcessor2_0 + | FusedAttnProcessor2_0 + | CustomDiffusionXFormersAttnProcessor + | CustomDiffusionAttnProcessor2_0 + | SlicedAttnProcessor + | SlicedAttnAddedKVProcessor + | SanaLinearAttnProcessor2_0 + | PAGCFGSanaLinearAttnProcessor2_0 + | PAGIdentitySanaLinearAttnProcessor2_0 + | SanaMultiscaleLinearAttention + | SanaMultiscaleAttnProcessor2_0 + | SanaMultiscaleAttentionProjection + | IPAdapterAttnProcessor + | IPAdapterAttnProcessor2_0 + | IPAdapterXFormersAttnProcessor + | SD3IPAdapterJointAttnProcessor2_0 + | PAGIdentitySelfAttnProcessor2_0 + | PAGCFGIdentitySelfAttnProcessor2_0 + | LoRAAttnProcessor + | LoRAAttnProcessor2_0 + | LoRAXFormersAttnProcessor + | LoRAAttnAddedKVProcessor +) diff --git a/src/diffusers/models/auto_model.py b/src/diffusers/models/auto_model.py index c96b4fa88c49..1c001e23fe00 100644 --- a/src/diffusers/models/auto_model.py +++ b/src/diffusers/models/auto_model.py @@ -13,12 +13,11 @@ # limitations under the License. import os -from typing import Optional, Union from huggingface_hub.utils import validate_hf_hub_args from ..configuration_utils import ConfigMixin -from ..utils import logging +from ..utils import DIFFUSERS_LOAD_ID_FIELDS, logging from ..utils.dynamic_modules_utils import get_class_from_dynamic_module, resolve_trust_remote_code @@ -31,13 +30,129 @@ class AutoModel(ConfigMixin): def __init__(self, *args, **kwargs): raise EnvironmentError( f"{self.__class__.__name__} is designed to be instantiated " - f"using the `{self.__class__.__name__}.from_pretrained(pretrained_model_name_or_path)` or " + f"using the `{self.__class__.__name__}.from_pretrained(pretrained_model_name_or_path)`, " + f"`{self.__class__.__name__}.from_config(config)`, or " f"`{self.__class__.__name__}.from_pipe(pipeline)` methods." ) + @classmethod + def from_config(cls, pretrained_model_name_or_path_or_dict: str | os.PathLike | dict | None = None, **kwargs): + r""" + Instantiate a model from a config dictionary or a pretrained model configuration file with random weights (no + pretrained weights are loaded). + + Parameters: + pretrained_model_name_or_path_or_dict (`str`, `os.PathLike`, or `dict`): + Can be either: + + - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model + configuration hosted on the Hub. + - A path to a *directory* (for example `./my_model_directory`) containing a model configuration + file. + - A config dictionary. + + cache_dir (`Union[str, os.PathLike]`, *optional*): + Path to a directory where a downloaded pretrained model configuration is cached if the standard cache + is not used. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force the (re-)download of the model configuration, overriding the cached version if + it exists. + proxies (`Dict[str, str]`, *optional*): + A dictionary of proxy servers to use by protocol or endpoint. + local_files_only(`bool`, *optional*, defaults to `False`): + Whether to only load local model configuration files or not. + token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. + trust_remote_code (`bool`, *optional*, defaults to `False`): + Whether to trust remote code. + subfolder (`str`, *optional*, defaults to `""`): + The subfolder location of a model file within a larger model repository on the Hub or locally. + + Returns: + A model object instantiated from the config with random weights. + + Example: + + ```py + from diffusers import AutoModel + + model = AutoModel.from_config("stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="unet") + ``` + """ + subfolder = kwargs.pop("subfolder", None) + trust_remote_code = kwargs.pop("trust_remote_code", False) + + hub_kwargs_names = [ + "cache_dir", + "force_download", + "local_files_only", + "proxies", + "revision", + "token", + ] + hub_kwargs = {name: kwargs.pop(name, None) for name in hub_kwargs_names} + + if pretrained_model_name_or_path_or_dict is None: + raise ValueError( + "Please provide a `pretrained_model_name_or_path_or_dict` as the first positional argument." + ) + + if isinstance(pretrained_model_name_or_path_or_dict, (str, os.PathLike)): + pretrained_model_name_or_path = pretrained_model_name_or_path_or_dict + config = cls.load_config(pretrained_model_name_or_path, subfolder=subfolder, **hub_kwargs) + else: + config = pretrained_model_name_or_path_or_dict + pretrained_model_name_or_path = config.get("_name_or_path", None) + + has_remote_code = "auto_map" in config and cls.__name__ in config["auto_map"] + trust_remote_code = resolve_trust_remote_code( + trust_remote_code, pretrained_model_name_or_path, has_remote_code + ) + + if has_remote_code and trust_remote_code: + class_ref = config["auto_map"][cls.__name__] + module_file, class_name = class_ref.split(".") + module_file = module_file + ".py" + model_cls = get_class_from_dynamic_module( + pretrained_model_name_or_path, + subfolder=subfolder, + module_file=module_file, + class_name=class_name, + **hub_kwargs, + ) + else: + if "_class_name" in config: + class_name = config["_class_name"] + library = "diffusers" + elif "model_type" in config: + class_name = "AutoModel" + library = "transformers" + else: + raise ValueError( + f"Couldn't find a model class associated with the config: {config}. Make sure the config " + "contains a `_class_name` or `model_type` key." + ) + + from ..pipelines.pipeline_loading_utils import ALL_IMPORTABLE_CLASSES, get_class_obj_and_candidates + + model_cls, _ = get_class_obj_and_candidates( + library_name=library, + class_name=class_name, + importable_classes=ALL_IMPORTABLE_CLASSES, + pipelines=None, + is_pipeline_module=False, + ) + + if model_cls is None: + raise ValueError(f"AutoModel can't find a model linked to {class_name}.") + + return model_cls.from_config(config, **kwargs) + @classmethod @validate_hf_hub_args - def from_pretrained(cls, pretrained_model_or_path: Optional[Union[str, os.PathLike]] = None, **kwargs): + def from_pretrained(cls, pretrained_model_or_path: str | os.PathLike | None = None, **kwargs): r""" Instantiate a pretrained PyTorch model from a pretrained model configuration. @@ -53,7 +168,7 @@ def from_pretrained(cls, pretrained_model_or_path: Optional[Union[str, os.PathLi - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved with [`~ModelMixin.save_pretrained`]. - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. torch_dtype (`torch.dtype`, *optional*): @@ -61,7 +176,7 @@ def from_pretrained(cls, pretrained_model_or_path: Optional[Union[str, os.PathLi force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. output_loading_info (`bool`, *optional*, defaults to `False`): @@ -83,7 +198,7 @@ def from_pretrained(cls, pretrained_model_or_path: Optional[Union[str, os.PathLi Mirror source to resolve accessibility issues if you're downloading a model in China. We do not guarantee the timeliness or safety of the source, and you should refer to the mirror site for more information. - device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*): + device_map (`str` or `dict[str, int | str | torch.device]`, *optional*): A map that specifies where each submodule should go. It doesn't need to be defined for each parameter/buffer name; once a given module name is inside, every submodule of it will be sent to the same device. Defaults to `None`, meaning that the model will be loaded on CPU. @@ -220,4 +335,11 @@ def from_pretrained(cls, pretrained_model_or_path: Optional[Union[str, os.PathLi raise ValueError(f"AutoModel can't find a model linked to {orig_class_name}.") kwargs = {**load_config_kwargs, **kwargs} - return model_cls.from_pretrained(pretrained_model_or_path, **kwargs) + model = model_cls.from_pretrained(pretrained_model_or_path, **kwargs) + + load_id_kwargs = {"pretrained_model_name_or_path": pretrained_model_or_path, **kwargs} + parts = [load_id_kwargs.get(field, "null") for field in DIFFUSERS_LOAD_ID_FIELDS] + load_id = "|".join("null" if p is None else p for p in parts) + model._diffusers_load_id = load_id + + return model diff --git a/src/diffusers/models/autoencoders/__init__.py b/src/diffusers/models/autoencoders/__init__.py index 56df27f93cd7..23665ee0532e 100644 --- a/src/diffusers/models/autoencoders/__init__.py +++ b/src/diffusers/models/autoencoders/__init__.py @@ -10,12 +10,15 @@ from .autoencoder_kl_hunyuanimage_refiner import AutoencoderKLHunyuanImageRefiner from .autoencoder_kl_hunyuanvideo15 import AutoencoderKLHunyuanVideo15 from .autoencoder_kl_ltx import AutoencoderKLLTXVideo +from .autoencoder_kl_ltx2 import AutoencoderKLLTX2Video +from .autoencoder_kl_ltx2_audio import AutoencoderKLLTX2Audio from .autoencoder_kl_magvit import AutoencoderKLMagvit from .autoencoder_kl_mochi import AutoencoderKLMochi from .autoencoder_kl_qwenimage import AutoencoderKLQwenImage from .autoencoder_kl_temporal_decoder import AutoencoderKLTemporalDecoder from .autoencoder_kl_wan import AutoencoderKLWan from .autoencoder_oobleck import AutoencoderOobleck +from .autoencoder_rae import AutoencoderRAE from .autoencoder_tiny import AutoencoderTiny from .consistency_decoder_vae import ConsistencyDecoderVAE from .vq_model import VQModel diff --git a/src/diffusers/models/autoencoders/autoencoder_asym_kl.py b/src/diffusers/models/autoencoders/autoencoder_asym_kl.py index fa49fcfe79f8..fbd9b3e459f7 100644 --- a/src/diffusers/models/autoencoders/autoencoder_asym_kl.py +++ b/src/diffusers/models/autoencoders/autoencoder_asym_kl.py @@ -11,8 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Tuple, Union - import torch import torch.nn as nn @@ -34,16 +32,16 @@ class AsymmetricAutoencoderKL(ModelMixin, AutoencoderMixin, ConfigMixin): Parameters: in_channels (int, *optional*, defaults to 3): Number of channels in the input image. out_channels (int, *optional*, defaults to 3): Number of channels in the output. - down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`): - Tuple of downsample block types. - down_block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`): - Tuple of down block output channels. + down_block_types (`tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`): + tuple of downsample block types. + down_block_out_channels (`tuple[int]`, *optional*, defaults to `(64,)`): + tuple of down block output channels. layers_per_down_block (`int`, *optional*, defaults to `1`): Number layers for down block. - up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`): - Tuple of upsample block types. - up_block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`): - Tuple of up block output channels. + up_block_types (`tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`): + tuple of upsample block types. + up_block_out_channels (`tuple[int]`, *optional*, defaults to `(64,)`): + tuple of up block output channels. layers_per_up_block (`int`, *optional*, defaults to `1`): Number layers for up block. act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. @@ -67,11 +65,11 @@ def __init__( self, in_channels: int = 3, out_channels: int = 3, - down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",), - down_block_out_channels: Tuple[int, ...] = (64,), + down_block_types: tuple[str, ...] = ("DownEncoderBlock2D",), + down_block_out_channels: tuple[int, ...] = (64,), layers_per_down_block: int = 1, - up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",), - up_block_out_channels: Tuple[int, ...] = (64,), + up_block_types: tuple[str, ...] = ("UpDecoderBlock2D",), + up_block_out_channels: tuple[int, ...] = (64,), layers_per_up_block: int = 1, act_fn: str = "silu", latent_channels: int = 4, @@ -111,7 +109,7 @@ def __init__( self.register_to_config(force_upcast=False) @apply_forward_hook - def encode(self, x: torch.Tensor, return_dict: bool = True) -> Union[AutoencoderKLOutput, Tuple[torch.Tensor]]: + def encode(self, x: torch.Tensor, return_dict: bool = True) -> AutoencoderKLOutput | tuple[torch.Tensor]: h = self.encoder(x) moments = self.quant_conv(h) posterior = DiagonalGaussianDistribution(moments) @@ -124,10 +122,10 @@ def encode(self, x: torch.Tensor, return_dict: bool = True) -> Union[Autoencoder def _decode( self, z: torch.Tensor, - image: Optional[torch.Tensor] = None, - mask: Optional[torch.Tensor] = None, + image: torch.Tensor | None = None, + mask: torch.Tensor | None = None, return_dict: bool = True, - ) -> Union[DecoderOutput, Tuple[torch.Tensor]]: + ) -> DecoderOutput | tuple[torch.Tensor]: z = self.post_quant_conv(z) dec = self.decoder(z, image, mask) @@ -140,11 +138,11 @@ def _decode( def decode( self, z: torch.Tensor, - generator: Optional[torch.Generator] = None, - image: Optional[torch.Tensor] = None, - mask: Optional[torch.Tensor] = None, + generator: torch.Generator | None = None, + image: torch.Tensor | None = None, + mask: torch.Tensor | None = None, return_dict: bool = True, - ) -> Union[DecoderOutput, Tuple[torch.Tensor]]: + ) -> DecoderOutput | tuple[torch.Tensor]: decoded = self._decode(z, image, mask).sample if not return_dict: @@ -155,11 +153,11 @@ def decode( def forward( self, sample: torch.Tensor, - mask: Optional[torch.Tensor] = None, + mask: torch.Tensor | None = None, sample_posterior: bool = False, return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[DecoderOutput, Tuple[torch.Tensor]]: + generator: torch.Generator | None = None, + ) -> DecoderOutput | tuple[torch.Tensor]: r""" Args: sample (`torch.Tensor`): Input sample. diff --git a/src/diffusers/models/autoencoders/autoencoder_dc.py b/src/diffusers/models/autoencoders/autoencoder_dc.py index ec301ef8ad51..02a83d79aba5 100644 --- a/src/diffusers/models/autoencoders/autoencoder_dc.py +++ b/src/diffusers/models/autoencoders/autoencoder_dc.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Tuple, Union import torch import torch.nn as nn @@ -68,7 +67,7 @@ def __init__( in_channels: int, mult: float = 1.0, attention_head_dim: int = 32, - qkv_multiscales: Tuple[int, ...] = (5,), + qkv_multiscales: tuple[int, ...] = (5,), norm_type: str = "batch_norm", ) -> None: super().__init__() @@ -102,14 +101,14 @@ def get_block( attention_head_dim: int, norm_type: str, act_fn: str, - qkv_mutliscales: Tuple[int, ...] = (), + qkv_multiscales: tuple[int, ...] = (), ): if block_type == "ResBlock": block = ResBlock(in_channels, out_channels, norm_type, act_fn) elif block_type == "EfficientViTBlock": block = EfficientViTBlock( - in_channels, attention_head_dim=attention_head_dim, norm_type=norm_type, qkv_multiscales=qkv_mutliscales + in_channels, attention_head_dim=attention_head_dim, norm_type=norm_type, qkv_multiscales=qkv_multiscales ) else: @@ -205,10 +204,10 @@ def __init__( in_channels: int, latent_channels: int, attention_head_dim: int = 32, - block_type: Union[str, Tuple[str]] = "ResBlock", - block_out_channels: Tuple[int, ...] = (128, 256, 512, 512, 1024, 1024), - layers_per_block: Tuple[int, ...] = (2, 2, 2, 2, 2, 2), - qkv_multiscales: Tuple[Tuple[int, ...], ...] = ((), (), (), (5,), (5,), (5,)), + block_type: str | tuple[str] = "ResBlock", + block_out_channels: tuple[int, ...] = (128, 256, 512, 512, 1024, 1024), + layers_per_block: tuple[int, ...] = (2, 2, 2, 2, 2, 2), + qkv_multiscales: tuple[tuple[int, ...], ...] = ((), (), (), (5,), (5,), (5,)), downsample_block_type: str = "pixel_unshuffle", out_shortcut: bool = True, ): @@ -247,7 +246,7 @@ def __init__( attention_head_dim=attention_head_dim, norm_type="rms_norm", act_fn="silu", - qkv_mutliscales=qkv_multiscales[i], + qkv_multiscales=qkv_multiscales[i], ) down_block_list.append(block) @@ -291,12 +290,12 @@ def __init__( in_channels: int, latent_channels: int, attention_head_dim: int = 32, - block_type: Union[str, Tuple[str]] = "ResBlock", - block_out_channels: Tuple[int, ...] = (128, 256, 512, 512, 1024, 1024), - layers_per_block: Tuple[int, ...] = (2, 2, 2, 2, 2, 2), - qkv_multiscales: Tuple[Tuple[int, ...], ...] = ((), (), (), (5,), (5,), (5,)), - norm_type: Union[str, Tuple[str]] = "rms_norm", - act_fn: Union[str, Tuple[str]] = "silu", + block_type: str | tuple[str] = "ResBlock", + block_out_channels: tuple[int, ...] = (128, 256, 512, 512, 1024, 1024), + layers_per_block: tuple[int, ...] = (2, 2, 2, 2, 2, 2), + qkv_multiscales: tuple[tuple[int, ...], ...] = ((), (), (), (5,), (5,), (5,)), + norm_type: str | tuple[str] = "rms_norm", + act_fn: str | tuple[str] = "silu", upsample_block_type: str = "pixel_shuffle", in_shortcut: bool = True, conv_act_fn: str = "relu", @@ -339,7 +338,7 @@ def __init__( attention_head_dim=attention_head_dim, norm_type=norm_type[i], act_fn=act_fn[i], - qkv_mutliscales=qkv_multiscales[i], + qkv_multiscales=qkv_multiscales[i], ) up_block_list.append(block) @@ -391,29 +390,29 @@ class AutoencoderDC(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModel The number of input channels in samples. latent_channels (`int`, defaults to `32`): The number of channels in the latent space representation. - encoder_block_types (`Union[str, Tuple[str]]`, defaults to `"ResBlock"`): + encoder_block_types (`str | tuple[str]`, defaults to `"ResBlock"`): The type(s) of block to use in the encoder. - decoder_block_types (`Union[str, Tuple[str]]`, defaults to `"ResBlock"`): + decoder_block_types (`str | tuple[str]`, defaults to `"ResBlock"`): The type(s) of block to use in the decoder. - encoder_block_out_channels (`Tuple[int, ...]`, defaults to `(128, 256, 512, 512, 1024, 1024)`): + encoder_block_out_channels (`tuple[int, ...]`, defaults to `(128, 256, 512, 512, 1024, 1024)`): The number of output channels for each block in the encoder. - decoder_block_out_channels (`Tuple[int, ...]`, defaults to `(128, 256, 512, 512, 1024, 1024)`): + decoder_block_out_channels (`tuple[int, ...]`, defaults to `(128, 256, 512, 512, 1024, 1024)`): The number of output channels for each block in the decoder. - encoder_layers_per_block (`Tuple[int]`, defaults to `(2, 2, 2, 3, 3, 3)`): + encoder_layers_per_block (`tuple[int]`, defaults to `(2, 2, 2, 3, 3, 3)`): The number of layers per block in the encoder. - decoder_layers_per_block (`Tuple[int]`, defaults to `(3, 3, 3, 3, 3, 3)`): + decoder_layers_per_block (`tuple[int]`, defaults to `(3, 3, 3, 3, 3, 3)`): The number of layers per block in the decoder. - encoder_qkv_multiscales (`Tuple[Tuple[int, ...], ...]`, defaults to `((), (), (), (5,), (5,), (5,))`): + encoder_qkv_multiscales (`tuple[tuple[int, ...], ...]`, defaults to `((), (), (), (5,), (5,), (5,))`): Multi-scale configurations for the encoder's QKV (query-key-value) transformations. - decoder_qkv_multiscales (`Tuple[Tuple[int, ...], ...]`, defaults to `((), (), (), (5,), (5,), (5,))`): + decoder_qkv_multiscales (`tuple[tuple[int, ...], ...]`, defaults to `((), (), (), (5,), (5,), (5,))`): Multi-scale configurations for the decoder's QKV (query-key-value) transformations. upsample_block_type (`str`, defaults to `"pixel_shuffle"`): The type of block to use for upsampling in the decoder. downsample_block_type (`str`, defaults to `"pixel_unshuffle"`): The type of block to use for downsampling in the encoder. - decoder_norm_types (`Union[str, Tuple[str]]`, defaults to `"rms_norm"`): + decoder_norm_types (`str | tuple[str]`, defaults to `"rms_norm"`): The normalization type(s) to use in the decoder. - decoder_act_fns (`Union[str, Tuple[str]]`, defaults to `"silu"`): + decoder_act_fns (`str | tuple[str]`, defaults to `"silu"`): The activation function(s) to use in the decoder. encoder_out_shortcut (`bool`, defaults to `True`): Whether to use shortcut at the end of the encoder. @@ -436,18 +435,18 @@ def __init__( in_channels: int = 3, latent_channels: int = 32, attention_head_dim: int = 32, - encoder_block_types: Union[str, Tuple[str]] = "ResBlock", - decoder_block_types: Union[str, Tuple[str]] = "ResBlock", - encoder_block_out_channels: Tuple[int, ...] = (128, 256, 512, 512, 1024, 1024), - decoder_block_out_channels: Tuple[int, ...] = (128, 256, 512, 512, 1024, 1024), - encoder_layers_per_block: Tuple[int, ...] = (2, 2, 2, 3, 3, 3), - decoder_layers_per_block: Tuple[int, ...] = (3, 3, 3, 3, 3, 3), - encoder_qkv_multiscales: Tuple[Tuple[int, ...], ...] = ((), (), (), (5,), (5,), (5,)), - decoder_qkv_multiscales: Tuple[Tuple[int, ...], ...] = ((), (), (), (5,), (5,), (5,)), + encoder_block_types: str | tuple[str] = "ResBlock", + decoder_block_types: str | tuple[str] = "ResBlock", + encoder_block_out_channels: tuple[int, ...] = (128, 256, 512, 512, 1024, 1024), + decoder_block_out_channels: tuple[int, ...] = (128, 256, 512, 512, 1024, 1024), + encoder_layers_per_block: tuple[int, ...] = (2, 2, 2, 3, 3, 3), + decoder_layers_per_block: tuple[int, ...] = (3, 3, 3, 3, 3, 3), + encoder_qkv_multiscales: tuple[tuple[int, ...], ...] = ((), (), (), (5,), (5,), (5,)), + decoder_qkv_multiscales: tuple[tuple[int, ...], ...] = ((), (), (), (5,), (5,), (5,)), upsample_block_type: str = "pixel_shuffle", downsample_block_type: str = "pixel_unshuffle", - decoder_norm_types: Union[str, Tuple[str]] = "rms_norm", - decoder_act_fns: Union[str, Tuple[str]] = "silu", + decoder_norm_types: str | tuple[str] = "rms_norm", + decoder_act_fns: str | tuple[str] = "silu", encoder_out_shortcut: bool = True, decoder_in_shortcut: bool = True, decoder_conv_act_fn: str = "relu", @@ -506,10 +505,10 @@ def __init__( def enable_tiling( self, - tile_sample_min_height: Optional[int] = None, - tile_sample_min_width: Optional[int] = None, - tile_sample_stride_height: Optional[float] = None, - tile_sample_stride_width: Optional[float] = None, + tile_sample_min_height: int | None = None, + tile_sample_min_width: int | None = None, + tile_sample_stride_height: float | None = None, + tile_sample_stride_width: float | None = None, ) -> None: r""" Enable tiled AE decoding. When this option is enabled, the AE will split the input tensor into tiles to compute @@ -547,7 +546,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor: return encoded @apply_forward_hook - def encode(self, x: torch.Tensor, return_dict: bool = True) -> Union[EncoderOutput, Tuple[torch.Tensor]]: + def encode(self, x: torch.Tensor, return_dict: bool = True) -> EncoderOutput | tuple[torch.Tensor]: r""" Encode a batch of images into latents. @@ -581,7 +580,7 @@ def _decode(self, z: torch.Tensor) -> torch.Tensor: return decoded @apply_forward_hook - def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, Tuple[torch.Tensor]]: + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]: r""" Decode a batch of images. @@ -665,7 +664,7 @@ def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> torch.Tenso return (encoded,) return EncoderOutput(latent=encoded) - def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: batch_size, num_channels, height, width = z.shape tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio diff --git a/src/diffusers/models/autoencoders/autoencoder_kl.py b/src/diffusers/models/autoencoders/autoencoder_kl.py index 95991dca3304..d2e7318f5679 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl.py @@ -11,8 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Tuple, Union - import torch import torch.nn as nn @@ -47,12 +45,12 @@ class AutoencoderKL( Parameters: in_channels (int, *optional*, defaults to 3): Number of channels in the input image. out_channels (int, *optional*, defaults to 3): Number of channels in the output. - down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`): - Tuple of downsample block types. - up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`): - Tuple of upsample block types. - block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`): - Tuple of block output channels. + down_block_types (`tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`): + tuple of downsample block types. + up_block_types (`tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`): + tuple of upsample block types. + block_out_channels (`tuple[int]`, *optional*, defaults to `(64,)`): + tuple of block output channels. act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space. sample_size (`int`, *optional*, defaults to `32`): Sample input size. @@ -81,18 +79,18 @@ def __init__( self, in_channels: int = 3, out_channels: int = 3, - down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",), - up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",), - block_out_channels: Tuple[int, ...] = (64,), + down_block_types: tuple[str] = ("DownEncoderBlock2D",), + up_block_types: tuple[str] = ("UpDecoderBlock2D",), + block_out_channels: tuple[int] = (64,), layers_per_block: int = 1, act_fn: str = "silu", latent_channels: int = 4, norm_num_groups: int = 32, sample_size: int = 32, scaling_factor: float = 0.18215, - shift_factor: Optional[float] = None, - latents_mean: Optional[Tuple[float]] = None, - latents_std: Optional[Tuple[float]] = None, + shift_factor: float | None = None, + latents_mean: tuple[float] | None = None, + latents_std: tuple[float] | None = None, force_upcast: bool = True, use_quant_conv: bool = True, use_post_quant_conv: bool = True, @@ -172,7 +170,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor: @apply_forward_hook def encode( self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: """ Encode a batch of images into latents. @@ -198,7 +196,7 @@ def encode( return AutoencoderKLOutput(latent_dist=posterior) - def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size): return self.tiled_decode(z, return_dict=return_dict) @@ -215,7 +213,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOut @apply_forward_hook def decode( self, z: torch.FloatTensor, return_dict: bool = True, generator=None - ) -> Union[DecoderOutput, torch.FloatTensor]: + ) -> DecoderOutput | torch.FloatTensor: """ Decode a batch of images. @@ -363,7 +361,7 @@ def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> Autoencoder return AutoencoderKLOutput(latent_dist=posterior) - def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: r""" Decode a batch of images using a tiled decoder. @@ -417,8 +415,8 @@ def forward( sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[DecoderOutput, torch.Tensor]: + generator: torch.Generator | None = None, + ) -> DecoderOutput | torch.Tensor: r""" Args: sample (`torch.Tensor`): Input sample. diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_allegro.py b/src/diffusers/models/autoencoders/autoencoder_kl_allegro.py index 6756586460d3..463f8f41bc10 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_allegro.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_allegro.py @@ -14,7 +14,6 @@ # limitations under the License. import math -from typing import Optional, Tuple, Union import torch import torch.nn as nn @@ -40,7 +39,7 @@ class AllegroTemporalConvLayer(nn.Module): def __init__( self, in_dim: int, - out_dim: Optional[int] = None, + out_dim: int | None = None, dropout: float = 0.0, norm_num_groups: int = 32, up_sample: bool = False, @@ -234,7 +233,7 @@ def __init__( output_scale_factor: float = 1.0, spatial_upsample: bool = True, temporal_upsample: bool = False, - temb_channels: Optional[int] = None, + temb_channels: int | None = None, ): super().__init__() @@ -417,14 +416,14 @@ def __init__( self, in_channels: int = 3, out_channels: int = 3, - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str, ...] = ( "AllegroDownBlock3D", "AllegroDownBlock3D", "AllegroDownBlock3D", "AllegroDownBlock3D", ), - block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), - temporal_downsample_blocks: Tuple[bool, ...] = [True, True, False, False], + block_out_channels: tuple[int, ...] = (128, 256, 512, 512), + temporal_downsample_blocks: tuple[bool, ...] = [True, True, False, False], layers_per_block: int = 2, norm_num_groups: int = 32, act_fn: str = "silu", @@ -544,14 +543,14 @@ def __init__( self, in_channels: int = 4, out_channels: int = 3, - up_block_types: Tuple[str, ...] = ( + up_block_types: tuple[str, ...] = ( "AllegroUpBlock3D", "AllegroUpBlock3D", "AllegroUpBlock3D", "AllegroUpBlock3D", ), - temporal_upsample_blocks: Tuple[bool, ...] = [False, True, True, False], - block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), + temporal_upsample_blocks: tuple[bool, ...] = [False, True, True, False], + block_out_channels: tuple[int, ...] = (128, 256, 512, 512), layers_per_block: int = 2, norm_num_groups: int = 32, act_fn: str = "silu", @@ -687,14 +686,14 @@ class AutoencoderKLAllegro(ModelMixin, AutoencoderMixin, ConfigMixin): Number of channels in the input image. out_channels (int, defaults to `3`): Number of channels in the output. - down_block_types (`Tuple[str, ...]`, defaults to `("AllegroDownBlock3D", "AllegroDownBlock3D", "AllegroDownBlock3D", "AllegroDownBlock3D")`): - Tuple of strings denoting which types of down blocks to use. - up_block_types (`Tuple[str, ...]`, defaults to `("AllegroUpBlock3D", "AllegroUpBlock3D", "AllegroUpBlock3D", "AllegroUpBlock3D")`): - Tuple of strings denoting which types of up blocks to use. - block_out_channels (`Tuple[int, ...]`, defaults to `(128, 256, 512, 512)`): - Tuple of integers denoting number of output channels in each block. - temporal_downsample_blocks (`Tuple[bool, ...]`, defaults to `(True, True, False, False)`): - Tuple of booleans denoting which blocks to enable temporal downsampling in. + down_block_types (`tuple[str, ...]`, defaults to `("AllegroDownBlock3D", "AllegroDownBlock3D", "AllegroDownBlock3D", "AllegroDownBlock3D")`): + tuple of strings denoting which types of down blocks to use. + up_block_types (`tuple[str, ...]`, defaults to `("AllegroUpBlock3D", "AllegroUpBlock3D", "AllegroUpBlock3D", "AllegroUpBlock3D")`): + tuple of strings denoting which types of up blocks to use. + block_out_channels (`tuple[int, ...]`, defaults to `(128, 256, 512, 512)`): + tuple of integers denoting number of output channels in each block. + temporal_downsample_blocks (`tuple[bool, ...]`, defaults to `(True, True, False, False)`): + tuple of booleans denoting which blocks to enable temporal downsampling in. latent_channels (`int`, defaults to `4`): Number of channels in latents. layers_per_block (`int`, defaults to `2`): @@ -727,21 +726,21 @@ def __init__( self, in_channels: int = 3, out_channels: int = 3, - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str, ...] = ( "AllegroDownBlock3D", "AllegroDownBlock3D", "AllegroDownBlock3D", "AllegroDownBlock3D", ), - up_block_types: Tuple[str, ...] = ( + up_block_types: tuple[str, ...] = ( "AllegroUpBlock3D", "AllegroUpBlock3D", "AllegroUpBlock3D", "AllegroUpBlock3D", ), - block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), - temporal_downsample_blocks: Tuple[bool, ...] = (True, True, False, False), - temporal_upsample_blocks: Tuple[bool, ...] = (False, True, True, False), + block_out_channels: tuple[int, ...] = (128, 256, 512, 512), + temporal_downsample_blocks: tuple[bool, ...] = (True, True, False, False), + temporal_upsample_blocks: tuple[bool, ...] = (False, True, True, False), latent_channels: int = 4, layers_per_block: int = 2, act_fn: str = "silu", @@ -807,7 +806,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor: @apply_forward_hook def encode( self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: r""" Encode a batch of videos into latents. @@ -842,7 +841,7 @@ def _decode(self, z: torch.Tensor) -> torch.Tensor: raise NotImplementedError("Decoding without tiling has not been implemented yet.") @apply_forward_hook - def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: """ Decode a batch of videos. @@ -1044,8 +1043,8 @@ def forward( sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[DecoderOutput, torch.Tensor]: + generator: torch.Generator | None = None, + ) -> DecoderOutput | torch.Tensor: r""" Args: sample (`torch.Tensor`): Input sample. diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py b/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py index 79433f7b9232..9921e3932465 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py @@ -13,8 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Dict, Optional, Tuple, Union - import numpy as np import torch import torch.nn as nn @@ -72,7 +70,7 @@ class CogVideoXCausalConv3d(nn.Module): Args: in_channels (`int`): Number of channels in the input tensor. out_channels (`int`): Number of output channels produced by the convolution. - kernel_size (`int` or `Tuple[int, int, int]`): Kernel size of the convolutional kernel. + kernel_size (`int` or `tuple[int, int, int]`): Kernel size of the convolutional kernel. stride (`int`, defaults to `1`): Stride of the convolution. dilation (`int`, defaults to `1`): Dilation rate of the convolution. pad_mode (`str`, defaults to `"constant"`): Padding mode. @@ -82,7 +80,7 @@ def __init__( self, in_channels: int, out_channels: int, - kernel_size: Union[int, Tuple[int, int, int]], + kernel_size: int | tuple[int, int, int], stride: int = 1, dilation: int = 1, pad_mode: str = "constant", @@ -123,7 +121,7 @@ def __init__( ) def fake_context_parallel_forward( - self, inputs: torch.Tensor, conv_cache: Optional[torch.Tensor] = None + self, inputs: torch.Tensor, conv_cache: torch.Tensor | None = None ) -> torch.Tensor: if self.pad_mode == "replicate": inputs = F.pad(inputs, self.time_causal_padding, mode="replicate") @@ -134,7 +132,7 @@ def fake_context_parallel_forward( inputs = torch.cat(cached_inputs + [inputs], dim=2) return inputs - def forward(self, inputs: torch.Tensor, conv_cache: Optional[torch.Tensor] = None) -> torch.Tensor: + def forward(self, inputs: torch.Tensor, conv_cache: torch.Tensor | None = None) -> torch.Tensor: inputs = self.fake_context_parallel_forward(inputs, conv_cache) if self.pad_mode == "replicate": @@ -174,7 +172,7 @@ def __init__( self.conv_b = CogVideoXCausalConv3d(zq_channels, f_channels, kernel_size=1, stride=1) def forward( - self, f: torch.Tensor, zq: torch.Tensor, conv_cache: Optional[Dict[str, torch.Tensor]] = None + self, f: torch.Tensor, zq: torch.Tensor, conv_cache: dict[str, torch.Tensor] | None = None ) -> torch.Tensor: new_conv_cache = {} conv_cache = conv_cache or {} @@ -227,14 +225,14 @@ class CogVideoXResnetBlock3D(nn.Module): def __init__( self, in_channels: int, - out_channels: Optional[int] = None, + out_channels: int | None = None, dropout: float = 0.0, temb_channels: int = 512, groups: int = 32, eps: float = 1e-6, non_linearity: str = "swish", conv_shortcut: bool = False, - spatial_norm_dim: Optional[int] = None, + spatial_norm_dim: int | None = None, pad_mode: str = "first", ): super().__init__() @@ -287,9 +285,9 @@ def __init__( def forward( self, inputs: torch.Tensor, - temb: Optional[torch.Tensor] = None, - zq: Optional[torch.Tensor] = None, - conv_cache: Optional[Dict[str, torch.Tensor]] = None, + temb: torch.Tensor | None = None, + zq: torch.Tensor | None = None, + conv_cache: dict[str, torch.Tensor] | None = None, ) -> torch.Tensor: new_conv_cache = {} conv_cache = conv_cache or {} @@ -409,9 +407,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - zq: Optional[torch.Tensor] = None, - conv_cache: Optional[Dict[str, torch.Tensor]] = None, + temb: torch.Tensor | None = None, + zq: torch.Tensor | None = None, + conv_cache: dict[str, torch.Tensor] | None = None, ) -> torch.Tensor: r"""Forward method of the `CogVideoXDownBlock3D` class.""" @@ -477,7 +475,7 @@ def __init__( resnet_eps: float = 1e-6, resnet_act_fn: str = "swish", resnet_groups: int = 32, - spatial_norm_dim: Optional[int] = None, + spatial_norm_dim: int | None = None, pad_mode: str = "first", ): super().__init__() @@ -504,9 +502,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - zq: Optional[torch.Tensor] = None, - conv_cache: Optional[Dict[str, torch.Tensor]] = None, + temb: torch.Tensor | None = None, + zq: torch.Tensor | None = None, + conv_cache: dict[str, torch.Tensor] | None = None, ) -> torch.Tensor: r"""Forward method of the `CogVideoXMidBlock3D` class.""" @@ -611,9 +609,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - zq: Optional[torch.Tensor] = None, - conv_cache: Optional[Dict[str, torch.Tensor]] = None, + temb: torch.Tensor | None = None, + zq: torch.Tensor | None = None, + conv_cache: dict[str, torch.Tensor] | None = None, ) -> torch.Tensor: r"""Forward method of the `CogVideoXUpBlock3D` class.""" @@ -652,10 +650,10 @@ class CogVideoXEncoder3D(nn.Module): The number of input channels. out_channels (`int`, *optional*, defaults to 3): The number of output channels. - down_block_types (`Tuple[str, ...]`, *optional*, defaults to `("DownEncoderBlock2D",)`): + down_block_types (`tuple[str, ...]`, *optional*, defaults to `("DownEncoderBlock2D",)`): The types of down blocks to use. See `~diffusers.models.unet_2d_blocks.get_down_block` for available options. - block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`): + block_out_channels (`tuple[int, ...]`, *optional*, defaults to `(64,)`): The number of output channels for each block. act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. See `~diffusers.models.activations.get_activation` for available options. @@ -671,13 +669,13 @@ def __init__( self, in_channels: int = 3, out_channels: int = 16, - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str, ...] = ( "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", ), - block_out_channels: Tuple[int, ...] = (128, 256, 256, 512), + block_out_channels: tuple[int, ...] = (128, 256, 256, 512), layers_per_block: int = 3, act_fn: str = "silu", norm_eps: float = 1e-6, @@ -743,8 +741,8 @@ def __init__( def forward( self, sample: torch.Tensor, - temb: Optional[torch.Tensor] = None, - conv_cache: Optional[Dict[str, torch.Tensor]] = None, + temb: torch.Tensor | None = None, + conv_cache: dict[str, torch.Tensor] | None = None, ) -> torch.Tensor: r"""The forward method of the `CogVideoXEncoder3D` class.""" @@ -805,9 +803,9 @@ class CogVideoXDecoder3D(nn.Module): The number of input channels. out_channels (`int`, *optional*, defaults to 3): The number of output channels. - up_block_types (`Tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`): + up_block_types (`tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`): The types of up blocks to use. See `~diffusers.models.unet_2d_blocks.get_up_block` for available options. - block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`): + block_out_channels (`tuple[int, ...]`, *optional*, defaults to `(64,)`): The number of output channels for each block. act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. See `~diffusers.models.activations.get_activation` for available options. @@ -823,13 +821,13 @@ def __init__( self, in_channels: int = 16, out_channels: int = 3, - up_block_types: Tuple[str, ...] = ( + up_block_types: tuple[str, ...] = ( "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", ), - block_out_channels: Tuple[int, ...] = (128, 256, 256, 512), + block_out_channels: tuple[int, ...] = (128, 256, 256, 512), layers_per_block: int = 3, act_fn: str = "silu", norm_eps: float = 1e-6, @@ -902,8 +900,8 @@ def __init__( def forward( self, sample: torch.Tensor, - temb: Optional[torch.Tensor] = None, - conv_cache: Optional[Dict[str, torch.Tensor]] = None, + temb: torch.Tensor | None = None, + conv_cache: dict[str, torch.Tensor] | None = None, ) -> torch.Tensor: r"""The forward method of the `CogVideoXDecoder3D` class.""" @@ -966,12 +964,12 @@ class AutoencoderKLCogVideoX(ModelMixin, AutoencoderMixin, ConfigMixin, FromOrig Parameters: in_channels (int, *optional*, defaults to 3): Number of channels in the input image. out_channels (int, *optional*, defaults to 3): Number of channels in the output. - down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`): - Tuple of downsample block types. - up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`): - Tuple of upsample block types. - block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`): - Tuple of block output channels. + down_block_types (`tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`): + tuple of downsample block types. + up_block_types (`tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`): + tuple of upsample block types. + block_out_channels (`tuple[int]`, *optional*, defaults to `(64,)`): + tuple of block output channels. act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. sample_size (`int`, *optional*, defaults to `32`): Sample input size. scaling_factor (`float`, *optional*, defaults to `1.15258426`): @@ -995,19 +993,19 @@ def __init__( self, in_channels: int = 3, out_channels: int = 3, - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str] = ( "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D", ), - up_block_types: Tuple[str, ...] = ( + up_block_types: tuple[str] = ( "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D", ), - block_out_channels: Tuple[int, ...] = (128, 256, 256, 512), + block_out_channels: tuple[int] = (128, 256, 256, 512), latent_channels: int = 16, layers_per_block: int = 3, act_fn: str = "silu", @@ -1017,9 +1015,9 @@ def __init__( sample_height: int = 480, sample_width: int = 720, scaling_factor: float = 1.15258426, - shift_factor: Optional[float] = None, - latents_mean: Optional[Tuple[float]] = None, - latents_std: Optional[Tuple[float]] = None, + shift_factor: float | None = None, + latents_mean: tuple[float] | None = None, + latents_std: tuple[float] | None = None, force_upcast: float = True, use_quant_conv: bool = False, use_post_quant_conv: bool = False, @@ -1090,10 +1088,10 @@ def __init__( def enable_tiling( self, - tile_sample_min_height: Optional[int] = None, - tile_sample_min_width: Optional[int] = None, - tile_overlap_factor_height: Optional[float] = None, - tile_overlap_factor_width: Optional[float] = None, + tile_sample_min_height: int | None = None, + tile_sample_min_width: int | None = None, + tile_overlap_factor_height: float | None = None, + tile_overlap_factor_width: float | None = None, ) -> None: r""" Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to @@ -1153,7 +1151,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor: @apply_forward_hook def encode( self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: """ Encode a batch of images into latents. @@ -1178,7 +1176,7 @@ def encode( return (posterior,) return AutoencoderKLOutput(latent_dist=posterior) - def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: batch_size, num_channels, num_frames, height, width = z.shape if self.use_tiling and (width > self.tile_latent_min_width or height > self.tile_latent_min_height): @@ -1207,7 +1205,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOut return DecoderOutput(sample=dec) @apply_forward_hook - def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: """ Decode a batch of images. @@ -1321,7 +1319,7 @@ def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: enc = torch.cat(result_rows, dim=3) return enc - def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: r""" Decode a batch of images using a tiled decoder. @@ -1409,8 +1407,8 @@ def forward( sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[torch.Tensor, torch.Tensor]: + generator: torch.Generator | None = None, + ) -> torch.Tensor | torch.Tensor: x = sample posterior = self.encode(x).latent_dist if sample_posterior: diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py b/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py index b17522d1c424..4fe1f62890be 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py @@ -12,8 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import math -from typing import List, Optional, Tuple, Union import torch import torch.nn as nn @@ -47,9 +48,9 @@ def __init__( self, in_channels: int = 1, out_channels: int = 1, - kernel_size: Union[int, Tuple[int, int, int]] = (3, 3, 3), - dilation: Union[int, Tuple[int, int, int]] = (1, 1, 1), - stride: Union[int, Tuple[int, int, int]] = (1, 1, 1), + kernel_size: int | tuple[int, int, int] = (3, 3, 3), + dilation: int | tuple[int, int, int] = (1, 1, 1), + stride: int | tuple[int, int, int] = (1, 1, 1), padding: int = 1, pad_mode: str = "constant", ) -> None: @@ -419,7 +420,7 @@ def __init__( attention_head_dim: int, num_groups: int = 1, dropout: float = 0.0, - processor: Union["CosmosSpatialAttentionProcessor2_0", "CosmosTemporalAttentionProcessor2_0"] = None, + processor: "CosmosSpatialAttentionProcessor2_0" | "CosmosTemporalAttentionProcessor2_0" = None, ) -> None: super().__init__() self.num_attention_heads = num_attention_heads @@ -438,7 +439,7 @@ def __init__( if self.processor is None: raise ValueError("CosmosCausalAttention requires a processor.") - def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor: return self.processor(self, hidden_states=hidden_states, attention_mask=attention_mask) @@ -450,7 +451,7 @@ def __init__(self): ) def __call__( - self, attn: CosmosCausalAttention, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None + self, attn: CosmosCausalAttention, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None ) -> torch.Tensor: batch_size, num_channels, num_frames, height, width = hidden_states.shape residual = hidden_states @@ -489,7 +490,7 @@ def __init__(self): ) def __call__( - self, attn: CosmosCausalAttention, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None + self, attn: CosmosCausalAttention, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None ) -> torch.Tensor: batch_size, num_channels, num_frames, height, width = hidden_states.shape residual = hidden_states @@ -711,9 +712,9 @@ def __init__( self, in_channels: int = 3, out_channels: int = 16, - block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), + block_out_channels: tuple[int, ...] = (128, 256, 512, 512), num_resnet_blocks: int = 2, - attention_resolutions: Tuple[int, ...] = (32,), + attention_resolutions: tuple[int, ...] = (32,), resolution: int = 1024, patch_size: int = 4, patch_type: str = "haar", @@ -795,9 +796,9 @@ def __init__( self, in_channels: int = 16, out_channels: int = 3, - block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), + block_out_channels: tuple[int, ...] = (128, 256, 512, 512), num_resnet_blocks: int = 2, - attention_resolutions: Tuple[int, ...] = (32,), + attention_resolutions: tuple[int, ...] = (32,), resolution: int = 1024, patch_size: int = 4, patch_type: str = "haar", @@ -886,12 +887,12 @@ class AutoencoderKLCosmos(ModelMixin, AutoencoderMixin, ConfigMixin): Number of output channels. latent_channels (`int`, defaults to `16`): Number of latent channels. - encoder_block_out_channels (`Tuple[int, ...]`, defaults to `(128, 256, 512, 512)`): + encoder_block_out_channels (`tuple[int, ...]`, defaults to `(128, 256, 512, 512)`): Number of output channels for each encoder down block. - decode_block_out_channels (`Tuple[int, ...]`, defaults to `(256, 512, 512, 512)`): + decode_block_out_channels (`tuple[int, ...]`, defaults to `(256, 512, 512, 512)`): Number of output channels for each decoder up block. - attention_resolutions (`Tuple[int, ...]`, defaults to `(32,)`): - List of image/video resolutions at which to apply attention. + attention_resolutions (`tuple[int, ...]`, defaults to `(32,)`): + list of image/video resolutions at which to apply attention. resolution (`int`, defaults to `1024`): Base image/video resolution used for computing whether a block should have attention layers. num_layers (`int`, defaults to `2`): @@ -924,9 +925,9 @@ def __init__( in_channels: int = 3, out_channels: int = 3, latent_channels: int = 16, - encoder_block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), - decode_block_out_channels: Tuple[int, ...] = (256, 512, 512, 512), - attention_resolutions: Tuple[int, ...] = (32,), + encoder_block_out_channels: tuple[int, ...] = (128, 256, 512, 512), + decode_block_out_channels: tuple[int, ...] = (256, 512, 512, 512), + attention_resolutions: tuple[int, ...] = (32,), resolution: int = 1024, num_layers: int = 2, patch_size: int = 4, @@ -934,8 +935,8 @@ def __init__( scaling_factor: float = 1.0, spatial_compression_ratio: int = 8, temporal_compression_ratio: int = 8, - latents_mean: Optional[List[float]] = LATENTS_MEAN, - latents_std: Optional[List[float]] = LATENTS_STD, + latents_mean: list[float] | None = LATENTS_MEAN, + latents_std: list[float] | None = LATENTS_STD, ) -> None: super().__init__() @@ -999,12 +1000,12 @@ def __init__( def enable_tiling( self, - tile_sample_min_height: Optional[int] = None, - tile_sample_min_width: Optional[int] = None, - tile_sample_min_num_frames: Optional[int] = None, - tile_sample_stride_height: Optional[float] = None, - tile_sample_stride_width: Optional[float] = None, - tile_sample_stride_num_frames: Optional[float] = None, + tile_sample_min_height: int | None = None, + tile_sample_min_width: int | None = None, + tile_sample_min_num_frames: int | None = None, + tile_sample_stride_height: float | None = None, + tile_sample_stride_width: float | None = None, + tile_sample_stride_num_frames: float | None = None, ) -> None: r""" Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to @@ -1050,7 +1051,7 @@ def encode(self, x: torch.Tensor, return_dict: bool = True) -> torch.Tensor: return (posterior,) return AutoencoderKLOutput(latent_dist=posterior) - def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, Tuple[torch.Tensor]]: + def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]: z = self.post_quant_conv(z) dec = self.decoder(z) @@ -1059,7 +1060,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOut return DecoderOutput(sample=dec) @apply_forward_hook - def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, Tuple[torch.Tensor]]: + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]: if self.use_slicing and z.shape[0] > 1: decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] decoded = torch.cat(decoded_slices) @@ -1075,8 +1076,8 @@ def forward( sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[Tuple[torch.Tensor], DecoderOutput]: + generator: torch.Generator | None = None, + ) -> tuple[torch.Tensor] | DecoderOutput: x = sample posterior = self.encode(x).latent_dist if sample_posterior: diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_flux2.py b/src/diffusers/models/autoencoders/autoencoder_kl_flux2.py index 3325d33c06bf..c1345d5de73f 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_flux2.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_flux2.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. import math -from typing import Optional, Tuple, Union import torch import torch.nn as nn @@ -48,11 +47,11 @@ class AutoencoderKLFlux2( Parameters: in_channels (int, *optional*, defaults to 3): Number of channels in the input image. out_channels (int, *optional*, defaults to 3): Number of channels in the output. - down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`): + down_block_types (`tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`): Tuple of downsample block types. - up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`): + up_block_types (`tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`): Tuple of upsample block types. - block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`): + block_out_channels (`tuple[int]`, *optional*, defaults to `(64,)`): Tuple of block output channels. act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space. @@ -74,19 +73,19 @@ def __init__( self, in_channels: int = 3, out_channels: int = 3, - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str, ...] = ( "DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D", ), - up_block_types: Tuple[str, ...] = ( + up_block_types: tuple[str, ...] = ( "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", ), - block_out_channels: Tuple[int, ...] = ( + block_out_channels: tuple[int, ...] = ( 128, 256, 512, @@ -103,7 +102,7 @@ def __init__( mid_block_add_attention: bool = True, batch_norm_eps: float = 1e-4, batch_norm_momentum: float = 0.1, - patch_size: Tuple[int, int] = (2, 2), + patch_size: tuple[int, int] = (2, 2), ): super().__init__() @@ -187,7 +186,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor: @apply_forward_hook def encode( self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: """ Encode a batch of images into latents. @@ -213,7 +212,7 @@ def encode( return AutoencoderKLOutput(latent_dist=posterior) - def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size): return self.tiled_decode(z, return_dict=return_dict) @@ -230,7 +229,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOut @apply_forward_hook def decode( self, z: torch.FloatTensor, return_dict: bool = True, generator=None - ) -> Union[DecoderOutput, torch.FloatTensor]: + ) -> DecoderOutput | torch.FloatTensor: """ Decode a batch of images. @@ -378,7 +377,7 @@ def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> Autoencoder return AutoencoderKLOutput(latent_dist=posterior) - def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: r""" Decode a batch of images using a tiled decoder. @@ -432,8 +431,8 @@ def forward( sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[DecoderOutput, torch.Tensor]: + generator: torch.Generator | None = None, + ) -> DecoderOutput | torch.Tensor: r""" Args: sample (`torch.Tensor`): Input sample. diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py index ddc0aed6b0ff..a19c267b6d36 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Tuple, Union - import numpy as np import torch import torch.nn as nn @@ -50,10 +48,10 @@ def __init__( self, in_channels: int, out_channels: int, - kernel_size: Union[int, Tuple[int, int, int]] = 3, - stride: Union[int, Tuple[int, int, int]] = 1, - padding: Union[int, Tuple[int, int, int]] = 0, - dilation: Union[int, Tuple[int, int, int]] = 1, + kernel_size: int | tuple[int, int, int] = 3, + stride: int | tuple[int, int, int] = 1, + padding: int | tuple[int, int, int] = 0, + dilation: int | tuple[int, int, int] = 1, bias: bool = True, pad_mode: str = "replicate", ) -> None: @@ -82,11 +80,11 @@ class HunyuanVideoUpsampleCausal3D(nn.Module): def __init__( self, in_channels: int, - out_channels: Optional[int] = None, + out_channels: int | None = None, kernel_size: int = 3, stride: int = 1, bias: bool = True, - upsample_factor: Tuple[float, float, float] = (2, 2, 2), + upsample_factor: tuple[float, float, float] = (2, 2, 2), ) -> None: super().__init__() @@ -124,7 +122,7 @@ class HunyuanVideoDownsampleCausal3D(nn.Module): def __init__( self, channels: int, - out_channels: Optional[int] = None, + out_channels: int | None = None, padding: int = 1, kernel_size: int = 3, bias: bool = True, @@ -144,7 +142,7 @@ class HunyuanVideoResnetBlockCausal3D(nn.Module): def __init__( self, in_channels: int, - out_channels: Optional[int] = None, + out_channels: int | None = None, dropout: float = 0.0, groups: int = 32, eps: float = 1e-6, @@ -357,7 +355,7 @@ def __init__( resnet_act_fn: str = "swish", resnet_groups: int = 32, add_upsample: bool = True, - upsample_scale_factor: Tuple[int, int, int] = (2, 2, 2), + upsample_scale_factor: tuple[int, int, int] = (2, 2, 2), ) -> None: super().__init__() resnets = [] @@ -418,13 +416,13 @@ def __init__( self, in_channels: int = 3, out_channels: int = 3, - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str, ...] = ( "HunyuanVideoDownBlock3D", "HunyuanVideoDownBlock3D", "HunyuanVideoDownBlock3D", "HunyuanVideoDownBlock3D", ), - block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), + block_out_channels: tuple[int, ...] = (128, 256, 512, 512), layers_per_block: int = 2, norm_num_groups: int = 32, act_fn: str = "silu", @@ -526,13 +524,13 @@ def __init__( self, in_channels: int = 3, out_channels: int = 3, - up_block_types: Tuple[str, ...] = ( + up_block_types: tuple[str, ...] = ( "HunyuanVideoUpBlock3D", "HunyuanVideoUpBlock3D", "HunyuanVideoUpBlock3D", "HunyuanVideoUpBlock3D", ), - block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), + block_out_channels: tuple[int, ...] = (128, 256, 512, 512), layers_per_block: int = 2, norm_num_groups: int = 32, act_fn: str = "silu", @@ -641,19 +639,19 @@ def __init__( in_channels: int = 3, out_channels: int = 3, latent_channels: int = 16, - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str, ...] = ( "HunyuanVideoDownBlock3D", "HunyuanVideoDownBlock3D", "HunyuanVideoDownBlock3D", "HunyuanVideoDownBlock3D", ), - up_block_types: Tuple[str, ...] = ( + up_block_types: tuple[str, ...] = ( "HunyuanVideoUpBlock3D", "HunyuanVideoUpBlock3D", "HunyuanVideoUpBlock3D", "HunyuanVideoUpBlock3D", ), - block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), + block_out_channels: tuple[int] = (128, 256, 512, 512), layers_per_block: int = 2, act_fn: str = "silu", norm_num_groups: int = 32, @@ -725,12 +723,12 @@ def __init__( def enable_tiling( self, - tile_sample_min_height: Optional[int] = None, - tile_sample_min_width: Optional[int] = None, - tile_sample_min_num_frames: Optional[int] = None, - tile_sample_stride_height: Optional[float] = None, - tile_sample_stride_width: Optional[float] = None, - tile_sample_stride_num_frames: Optional[float] = None, + tile_sample_min_height: int | None = None, + tile_sample_min_width: int | None = None, + tile_sample_min_num_frames: int | None = None, + tile_sample_stride_height: float | None = None, + tile_sample_stride_width: float | None = None, + tile_sample_stride_num_frames: float | None = None, ) -> None: r""" Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to @@ -779,7 +777,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor: @apply_forward_hook def encode( self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: r""" Encode a batch of images into latents. @@ -804,7 +802,7 @@ def encode( return (posterior,) return AutoencoderKLOutput(latent_dist=posterior) - def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: batch_size, num_channels, num_frames, height, width = z.shape tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio @@ -825,7 +823,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOut return DecoderOutput(sample=dec) @apply_forward_hook - def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: r""" Decode a batch of images. @@ -924,7 +922,7 @@ def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput: enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width] return enc - def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: r""" Decode a batch of images using a tiled decoder. @@ -1013,7 +1011,7 @@ def _temporal_tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput: enc = torch.cat(result_row, dim=2)[:, :, :latent_num_frames] return enc - def _temporal_tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def _temporal_tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: batch_size, num_channels, num_frames, height, width = z.shape num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1 @@ -1054,8 +1052,8 @@ def forward( sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[DecoderOutput, torch.Tensor]: + generator: torch.Generator | None = None, + ) -> DecoderOutput | torch.Tensor: r""" Args: sample (`torch.Tensor`): Input sample. diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage.py b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage.py index 616d0d415840..6922ac853554 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Tuple, Union import numpy as np import torch @@ -27,7 +26,7 @@ from ..activations import get_activation from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin -from .vae import DecoderOutput, DiagonalGaussianDistribution +from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -238,7 +237,7 @@ def __init__( self, in_channels: int, z_channels: int, - block_out_channels: Tuple[int, ...], + block_out_channels: tuple[int, ...], num_res_blocks: int, spatial_compression_ratio: int, non_linearity: str = "silu", @@ -329,7 +328,7 @@ class HunyuanImageDecoder2D(nn.Module): Number of latent channels. out_channels : int Number of output channels. - block_out_channels : Tuple[int, ...] + block_out_channels : tuple[int, ...] Output channels for each block. num_res_blocks : int Number of residual blocks per block. @@ -344,7 +343,7 @@ def __init__( self, z_channels: int, out_channels: int, - block_out_channels: Tuple[int, ...], + block_out_channels: tuple[int, ...], num_res_blocks: int, spatial_compression_ratio: int, upsample_match_channel: bool = True, @@ -410,7 +409,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return h -class AutoencoderKLHunyuanImage(ModelMixin, ConfigMixin, FromOriginalModelMixin): +class AutoencoderKLHunyuanImage(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): r""" A VAE model for 2D images with spatial tiling support. @@ -427,7 +426,7 @@ def __init__( in_channels: int, out_channels: int, latent_channels: int, - block_out_channels: Tuple[int, ...], + block_out_channels: tuple[int, ...], layers_per_block: int, spatial_compression_ratio: int, sample_size: int, @@ -467,8 +466,8 @@ def __init__( def enable_tiling( self, - tile_sample_min_size: Optional[int] = None, - tile_overlap_factor: Optional[float] = None, + tile_sample_min_size: int | None = None, + tile_overlap_factor: float | None = None, ) -> None: r""" Enable spatial tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles @@ -486,27 +485,6 @@ def enable_tiling( self.tile_overlap_factor = tile_overlap_factor or self.tile_overlap_factor self.tile_latent_min_size = self.tile_sample_min_size // self.config.spatial_compression_ratio - def disable_tiling(self) -> None: - r""" - Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_tiling = False - - def enable_slicing(self) -> None: - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.use_slicing = True - - def disable_slicing(self) -> None: - r""" - Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_slicing = False - def _encode(self, x: torch.Tensor): batch_size, num_channels, height, width = x.shape @@ -521,7 +499,7 @@ def _encode(self, x: torch.Tensor): @apply_forward_hook def encode( self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: r""" Encode a batch of images into latents. @@ -560,7 +538,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True): return DecoderOutput(sample=dec) @apply_forward_hook - def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: r""" Decode a batch of images. @@ -641,7 +619,7 @@ def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: return moments - def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: """ Decode latent using spatial tiling strategy. @@ -691,8 +669,8 @@ def forward( sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[DecoderOutput, torch.Tensor]: + generator: torch.Generator | None = None, + ) -> DecoderOutput | torch.Tensor: """ Args: sample (`torch.Tensor`): Input sample. diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py index 2249063a9f00..81957e2feed4 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Tuple, Union import numpy as np import torch @@ -26,7 +25,7 @@ from ..activations import get_activation from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin -from .vae import DecoderOutput, DiagonalGaussianDistribution +from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -37,10 +36,10 @@ def __init__( self, in_channels: int, out_channels: int, - kernel_size: Union[int, Tuple[int, int, int]] = 3, - stride: Union[int, Tuple[int, int, int]] = 1, - padding: Union[int, Tuple[int, int, int]] = 0, - dilation: Union[int, Tuple[int, int, int]] = 1, + kernel_size: int | tuple[int, int, int] = 3, + stride: int | tuple[int, int, int] = 1, + padding: int | tuple[int, int, int] = 0, + dilation: int | tuple[int, int, int] = 1, bias: bool = True, pad_mode: str = "replicate", ) -> None: @@ -226,7 +225,7 @@ class HunyuanImageRefinerResnetBlock(nn.Module): def __init__( self, in_channels: int, - out_channels: Optional[int] = None, + out_channels: int | None = None, non_linearity: str = "swish", ) -> None: super().__init__() @@ -315,7 +314,7 @@ def __init__( in_channels: int, out_channels: int, num_layers: int = 1, - downsample_out_channels: Optional[int] = None, + downsample_out_channels: int | None = None, add_temporal_downsample: int = True, ) -> None: super().__init__() @@ -364,7 +363,7 @@ def __init__( in_channels: int, out_channels: int, num_layers: int = 1, - upsample_out_channels: Optional[int] = None, + upsample_out_channels: int | None = None, add_temporal_upsample: bool = True, ) -> None: super().__init__() @@ -422,7 +421,7 @@ def __init__( self, in_channels: int = 3, out_channels: int = 64, - block_out_channels: Tuple[int, ...] = (128, 256, 512, 1024, 1024), + block_out_channels: tuple[int, ...] = (128, 256, 512, 1024, 1024), layers_per_block: int = 2, temporal_compression_ratio: int = 4, spatial_compression_ratio: int = 16, @@ -509,7 +508,7 @@ def __init__( self, in_channels: int = 32, out_channels: int = 3, - block_out_channels: Tuple[int, ...] = (1024, 1024, 512, 256, 128), + block_out_channels: tuple[int, ...] = (1024, 1024, 512, 256, 128), layers_per_block: int = 2, spatial_compression_ratio: int = 16, temporal_compression_ratio: int = 4, @@ -584,7 +583,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return hidden_states -class AutoencoderKLHunyuanImageRefiner(ModelMixin, ConfigMixin): +class AutoencoderKLHunyuanImageRefiner(ModelMixin, AutoencoderMixin, ConfigMixin): r""" A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Used for HunyuanImage-2.1 Refiner. @@ -601,7 +600,7 @@ def __init__( in_channels: int = 3, out_channels: int = 3, latent_channels: int = 32, - block_out_channels: Tuple[int, ...] = (128, 256, 512, 1024, 1024), + block_out_channels: tuple[int, ...] = (128, 256, 512, 1024, 1024), layers_per_block: int = 2, spatial_compression_ratio: int = 16, temporal_compression_ratio: int = 4, @@ -655,11 +654,11 @@ def __init__( def enable_tiling( self, - tile_sample_min_height: Optional[int] = None, - tile_sample_min_width: Optional[int] = None, - tile_sample_stride_height: Optional[float] = None, - tile_sample_stride_width: Optional[float] = None, - tile_overlap_factor: Optional[float] = None, + tile_sample_min_height: int | None = None, + tile_sample_min_width: int | None = None, + tile_sample_stride_height: float | None = None, + tile_sample_stride_width: float | None = None, + tile_overlap_factor: float | None = None, ) -> None: r""" Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to @@ -685,27 +684,6 @@ def enable_tiling( self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width self.tile_overlap_factor = tile_overlap_factor or self.tile_overlap_factor - def disable_tiling(self) -> None: - r""" - Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_tiling = False - - def enable_slicing(self) -> None: - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.use_slicing = True - - def disable_slicing(self) -> None: - r""" - Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_slicing = False - def _encode(self, x: torch.Tensor) -> torch.Tensor: _, _, _, height, width = x.shape @@ -718,7 +696,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor: @apply_forward_hook def encode( self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: r""" Encode a batch of images into latents. @@ -756,7 +734,7 @@ def _decode(self, z: torch.Tensor) -> torch.Tensor: return dec @apply_forward_hook - def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: r""" Decode a batch of images. @@ -914,8 +892,8 @@ def forward( sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[DecoderOutput, torch.Tensor]: + generator: torch.Generator | None = None, + ) -> DecoderOutput | torch.Tensor: r""" Args: sample (`torch.Tensor`): Input sample. diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanvideo15.py b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanvideo15.py index 4b1beb74a3bc..2c38b174a100 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanvideo15.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanvideo15.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Tuple, Union import numpy as np import torch @@ -26,7 +25,7 @@ from ..activations import get_activation from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin -from .vae import DecoderOutput, DiagonalGaussianDistribution +from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -37,10 +36,10 @@ def __init__( self, in_channels: int, out_channels: int, - kernel_size: Union[int, Tuple[int, int, int]] = 3, - stride: Union[int, Tuple[int, int, int]] = 1, - padding: Union[int, Tuple[int, int, int]] = 0, - dilation: Union[int, Tuple[int, int, int]] = 1, + kernel_size: int | tuple[int, int, int] = 3, + stride: int | tuple[int, int, int] = 1, + padding: int | tuple[int, int, int] = 0, + dilation: int | tuple[int, int, int] = 1, bias: bool = True, pad_mode: str = "replicate", ) -> None: @@ -268,7 +267,7 @@ class HunyuanVideo15ResnetBlock(nn.Module): def __init__( self, in_channels: int, - out_channels: Optional[int] = None, + out_channels: int | None = None, non_linearity: str = "swish", ) -> None: super().__init__() @@ -357,7 +356,7 @@ def __init__( in_channels: int, out_channels: int, num_layers: int = 1, - downsample_out_channels: Optional[int] = None, + downsample_out_channels: int | None = None, add_temporal_downsample: int = True, ) -> None: super().__init__() @@ -406,7 +405,7 @@ def __init__( in_channels: int, out_channels: int, num_layers: int = 1, - upsample_out_channels: Optional[int] = None, + upsample_out_channels: int | None = None, add_temporal_upsample: bool = True, ) -> None: super().__init__() @@ -464,7 +463,7 @@ def __init__( self, in_channels: int = 3, out_channels: int = 64, - block_out_channels: Tuple[int, ...] = (128, 256, 512, 1024, 1024), + block_out_channels: tuple[int, ...] = (128, 256, 512, 1024, 1024), layers_per_block: int = 2, temporal_compression_ratio: int = 4, spatial_compression_ratio: int = 16, @@ -550,7 +549,7 @@ def __init__( self, in_channels: int = 32, out_channels: int = 3, - block_out_channels: Tuple[int, ...] = (1024, 1024, 512, 256, 128), + block_out_channels: tuple[int, ...] = (1024, 1024, 512, 256, 128), layers_per_block: int = 2, spatial_compression_ratio: int = 16, temporal_compression_ratio: int = 4, @@ -625,7 +624,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return hidden_states -class AutoencoderKLHunyuanVideo15(ModelMixin, ConfigMixin): +class AutoencoderKLHunyuanVideo15(ModelMixin, AutoencoderMixin, ConfigMixin): r""" A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Used for HunyuanVideo-1.5. @@ -642,7 +641,7 @@ def __init__( in_channels: int = 3, out_channels: int = 3, latent_channels: int = 32, - block_out_channels: Tuple[int] = (128, 256, 512, 1024, 1024), + block_out_channels: tuple[int] = (128, 256, 512, 1024, 1024), layers_per_block: int = 2, spatial_compression_ratio: int = 16, temporal_compression_ratio: int = 4, @@ -695,11 +694,11 @@ def __init__( def enable_tiling( self, - tile_sample_min_height: Optional[int] = None, - tile_sample_min_width: Optional[int] = None, - tile_latent_min_height: Optional[int] = None, - tile_latent_min_width: Optional[int] = None, - tile_overlap_factor: Optional[float] = None, + tile_sample_min_height: int | None = None, + tile_sample_min_width: int | None = None, + tile_latent_min_height: int | None = None, + tile_latent_min_width: int | None = None, + tile_overlap_factor: float | None = None, ) -> None: r""" Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to @@ -723,27 +722,6 @@ def enable_tiling( self.tile_latent_min_width = tile_latent_min_width or self.tile_latent_min_width self.tile_overlap_factor = tile_overlap_factor or self.tile_overlap_factor - def disable_tiling(self) -> None: - r""" - Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_tiling = False - - def enable_slicing(self) -> None: - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.use_slicing = True - - def disable_slicing(self) -> None: - r""" - Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing - decoding in one step. - """ - self.use_slicing = False - def _encode(self, x: torch.Tensor) -> torch.Tensor: _, _, _, height, width = x.shape @@ -756,7 +734,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor: @apply_forward_hook def encode( self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: r""" Encode a batch of images into latents. @@ -792,7 +770,7 @@ def _decode(self, z: torch.Tensor) -> torch.Tensor: return dec @apply_forward_hook - def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: r""" Decode a batch of images. @@ -947,8 +925,8 @@ def forward( sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[DecoderOutput, torch.Tensor]: + generator: torch.Generator | None = None, + ) -> DecoderOutput | torch.Tensor: r""" Args: sample (`torch.Tensor`): Input sample. diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_ltx.py b/src/diffusers/models/autoencoders/autoencoder_kl_ltx.py index 47f2081b7e45..a7acc105e9ec 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_ltx.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_ltx.py @@ -13,8 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Tuple, Union - import torch import torch.nn as nn @@ -34,9 +32,9 @@ def __init__( self, in_channels: int, out_channels: int, - kernel_size: Union[int, Tuple[int, int, int]] = 3, - stride: Union[int, Tuple[int, int, int]] = 1, - dilation: Union[int, Tuple[int, int, int]] = 1, + kernel_size: int | tuple[int, int, int] = 3, + stride: int | tuple[int, int, int] = 1, + dilation: int | tuple[int, int, int] = 1, groups: int = 1, padding_mode: str = "zeros", is_causal: bool = True, @@ -104,7 +102,7 @@ class LTXVideoResnetBlock3d(nn.Module): def __init__( self, in_channels: int, - out_channels: Optional[int] = None, + out_channels: int | None = None, dropout: float = 0.0, eps: float = 1e-6, elementwise_affine: bool = False, @@ -149,7 +147,7 @@ def __init__( self.scale_shift_table = nn.Parameter(torch.randn(4, in_channels) / in_channels**0.5) def forward( - self, inputs: torch.Tensor, temb: Optional[torch.Tensor] = None, generator: Optional[torch.Generator] = None + self, inputs: torch.Tensor, temb: torch.Tensor | None = None, generator: torch.Generator | None = None ) -> torch.Tensor: hidden_states = inputs @@ -201,7 +199,7 @@ def __init__( self, in_channels: int, out_channels: int, - stride: Union[int, Tuple[int, int, int]] = 1, + stride: int | tuple[int, int, int] = 1, is_causal: bool = True, padding_mode: str = "zeros", ) -> None: @@ -249,7 +247,7 @@ class LTXVideoUpsampler3d(nn.Module): def __init__( self, in_channels: int, - stride: Union[int, Tuple[int, int, int]] = 1, + stride: int | tuple[int, int, int] = 1, is_causal: bool = True, residual: bool = False, upscale_factor: int = 1, @@ -326,7 +324,7 @@ class LTXVideoDownBlock3D(nn.Module): def __init__( self, in_channels: int, - out_channels: Optional[int] = None, + out_channels: int | None = None, num_layers: int = 1, dropout: float = 0.0, resnet_eps: float = 1e-6, @@ -382,8 +380,8 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - generator: Optional[torch.Generator] = None, + temb: torch.Tensor | None = None, + generator: torch.Generator | None = None, ) -> torch.Tensor: r"""Forward method of the `LTXDownBlock3D` class.""" @@ -432,7 +430,7 @@ class LTXVideo095DownBlock3D(nn.Module): def __init__( self, in_channels: int, - out_channels: Optional[int] = None, + out_channels: int | None = None, num_layers: int = 1, dropout: float = 0.0, resnet_eps: float = 1e-6, @@ -497,8 +495,8 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - generator: Optional[torch.Generator] = None, + temb: torch.Tensor | None = None, + generator: torch.Generator | None = None, ) -> torch.Tensor: r"""Forward method of the `LTXDownBlock3D` class.""" @@ -575,8 +573,8 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - generator: Optional[torch.Generator] = None, + temb: torch.Tensor | None = None, + generator: torch.Generator | None = None, ) -> torch.Tensor: r"""Forward method of the `LTXMidBlock3D` class.""" @@ -628,7 +626,7 @@ class LTXVideoUpBlock3d(nn.Module): def __init__( self, in_channels: int, - out_channels: Optional[int] = None, + out_channels: int | None = None, num_layers: int = 1, dropout: float = 0.0, resnet_eps: float = 1e-6, @@ -696,8 +694,8 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - generator: Optional[torch.Generator] = None, + temb: torch.Tensor | None = None, + generator: torch.Generator | None = None, ) -> torch.Tensor: if self.conv_in is not None: hidden_states = self.conv_in(hidden_states, temb, generator) @@ -735,11 +733,11 @@ class LTXVideoEncoder3d(nn.Module): Number of input channels. out_channels (`int`, defaults to 128): Number of latent channels. - block_out_channels (`Tuple[int, ...]`, defaults to `(128, 256, 512, 512)`): + block_out_channels (`tuple[int, ...]`, defaults to `(128, 256, 512, 512)`): The number of output channels for each block. - spatio_temporal_scaling (`Tuple[bool, ...], defaults to `(True, True, True, False)`: + spatio_temporal_scaling (`tuple[bool, ...], defaults to `(True, True, True, False)`: Whether a block should contain spatio-temporal downscaling layers or not. - layers_per_block (`Tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`): + layers_per_block (`tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`): The number of layers per block. patch_size (`int`, defaults to `4`): The size of spatial patches. @@ -755,16 +753,16 @@ def __init__( self, in_channels: int = 3, out_channels: int = 128, - block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), - down_block_types: Tuple[str, ...] = ( + block_out_channels: tuple[int, ...] = (128, 256, 512, 512), + down_block_types: tuple[str, ...] = ( "LTXVideoDownBlock3D", "LTXVideoDownBlock3D", "LTXVideoDownBlock3D", "LTXVideoDownBlock3D", ), - spatio_temporal_scaling: Tuple[bool, ...] = (True, True, True, False), - layers_per_block: Tuple[int, ...] = (4, 3, 3, 3, 4), - downsample_type: Tuple[str, ...] = ("conv", "conv", "conv", "conv"), + spatio_temporal_scaling: tuple[bool, ...] = (True, True, True, False), + layers_per_block: tuple[int, ...] = (4, 3, 3, 3, 4), + downsample_type: tuple[str, ...] = ("conv", "conv", "conv", "conv"), patch_size: int = 4, patch_size_t: int = 1, resnet_norm_eps: float = 1e-6, @@ -888,11 +886,11 @@ class LTXVideoDecoder3d(nn.Module): Number of latent channels. out_channels (`int`, defaults to 3): Number of output channels. - block_out_channels (`Tuple[int, ...]`, defaults to `(128, 256, 512, 512)`): + block_out_channels (`tuple[int, ...]`, defaults to `(128, 256, 512, 512)`): The number of output channels for each block. - spatio_temporal_scaling (`Tuple[bool, ...], defaults to `(True, True, True, False)`: + spatio_temporal_scaling (`tuple[bool, ...], defaults to `(True, True, True, False)`: Whether a block should contain spatio-temporal upscaling layers or not. - layers_per_block (`Tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`): + layers_per_block (`tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`): The number of layers per block. patch_size (`int`, defaults to `4`): The size of spatial patches. @@ -910,17 +908,17 @@ def __init__( self, in_channels: int = 128, out_channels: int = 3, - block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), - spatio_temporal_scaling: Tuple[bool, ...] = (True, True, True, False), - layers_per_block: Tuple[int, ...] = (4, 3, 3, 3, 4), + block_out_channels: tuple[int, ...] = (128, 256, 512, 512), + spatio_temporal_scaling: tuple[bool, ...] = (True, True, True, False), + layers_per_block: tuple[int, ...] = (4, 3, 3, 3, 4), patch_size: int = 4, patch_size_t: int = 1, resnet_norm_eps: float = 1e-6, is_causal: bool = False, - inject_noise: Tuple[bool, ...] = (False, False, False, False), + inject_noise: tuple[bool, ...] = (False, False, False, False), timestep_conditioning: bool = False, - upsample_residual: Tuple[bool, ...] = (False, False, False, False), - upsample_factor: Tuple[bool, ...] = (1, 1, 1, 1), + upsample_residual: tuple[bool, ...] = (False, False, False, False), + upsample_factor: tuple[bool, ...] = (1, 1, 1, 1), ) -> None: super().__init__() @@ -989,7 +987,7 @@ def __init__( self.gradient_checkpointing = False - def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor: hidden_states = self.conv_in(hidden_states) if self.timestep_scale_multiplier is not None: @@ -1049,11 +1047,11 @@ class AutoencoderKLLTXVideo(ModelMixin, AutoencoderMixin, ConfigMixin, FromOrigi Number of output channels. latent_channels (`int`, defaults to `128`): Number of latent channels. - block_out_channels (`Tuple[int, ...]`, defaults to `(128, 256, 512, 512)`): + block_out_channels (`tuple[int, ...]`, defaults to `(128, 256, 512, 512)`): The number of output channels for each block. - spatio_temporal_scaling (`Tuple[bool, ...], defaults to `(True, True, True, False)`: + spatio_temporal_scaling (`tuple[bool, ...], defaults to `(True, True, True, False)`: Whether a block should contain spatio-temporal downscaling or not. - layers_per_block (`Tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`): + layers_per_block (`tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`): The number of layers per block. patch_size (`int`, defaults to `4`): The size of spatial patches. @@ -1082,22 +1080,22 @@ def __init__( in_channels: int = 3, out_channels: int = 3, latent_channels: int = 128, - block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), - down_block_types: Tuple[str, ...] = ( + block_out_channels: tuple[int, ...] = (128, 256, 512, 512), + down_block_types: tuple[str, ...] = ( "LTXVideoDownBlock3D", "LTXVideoDownBlock3D", "LTXVideoDownBlock3D", "LTXVideoDownBlock3D", ), - decoder_block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), - layers_per_block: Tuple[int, ...] = (4, 3, 3, 3, 4), - decoder_layers_per_block: Tuple[int, ...] = (4, 3, 3, 3, 4), - spatio_temporal_scaling: Tuple[bool, ...] = (True, True, True, False), - decoder_spatio_temporal_scaling: Tuple[bool, ...] = (True, True, True, False), - decoder_inject_noise: Tuple[bool, ...] = (False, False, False, False, False), - downsample_type: Tuple[str, ...] = ("conv", "conv", "conv", "conv"), - upsample_residual: Tuple[bool, ...] = (False, False, False, False), - upsample_factor: Tuple[int, ...] = (1, 1, 1, 1), + decoder_block_out_channels: tuple[int, ...] = (128, 256, 512, 512), + layers_per_block: tuple[int, ...] = (4, 3, 3, 3, 4), + decoder_layers_per_block: tuple[int, ...] = (4, 3, 3, 3, 4), + spatio_temporal_scaling: tuple[bool, ...] = (True, True, True, False), + decoder_spatio_temporal_scaling: tuple[bool, ...] = (True, True, True, False), + decoder_inject_noise: tuple[bool, ...] = (False, False, False, False, False), + downsample_type: tuple[str, ...] = ("conv", "conv", "conv", "conv"), + upsample_residual: tuple[bool, ...] = (False, False, False, False), + upsample_factor: tuple[int, ...] = (1, 1, 1, 1), timestep_conditioning: bool = False, patch_size: int = 4, patch_size_t: int = 1, @@ -1187,12 +1185,12 @@ def __init__( def enable_tiling( self, - tile_sample_min_height: Optional[int] = None, - tile_sample_min_width: Optional[int] = None, - tile_sample_min_num_frames: Optional[int] = None, - tile_sample_stride_height: Optional[float] = None, - tile_sample_stride_width: Optional[float] = None, - tile_sample_stride_num_frames: Optional[float] = None, + tile_sample_min_height: int | None = None, + tile_sample_min_width: int | None = None, + tile_sample_min_num_frames: int | None = None, + tile_sample_stride_height: float | None = None, + tile_sample_stride_width: float | None = None, + tile_sample_stride_num_frames: float | None = None, ) -> None: r""" Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to @@ -1235,7 +1233,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor: @apply_forward_hook def encode( self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: """ Encode a batch of images into latents. @@ -1260,8 +1258,8 @@ def encode( return AutoencoderKLOutput(latent_dist=posterior) def _decode( - self, z: torch.Tensor, temb: Optional[torch.Tensor] = None, return_dict: bool = True - ) -> Union[DecoderOutput, torch.Tensor]: + self, z: torch.Tensor, temb: torch.Tensor | None = None, return_dict: bool = True + ) -> DecoderOutput | torch.Tensor: batch_size, num_channels, num_frames, height, width = z.shape tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio @@ -1282,8 +1280,8 @@ def _decode( @apply_forward_hook def decode( - self, z: torch.Tensor, temb: Optional[torch.Tensor] = None, return_dict: bool = True - ) -> Union[DecoderOutput, torch.Tensor]: + self, z: torch.Tensor, temb: torch.Tensor | None = None, return_dict: bool = True + ) -> DecoderOutput | torch.Tensor: """ Decode a batch of images. @@ -1389,8 +1387,8 @@ def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: return enc def tiled_decode( - self, z: torch.Tensor, temb: Optional[torch.Tensor], return_dict: bool = True - ) -> Union[DecoderOutput, torch.Tensor]: + self, z: torch.Tensor, temb: torch.Tensor | None, return_dict: bool = True + ) -> DecoderOutput | torch.Tensor: r""" Decode a batch of images using a tiled decoder. @@ -1479,8 +1477,8 @@ def _temporal_tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput: return enc def _temporal_tiled_decode( - self, z: torch.Tensor, temb: Optional[torch.Tensor], return_dict: bool = True - ) -> Union[DecoderOutput, torch.Tensor]: + self, z: torch.Tensor, temb: torch.Tensor | None, return_dict: bool = True + ) -> DecoderOutput | torch.Tensor: batch_size, num_channels, num_frames, height, width = z.shape num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1 @@ -1519,11 +1517,11 @@ def _temporal_tiled_decode( def forward( self, sample: torch.Tensor, - temb: Optional[torch.Tensor] = None, + temb: torch.Tensor | None = None, sample_posterior: bool = False, return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[torch.Tensor, torch.Tensor]: + generator: torch.Generator | None = None, + ) -> torch.Tensor | torch.Tensor: x = sample posterior = self.encode(x).latent_dist if sample_posterior: diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py b/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py new file mode 100644 index 000000000000..7c04bd715c25 --- /dev/null +++ b/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py @@ -0,0 +1,1520 @@ +# Copyright 2025 The Lightricks team and The HuggingFace Team. +# All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch +import torch.nn as nn + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin +from ...utils.accelerate_utils import apply_forward_hook +from ..activations import get_activation +from ..embeddings import PixArtAlphaCombinedTimestepSizeEmbeddings +from ..modeling_outputs import AutoencoderKLOutput +from ..modeling_utils import ModelMixin +from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution + + +class PerChannelRMSNorm(nn.Module): + """ + Per-pixel (per-location) RMS normalization layer. + + For each element along the chosen dimension, this layer normalizes the tensor by the root-mean-square of its values + across that dimension: + + y = x / sqrt(mean(x^2, dim=dim, keepdim=True) + eps) + """ + + def __init__(self, channel_dim: int = 1, eps: float = 1e-8) -> None: + """ + Args: + dim: Dimension along which to compute the RMS (typically channels). + eps: Small constant added for numerical stability. + """ + super().__init__() + self.channel_dim = channel_dim + self.eps = eps + + def forward(self, x: torch.Tensor, channel_dim: int | None = None) -> torch.Tensor: + """ + Apply RMS normalization along the configured dimension. + """ + channel_dim = channel_dim or self.channel_dim + # Compute mean of squared values along `dim`, keep dimensions for broadcasting. + mean_sq = torch.mean(x**2, dim=self.channel_dim, keepdim=True) + # Normalize by the root-mean-square (RMS). + rms = torch.sqrt(mean_sq + self.eps) + return x / rms + + +# Like LTXCausalConv3d, but whether causal inference is performed can be specified at runtime +class LTX2VideoCausalConv3d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int, int] = 3, + stride: int | tuple[int, int, int] = 1, + dilation: int | tuple[int, int, int] = 1, + groups: int = 1, + spatial_padding_mode: str = "zeros", + ): + super().__init__() + + self.in_channels = in_channels + self.out_channels = out_channels + self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size, kernel_size) + + dilation = dilation if isinstance(dilation, tuple) else (dilation, 1, 1) + stride = stride if isinstance(stride, tuple) else (stride, stride, stride) + height_pad = self.kernel_size[1] // 2 + width_pad = self.kernel_size[2] // 2 + padding = (0, height_pad, width_pad) + + self.conv = nn.Conv3d( + in_channels, + out_channels, + self.kernel_size, + stride=stride, + dilation=dilation, + groups=groups, + padding=padding, + padding_mode=spatial_padding_mode, + ) + + def forward(self, hidden_states: torch.Tensor, causal: bool = True) -> torch.Tensor: + time_kernel_size = self.kernel_size[0] + + if causal: + pad_left = hidden_states[:, :, :1, :, :].repeat((1, 1, time_kernel_size - 1, 1, 1)) + hidden_states = torch.concatenate([pad_left, hidden_states], dim=2) + else: + pad_left = hidden_states[:, :, :1, :, :].repeat((1, 1, (time_kernel_size - 1) // 2, 1, 1)) + pad_right = hidden_states[:, :, -1:, :, :].repeat((1, 1, (time_kernel_size - 1) // 2, 1, 1)) + hidden_states = torch.concatenate([pad_left, hidden_states, pad_right], dim=2) + + hidden_states = self.conv(hidden_states) + return hidden_states + + +# Like LTXVideoResnetBlock3d, but uses new causal Conv3d, normal Conv3d for the conv_shortcut, and the spatial padding +# mode is configurable +class LTX2VideoResnetBlock3d(nn.Module): + r""" + A 3D ResNet block used in the LTX 2.0 audiovisual model. + + Args: + in_channels (`int`): + Number of input channels. + out_channels (`int`, *optional*): + Number of output channels. If None, defaults to `in_channels`. + dropout (`float`, defaults to `0.0`): + Dropout rate. + eps (`float`, defaults to `1e-6`): + Epsilon value for normalization layers. + elementwise_affine (`bool`, defaults to `False`): + Whether to enable elementwise affinity in the normalization layers. + non_linearity (`str`, defaults to `"swish"`): + Activation function to use. + conv_shortcut (bool, defaults to `False`): + Whether or not to use a convolution shortcut. + """ + + def __init__( + self, + in_channels: int, + out_channels: int | None = None, + dropout: float = 0.0, + eps: float = 1e-6, + elementwise_affine: bool = False, + non_linearity: str = "swish", + inject_noise: bool = False, + timestep_conditioning: bool = False, + spatial_padding_mode: str = "zeros", + ) -> None: + super().__init__() + + out_channels = out_channels or in_channels + + self.nonlinearity = get_activation(non_linearity) + + self.norm1 = PerChannelRMSNorm() + self.conv1 = LTX2VideoCausalConv3d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + spatial_padding_mode=spatial_padding_mode, + ) + + self.norm2 = PerChannelRMSNorm() + self.dropout = nn.Dropout(dropout) + self.conv2 = LTX2VideoCausalConv3d( + in_channels=out_channels, + out_channels=out_channels, + kernel_size=3, + spatial_padding_mode=spatial_padding_mode, + ) + + self.norm3 = None + self.conv_shortcut = None + if in_channels != out_channels: + self.norm3 = nn.LayerNorm(in_channels, eps=eps, elementwise_affine=True, bias=True) + # LTX 2.0 uses a normal nn.Conv3d here rather than LTXVideoCausalConv3d + self.conv_shortcut = nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1) + + self.per_channel_scale1 = None + self.per_channel_scale2 = None + if inject_noise: + self.per_channel_scale1 = nn.Parameter(torch.zeros(in_channels, 1, 1)) + self.per_channel_scale2 = nn.Parameter(torch.zeros(in_channels, 1, 1)) + + self.scale_shift_table = None + if timestep_conditioning: + self.scale_shift_table = nn.Parameter(torch.randn(4, in_channels) / in_channels**0.5) + + def forward( + self, + inputs: torch.Tensor, + temb: torch.Tensor | None = None, + generator: torch.Generator | None = None, + causal: bool = True, + ) -> torch.Tensor: + hidden_states = inputs + + hidden_states = self.norm1(hidden_states) + + if self.scale_shift_table is not None: + temb = temb.unflatten(1, (4, -1)) + self.scale_shift_table[None, ..., None, None, None] + shift_1, scale_1, shift_2, scale_2 = temb.unbind(dim=1) + hidden_states = hidden_states * (1 + scale_1) + shift_1 + + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.conv1(hidden_states, causal=causal) + + if self.per_channel_scale1 is not None: + spatial_shape = hidden_states.shape[-2:] + spatial_noise = torch.randn( + spatial_shape, generator=generator, device=hidden_states.device, dtype=hidden_states.dtype + )[None] + hidden_states = hidden_states + (spatial_noise * self.per_channel_scale1)[None, :, None, ...] + + hidden_states = self.norm2(hidden_states) + + if self.scale_shift_table is not None: + hidden_states = hidden_states * (1 + scale_2) + shift_2 + + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = self.conv2(hidden_states, causal=causal) + + if self.per_channel_scale2 is not None: + spatial_shape = hidden_states.shape[-2:] + spatial_noise = torch.randn( + spatial_shape, generator=generator, device=hidden_states.device, dtype=hidden_states.dtype + )[None] + hidden_states = hidden_states + (spatial_noise * self.per_channel_scale2)[None, :, None, ...] + + if self.norm3 is not None: + inputs = self.norm3(inputs.movedim(1, -1)).movedim(-1, 1) + + if self.conv_shortcut is not None: + inputs = self.conv_shortcut(inputs) + + hidden_states = hidden_states + inputs + return hidden_states + + +# Like LTX 1.0 LTXVideoDownsampler3d, but uses new causal Conv3d +class LTXVideoDownsampler3d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + stride: int | tuple[int, int, int] = 1, + spatial_padding_mode: str = "zeros", + ) -> None: + super().__init__() + + self.stride = stride if isinstance(stride, tuple) else (stride, stride, stride) + self.group_size = (in_channels * stride[0] * stride[1] * stride[2]) // out_channels + + out_channels = out_channels // (self.stride[0] * self.stride[1] * self.stride[2]) + + self.conv = LTX2VideoCausalConv3d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=1, + spatial_padding_mode=spatial_padding_mode, + ) + + def forward(self, hidden_states: torch.Tensor, causal: bool = True) -> torch.Tensor: + hidden_states = torch.cat([hidden_states[:, :, : self.stride[0] - 1], hidden_states], dim=2) + + residual = ( + hidden_states.unflatten(4, (-1, self.stride[2])) + .unflatten(3, (-1, self.stride[1])) + .unflatten(2, (-1, self.stride[0])) + ) + residual = residual.permute(0, 1, 3, 5, 7, 2, 4, 6).flatten(1, 4) + residual = residual.unflatten(1, (-1, self.group_size)) + residual = residual.mean(dim=2) + + hidden_states = self.conv(hidden_states, causal=causal) + hidden_states = ( + hidden_states.unflatten(4, (-1, self.stride[2])) + .unflatten(3, (-1, self.stride[1])) + .unflatten(2, (-1, self.stride[0])) + ) + hidden_states = hidden_states.permute(0, 1, 3, 5, 7, 2, 4, 6).flatten(1, 4) + hidden_states = hidden_states + residual + + return hidden_states + + +# Like LTX 1.0 LTXVideoUpsampler3d, but uses new causal Conv3d +class LTXVideoUpsampler3d(nn.Module): + def __init__( + self, + in_channels: int, + stride: int | tuple[int, int, int] = 1, + residual: bool = False, + upscale_factor: int = 1, + spatial_padding_mode: str = "zeros", + ) -> None: + super().__init__() + + self.stride = stride if isinstance(stride, tuple) else (stride, stride, stride) + self.residual = residual + self.upscale_factor = upscale_factor + + out_channels = (in_channels * stride[0] * stride[1] * stride[2]) // upscale_factor + + self.conv = LTX2VideoCausalConv3d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=1, + spatial_padding_mode=spatial_padding_mode, + ) + + def forward(self, hidden_states: torch.Tensor, causal: bool = True) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = hidden_states.shape + + if self.residual: + residual = hidden_states.reshape( + batch_size, -1, self.stride[0], self.stride[1], self.stride[2], num_frames, height, width + ) + residual = residual.permute(0, 1, 5, 2, 6, 3, 7, 4).flatten(6, 7).flatten(4, 5).flatten(2, 3) + repeats = (self.stride[0] * self.stride[1] * self.stride[2]) // self.upscale_factor + residual = residual.repeat(1, repeats, 1, 1, 1) + residual = residual[:, :, self.stride[0] - 1 :] + + hidden_states = self.conv(hidden_states, causal=causal) + hidden_states = hidden_states.reshape( + batch_size, -1, self.stride[0], self.stride[1], self.stride[2], num_frames, height, width + ) + hidden_states = hidden_states.permute(0, 1, 5, 2, 6, 3, 7, 4).flatten(6, 7).flatten(4, 5).flatten(2, 3) + hidden_states = hidden_states[:, :, self.stride[0] - 1 :] + + if self.residual: + hidden_states = hidden_states + residual + + return hidden_states + + +# Like LTX 1.0 LTXVideo095DownBlock3D, but with the updated LTX2VideoResnetBlock3d +class LTX2VideoDownBlock3D(nn.Module): + r""" + Down block used in the LTXVideo model. + + Args: + in_channels (`int`): + Number of input channels. + out_channels (`int`, *optional*): + Number of output channels. If None, defaults to `in_channels`. + num_layers (`int`, defaults to `1`): + Number of resnet layers. + dropout (`float`, defaults to `0.0`): + Dropout rate. + resnet_eps (`float`, defaults to `1e-6`): + Epsilon value for normalization layers. + resnet_act_fn (`str`, defaults to `"swish"`): + Activation function to use. + spatio_temporal_scale (`bool`, defaults to `True`): + Whether or not to use a downsampling layer. If not used, output dimension would be same as input dimension. + Whether or not to downsample across temporal dimension. + is_causal (`bool`, defaults to `True`): + Whether this layer behaves causally (future frames depend only on past frames) or not. + """ + + _supports_gradient_checkpointing = True + + def __init__( + self, + in_channels: int, + out_channels: int | None = None, + num_layers: int = 1, + dropout: float = 0.0, + resnet_eps: float = 1e-6, + resnet_act_fn: str = "swish", + spatio_temporal_scale: bool = True, + downsample_type: str = "conv", + spatial_padding_mode: str = "zeros", + ): + super().__init__() + + out_channels = out_channels or in_channels + + resnets = [] + for _ in range(num_layers): + resnets.append( + LTX2VideoResnetBlock3d( + in_channels=in_channels, + out_channels=in_channels, + dropout=dropout, + eps=resnet_eps, + non_linearity=resnet_act_fn, + spatial_padding_mode=spatial_padding_mode, + ) + ) + self.resnets = nn.ModuleList(resnets) + + self.downsamplers = None + if spatio_temporal_scale: + self.downsamplers = nn.ModuleList() + + if downsample_type == "conv": + self.downsamplers.append( + LTX2VideoCausalConv3d( + in_channels=in_channels, + out_channels=in_channels, + kernel_size=3, + stride=(2, 2, 2), + spatial_padding_mode=spatial_padding_mode, + ) + ) + elif downsample_type == "spatial": + self.downsamplers.append( + LTXVideoDownsampler3d( + in_channels=in_channels, + out_channels=out_channels, + stride=(1, 2, 2), + spatial_padding_mode=spatial_padding_mode, + ) + ) + elif downsample_type == "temporal": + self.downsamplers.append( + LTXVideoDownsampler3d( + in_channels=in_channels, + out_channels=out_channels, + stride=(2, 1, 1), + spatial_padding_mode=spatial_padding_mode, + ) + ) + elif downsample_type == "spatiotemporal": + self.downsamplers.append( + LTXVideoDownsampler3d( + in_channels=in_channels, + out_channels=out_channels, + stride=(2, 2, 2), + spatial_padding_mode=spatial_padding_mode, + ) + ) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor | None = None, + generator: torch.Generator | None = None, + causal: bool = True, + ) -> torch.Tensor: + r"""Forward method of the `LTXDownBlock3D` class.""" + + for i, resnet in enumerate(self.resnets): + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb, generator, causal) + else: + hidden_states = resnet(hidden_states, temb, generator, causal=causal) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, causal=causal) + + return hidden_states + + +# Adapted from diffusers.models.autoencoders.autoencoder_kl_cogvideox.CogVideoMidBlock3d +# Like LTX 1.0 LTXVideoMidBlock3d, but with the updated LTX2VideoResnetBlock3d +class LTX2VideoMidBlock3d(nn.Module): + r""" + A middle block used in the LTXVideo model. + + Args: + in_channels (`int`): + Number of input channels. + num_layers (`int`, defaults to `1`): + Number of resnet layers. + dropout (`float`, defaults to `0.0`): + Dropout rate. + resnet_eps (`float`, defaults to `1e-6`): + Epsilon value for normalization layers. + resnet_act_fn (`str`, defaults to `"swish"`): + Activation function to use. + is_causal (`bool`, defaults to `True`): + Whether this layer behaves causally (future frames depend only on past frames) or not. + """ + + _supports_gradient_checkpointing = True + + def __init__( + self, + in_channels: int, + num_layers: int = 1, + dropout: float = 0.0, + resnet_eps: float = 1e-6, + resnet_act_fn: str = "swish", + inject_noise: bool = False, + timestep_conditioning: bool = False, + spatial_padding_mode: str = "zeros", + ) -> None: + super().__init__() + + self.time_embedder = None + if timestep_conditioning: + self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(in_channels * 4, 0) + + resnets = [] + for _ in range(num_layers): + resnets.append( + LTX2VideoResnetBlock3d( + in_channels=in_channels, + out_channels=in_channels, + dropout=dropout, + eps=resnet_eps, + non_linearity=resnet_act_fn, + inject_noise=inject_noise, + timestep_conditioning=timestep_conditioning, + spatial_padding_mode=spatial_padding_mode, + ) + ) + self.resnets = nn.ModuleList(resnets) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor | None = None, + generator: torch.Generator | None = None, + causal: bool = True, + ) -> torch.Tensor: + r"""Forward method of the `LTXMidBlock3D` class.""" + + if self.time_embedder is not None: + temb = self.time_embedder( + timestep=temb.flatten(), + resolution=None, + aspect_ratio=None, + batch_size=hidden_states.size(0), + hidden_dtype=hidden_states.dtype, + ) + temb = temb.view(hidden_states.size(0), -1, 1, 1, 1) + + for i, resnet in enumerate(self.resnets): + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb, generator, causal) + else: + hidden_states = resnet(hidden_states, temb, generator, causal=causal) + + return hidden_states + + +# Like LTXVideoUpBlock3d but with no conv_in and the updated LTX2VideoResnetBlock3d +class LTX2VideoUpBlock3d(nn.Module): + r""" + Up block used in the LTXVideo model. + + Args: + in_channels (`int`): + Number of input channels. + out_channels (`int`, *optional*): + Number of output channels. If None, defaults to `in_channels`. + num_layers (`int`, defaults to `1`): + Number of resnet layers. + dropout (`float`, defaults to `0.0`): + Dropout rate. + resnet_eps (`float`, defaults to `1e-6`): + Epsilon value for normalization layers. + resnet_act_fn (`str`, defaults to `"swish"`): + Activation function to use. + spatio_temporal_scale (`bool`, defaults to `True`): + Whether or not to use a downsampling layer. If not used, output dimension would be same as input dimension. + Whether or not to downsample across temporal dimension. + is_causal (`bool`, defaults to `True`): + Whether this layer behaves causally (future frames depend only on past frames) or not. + """ + + _supports_gradient_checkpointing = True + + def __init__( + self, + in_channels: int, + out_channels: int | None = None, + num_layers: int = 1, + dropout: float = 0.0, + resnet_eps: float = 1e-6, + resnet_act_fn: str = "swish", + spatio_temporal_scale: bool = True, + inject_noise: bool = False, + timestep_conditioning: bool = False, + upsample_residual: bool = False, + upscale_factor: int = 1, + spatial_padding_mode: str = "zeros", + ): + super().__init__() + + out_channels = out_channels or in_channels + + self.time_embedder = None + if timestep_conditioning: + self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(in_channels * 4, 0) + + self.conv_in = None + if in_channels != out_channels: + self.conv_in = LTX2VideoResnetBlock3d( + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + eps=resnet_eps, + non_linearity=resnet_act_fn, + inject_noise=inject_noise, + timestep_conditioning=timestep_conditioning, + spatial_padding_mode=spatial_padding_mode, + ) + + self.upsamplers = None + if spatio_temporal_scale: + self.upsamplers = nn.ModuleList( + [ + LTXVideoUpsampler3d( + out_channels * upscale_factor, + stride=(2, 2, 2), + residual=upsample_residual, + upscale_factor=upscale_factor, + spatial_padding_mode=spatial_padding_mode, + ) + ] + ) + + resnets = [] + for _ in range(num_layers): + resnets.append( + LTX2VideoResnetBlock3d( + in_channels=out_channels, + out_channels=out_channels, + dropout=dropout, + eps=resnet_eps, + non_linearity=resnet_act_fn, + inject_noise=inject_noise, + timestep_conditioning=timestep_conditioning, + spatial_padding_mode=spatial_padding_mode, + ) + ) + self.resnets = nn.ModuleList(resnets) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor | None = None, + generator: torch.Generator | None = None, + causal: bool = True, + ) -> torch.Tensor: + if self.conv_in is not None: + hidden_states = self.conv_in(hidden_states, temb, generator, causal=causal) + + if self.time_embedder is not None: + temb = self.time_embedder( + timestep=temb.flatten(), + resolution=None, + aspect_ratio=None, + batch_size=hidden_states.size(0), + hidden_dtype=hidden_states.dtype, + ) + temb = temb.view(hidden_states.size(0), -1, 1, 1, 1) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, causal=causal) + + for i, resnet in enumerate(self.resnets): + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb, generator, causal) + else: + hidden_states = resnet(hidden_states, temb, generator, causal=causal) + + return hidden_states + + +# Like LTX 1.0 LTXVideoEncoder3d but with different default args - the spatiotemporal downsampling pattern is +# different, as is the layers_per_block (the 2.0 VAE is bigger) +class LTX2VideoEncoder3d(nn.Module): + r""" + The `LTXVideoEncoder3d` layer of a variational autoencoder that encodes input video samples to its latent + representation. + + Args: + in_channels (`int`, defaults to 3): + Number of input channels. + out_channels (`int`, defaults to 128): + Number of latent channels. + block_out_channels (`tuple[int, ...]`, defaults to `(256, 512, 1024, 2048)`): + The number of output channels for each block. + spatio_temporal_scaling (`tuple[bool, ...], defaults to `(True, True, True, True)`: + Whether a block should contain spatio-temporal downscaling layers or not. + layers_per_block (`tuple[int, ...]`, defaults to `(4, 6, 6, 2, 2)`): + The number of layers per block. + downsample_type (`tuple[str, ...]`, defaults to `("spatial", "temporal", "spatiotemporal", "spatiotemporal")`): + The spatiotemporal downsampling pattern per block. Per-layer values can be + - `"spatial"` (downsample spatial dims by 2x) + - `"temporal"` (downsample temporal dim by 2x) + - `"spatiotemporal"` (downsample both spatial and temporal dims by 2x) + patch_size (`int`, defaults to `4`): + The size of spatial patches. + patch_size_t (`int`, defaults to `1`): + The size of temporal patches. + resnet_norm_eps (`float`, defaults to `1e-6`): + Epsilon value for ResNet normalization layers. + is_causal (`bool`, defaults to `True`): + Whether this layer behaves causally (future frames depend only on past frames) or not. + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 128, + block_out_channels: tuple[int, ...] = (256, 512, 1024, 2048), + down_block_types: tuple[str, ...] = ( + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + ), + spatio_temporal_scaling: tuple[bool, ...] = (True, True, True, True), + layers_per_block: tuple[int, ...] = (4, 6, 6, 2, 2), + downsample_type: tuple[str, ...] = ("spatial", "temporal", "spatiotemporal", "spatiotemporal"), + patch_size: int = 4, + patch_size_t: int = 1, + resnet_norm_eps: float = 1e-6, + is_causal: bool = True, + spatial_padding_mode: str = "zeros", + ): + super().__init__() + + self.patch_size = patch_size + self.patch_size_t = patch_size_t + self.in_channels = in_channels * patch_size**2 + self.is_causal = is_causal + + output_channel = out_channels + + self.conv_in = LTX2VideoCausalConv3d( + in_channels=self.in_channels, + out_channels=output_channel, + kernel_size=3, + stride=1, + spatial_padding_mode=spatial_padding_mode, + ) + + # down blocks + num_block_out_channels = len(block_out_channels) + self.down_blocks = nn.ModuleList([]) + for i in range(num_block_out_channels): + input_channel = output_channel + output_channel = block_out_channels[i] + + if down_block_types[i] == "LTX2VideoDownBlock3D": + down_block = LTX2VideoDownBlock3D( + in_channels=input_channel, + out_channels=output_channel, + num_layers=layers_per_block[i], + resnet_eps=resnet_norm_eps, + spatio_temporal_scale=spatio_temporal_scaling[i], + downsample_type=downsample_type[i], + spatial_padding_mode=spatial_padding_mode, + ) + else: + raise ValueError(f"Unknown down block type: {down_block_types[i]}") + + self.down_blocks.append(down_block) + + # mid block + self.mid_block = LTX2VideoMidBlock3d( + in_channels=output_channel, + num_layers=layers_per_block[-1], + resnet_eps=resnet_norm_eps, + spatial_padding_mode=spatial_padding_mode, + ) + + # out + self.norm_out = PerChannelRMSNorm() + self.conv_act = nn.SiLU() + self.conv_out = LTX2VideoCausalConv3d( + in_channels=output_channel, + out_channels=out_channels + 1, + kernel_size=3, + stride=1, + spatial_padding_mode=spatial_padding_mode, + ) + + self.gradient_checkpointing = False + + def forward(self, hidden_states: torch.Tensor, causal: bool | None = None) -> torch.Tensor: + r"""The forward method of the `LTXVideoEncoder3d` class.""" + + p = self.patch_size + p_t = self.patch_size_t + + batch_size, num_channels, num_frames, height, width = hidden_states.shape + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p + post_patch_width = width // p + causal = causal or self.is_causal + + hidden_states = hidden_states.reshape( + batch_size, num_channels, post_patch_num_frames, p_t, post_patch_height, p, post_patch_width, p + ) + # Thanks for driving me insane with the weird patching order :( + hidden_states = hidden_states.permute(0, 1, 3, 7, 5, 2, 4, 6).flatten(1, 4) + hidden_states = self.conv_in(hidden_states, causal=causal) + + if torch.is_grad_enabled() and self.gradient_checkpointing: + for down_block in self.down_blocks: + hidden_states = self._gradient_checkpointing_func(down_block, hidden_states, None, None, causal) + + hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states, None, None, causal) + else: + for down_block in self.down_blocks: + hidden_states = down_block(hidden_states, causal=causal) + + hidden_states = self.mid_block(hidden_states, causal=causal) + + hidden_states = self.norm_out(hidden_states) + hidden_states = self.conv_act(hidden_states) + hidden_states = self.conv_out(hidden_states, causal=causal) + + last_channel = hidden_states[:, -1:] + last_channel = last_channel.repeat(1, hidden_states.size(1) - 2, 1, 1, 1) + hidden_states = torch.cat([hidden_states, last_channel], dim=1) + + return hidden_states + + +# Like LTX 1.0 LTXVideoDecoder3d, but has only 3 symmetric up blocks which are causal and residual with upsample_factor 2 +class LTX2VideoDecoder3d(nn.Module): + r""" + The `LTXVideoDecoder3d` layer of a variational autoencoder that decodes its latent representation into an output + sample. + + Args: + in_channels (`int`, defaults to 128): + Number of latent channels. + out_channels (`int`, defaults to 3): + Number of output channels. + block_out_channels (`tuple[int, ...]`, defaults to `(128, 256, 512, 512)`): + The number of output channels for each block. + spatio_temporal_scaling (`tuple[bool, ...], defaults to `(True, True, True, False)`: + Whether a block should contain spatio-temporal upscaling layers or not. + layers_per_block (`tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`): + The number of layers per block. + patch_size (`int`, defaults to `4`): + The size of spatial patches. + patch_size_t (`int`, defaults to `1`): + The size of temporal patches. + resnet_norm_eps (`float`, defaults to `1e-6`): + Epsilon value for ResNet normalization layers. + is_causal (`bool`, defaults to `False`): + Whether this layer behaves causally (future frames depend only on past frames) or not. + timestep_conditioning (`bool`, defaults to `False`): + Whether to condition the model on timesteps. + """ + + def __init__( + self, + in_channels: int = 128, + out_channels: int = 3, + block_out_channels: tuple[int, ...] = (256, 512, 1024), + spatio_temporal_scaling: tuple[bool, ...] = (True, True, True), + layers_per_block: tuple[int, ...] = (5, 5, 5, 5), + patch_size: int = 4, + patch_size_t: int = 1, + resnet_norm_eps: float = 1e-6, + is_causal: bool = False, + inject_noise: tuple[bool, ...] = (False, False, False), + timestep_conditioning: bool = False, + upsample_residual: tuple[bool, ...] = (True, True, True), + upsample_factor: tuple[bool, ...] = (2, 2, 2), + spatial_padding_mode: str = "reflect", + ) -> None: + super().__init__() + + self.patch_size = patch_size + self.patch_size_t = patch_size_t + self.out_channels = out_channels * patch_size**2 + self.is_causal = is_causal + + block_out_channels = tuple(reversed(block_out_channels)) + spatio_temporal_scaling = tuple(reversed(spatio_temporal_scaling)) + layers_per_block = tuple(reversed(layers_per_block)) + inject_noise = tuple(reversed(inject_noise)) + upsample_residual = tuple(reversed(upsample_residual)) + upsample_factor = tuple(reversed(upsample_factor)) + output_channel = block_out_channels[0] + + self.conv_in = LTX2VideoCausalConv3d( + in_channels=in_channels, + out_channels=output_channel, + kernel_size=3, + stride=1, + spatial_padding_mode=spatial_padding_mode, + ) + + self.mid_block = LTX2VideoMidBlock3d( + in_channels=output_channel, + num_layers=layers_per_block[0], + resnet_eps=resnet_norm_eps, + inject_noise=inject_noise[0], + timestep_conditioning=timestep_conditioning, + spatial_padding_mode=spatial_padding_mode, + ) + + # up blocks + num_block_out_channels = len(block_out_channels) + self.up_blocks = nn.ModuleList([]) + for i in range(num_block_out_channels): + input_channel = output_channel // upsample_factor[i] + output_channel = block_out_channels[i] // upsample_factor[i] + + up_block = LTX2VideoUpBlock3d( + in_channels=input_channel, + out_channels=output_channel, + num_layers=layers_per_block[i + 1], + resnet_eps=resnet_norm_eps, + spatio_temporal_scale=spatio_temporal_scaling[i], + inject_noise=inject_noise[i + 1], + timestep_conditioning=timestep_conditioning, + upsample_residual=upsample_residual[i], + upscale_factor=upsample_factor[i], + spatial_padding_mode=spatial_padding_mode, + ) + + self.up_blocks.append(up_block) + + # out + self.norm_out = PerChannelRMSNorm() + self.conv_act = nn.SiLU() + self.conv_out = LTX2VideoCausalConv3d( + in_channels=output_channel, + out_channels=self.out_channels, + kernel_size=3, + stride=1, + spatial_padding_mode=spatial_padding_mode, + ) + + # timestep embedding + self.time_embedder = None + self.scale_shift_table = None + self.timestep_scale_multiplier = None + if timestep_conditioning: + self.timestep_scale_multiplier = nn.Parameter(torch.tensor(1000.0, dtype=torch.float32)) + self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(output_channel * 2, 0) + self.scale_shift_table = nn.Parameter(torch.randn(2, output_channel) / output_channel**0.5) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor | None = None, + causal: bool | None = None, + ) -> torch.Tensor: + causal = causal or self.is_causal + + hidden_states = self.conv_in(hidden_states, causal=causal) + + if self.timestep_scale_multiplier is not None: + temb = temb * self.timestep_scale_multiplier + + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states, temb, None, causal) + + for up_block in self.up_blocks: + hidden_states = self._gradient_checkpointing_func(up_block, hidden_states, temb, None, causal) + else: + hidden_states = self.mid_block(hidden_states, temb, causal=causal) + + for up_block in self.up_blocks: + hidden_states = up_block(hidden_states, temb, causal=causal) + + hidden_states = self.norm_out(hidden_states) + + if self.time_embedder is not None: + temb = self.time_embedder( + timestep=temb.flatten(), + resolution=None, + aspect_ratio=None, + batch_size=hidden_states.size(0), + hidden_dtype=hidden_states.dtype, + ) + temb = temb.view(hidden_states.size(0), -1, 1, 1, 1).unflatten(1, (2, -1)) + temb = temb + self.scale_shift_table[None, ..., None, None, None] + shift, scale = temb.unbind(dim=1) + hidden_states = hidden_states * (1 + scale) + shift + + hidden_states = self.conv_act(hidden_states) + hidden_states = self.conv_out(hidden_states, causal=causal) + + p = self.patch_size + p_t = self.patch_size_t + + batch_size, num_channels, num_frames, height, width = hidden_states.shape + hidden_states = hidden_states.reshape(batch_size, -1, p_t, p, p, num_frames, height, width) + hidden_states = hidden_states.permute(0, 1, 5, 2, 6, 4, 7, 3).flatten(6, 7).flatten(4, 5).flatten(2, 3) + + return hidden_states + + +class AutoencoderKLLTX2Video(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): + r""" + A VAE model with KL loss for encoding images into latents and decoding latent representations into images. Used in + [LTX-2](https://huggingface.co/Lightricks/LTX-2). + + This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented + for all models (such as downloading or saving). + + Args: + in_channels (`int`, defaults to `3`): + Number of input channels. + out_channels (`int`, defaults to `3`): + Number of output channels. + latent_channels (`int`, defaults to `128`): + Number of latent channels. + block_out_channels (`tuple[int, ...]`, defaults to `(128, 256, 512, 512)`): + The number of output channels for each block. + spatio_temporal_scaling (`tuple[bool, ...], defaults to `(True, True, True, False)`: + Whether a block should contain spatio-temporal downscaling or not. + layers_per_block (`tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`): + The number of layers per block. + patch_size (`int`, defaults to `4`): + The size of spatial patches. + patch_size_t (`int`, defaults to `1`): + The size of temporal patches. + resnet_norm_eps (`float`, defaults to `1e-6`): + Epsilon value for ResNet normalization layers. + scaling_factor (`float`, *optional*, defaults to `1.0`): + The component-wise standard deviation of the trained latent space computed using the first batch of the + training set. This is used to scale the latent space to have unit variance when training the diffusion + model. The latents are scaled with the formula `z = z * scaling_factor` before being passed to the + diffusion model. When decoding, the latents are scaled back to the original scale with the formula: `z = 1 + / scaling_factor * z`. For more details, refer to sections 4.3.2 and D.1 of the [High-Resolution Image + Synthesis with Latent Diffusion Models](https://huggingface.co/papers/2112.10752) paper. + encoder_causal (`bool`, defaults to `True`): + Whether the encoder should behave causally (future frames depend only on past frames) or not. + decoder_causal (`bool`, defaults to `False`): + Whether the decoder should behave causally (future frames depend only on past frames) or not. + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + latent_channels: int = 128, + block_out_channels: tuple[int, ...] = (256, 512, 1024, 2048), + down_block_types: tuple[str, ...] = ( + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + ), + decoder_block_out_channels: tuple[int, ...] = (256, 512, 1024), + layers_per_block: tuple[int, ...] = (4, 6, 6, 2, 2), + decoder_layers_per_block: tuple[int, ...] = (5, 5, 5, 5), + spatio_temporal_scaling: tuple[bool, ...] = (True, True, True, True), + decoder_spatio_temporal_scaling: tuple[bool, ...] = (True, True, True), + decoder_inject_noise: tuple[bool, ...] = (False, False, False, False), + downsample_type: tuple[str, ...] = ("spatial", "temporal", "spatiotemporal", "spatiotemporal"), + upsample_residual: tuple[bool, ...] = (True, True, True), + upsample_factor: tuple[int, ...] = (2, 2, 2), + timestep_conditioning: bool = False, + patch_size: int = 4, + patch_size_t: int = 1, + resnet_norm_eps: float = 1e-6, + scaling_factor: float = 1.0, + encoder_causal: bool = True, + decoder_causal: bool = True, + encoder_spatial_padding_mode: str = "zeros", + decoder_spatial_padding_mode: str = "reflect", + spatial_compression_ratio: int = None, + temporal_compression_ratio: int = None, + ) -> None: + super().__init__() + + self.encoder = LTX2VideoEncoder3d( + in_channels=in_channels, + out_channels=latent_channels, + block_out_channels=block_out_channels, + down_block_types=down_block_types, + spatio_temporal_scaling=spatio_temporal_scaling, + layers_per_block=layers_per_block, + downsample_type=downsample_type, + patch_size=patch_size, + patch_size_t=patch_size_t, + resnet_norm_eps=resnet_norm_eps, + is_causal=encoder_causal, + spatial_padding_mode=encoder_spatial_padding_mode, + ) + self.decoder = LTX2VideoDecoder3d( + in_channels=latent_channels, + out_channels=out_channels, + block_out_channels=decoder_block_out_channels, + spatio_temporal_scaling=decoder_spatio_temporal_scaling, + layers_per_block=decoder_layers_per_block, + patch_size=patch_size, + patch_size_t=patch_size_t, + resnet_norm_eps=resnet_norm_eps, + is_causal=decoder_causal, + timestep_conditioning=timestep_conditioning, + inject_noise=decoder_inject_noise, + upsample_residual=upsample_residual, + upsample_factor=upsample_factor, + spatial_padding_mode=decoder_spatial_padding_mode, + ) + + latents_mean = torch.zeros((latent_channels,), requires_grad=False) + latents_std = torch.ones((latent_channels,), requires_grad=False) + self.register_buffer("latents_mean", latents_mean, persistent=True) + self.register_buffer("latents_std", latents_std, persistent=True) + + self.spatial_compression_ratio = ( + patch_size * 2 ** sum(spatio_temporal_scaling) + if spatial_compression_ratio is None + else spatial_compression_ratio + ) + self.temporal_compression_ratio = ( + patch_size_t * 2 ** sum(spatio_temporal_scaling) + if temporal_compression_ratio is None + else temporal_compression_ratio + ) + + # When decoding a batch of video latents at a time, one can save memory by slicing across the batch dimension + # to perform decoding of a single video latent at a time. + self.use_slicing = False + + # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent + # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the + # intermediate tiles together, the memory requirement can be lowered. + self.use_tiling = False + + # When decoding temporally long video latents, the memory requirement is very high. By decoding latent frames + # at a fixed frame batch size (based on `self.num_latent_frames_batch_sizes`), the memory requirement can be lowered. + self.use_framewise_encoding = False + self.use_framewise_decoding = False + + # This can be configured based on the amount of GPU memory available. + # `16` for sample frames and `2` for latent frames are sensible defaults for consumer GPUs. + # Setting it to higher values results in higher memory usage. + self.num_sample_frames_batch_size = 16 + self.num_latent_frames_batch_size = 2 + + # The minimal tile height and width for spatial tiling to be used + self.tile_sample_min_height = 512 + self.tile_sample_min_width = 512 + self.tile_sample_min_num_frames = 16 + + # The minimal distance between two spatial tiles + self.tile_sample_stride_height = 448 + self.tile_sample_stride_width = 448 + self.tile_sample_stride_num_frames = 8 + + def enable_tiling( + self, + tile_sample_min_height: int | None = None, + tile_sample_min_width: int | None = None, + tile_sample_min_num_frames: int | None = None, + tile_sample_stride_height: float | None = None, + tile_sample_stride_width: float | None = None, + tile_sample_stride_num_frames: float | None = None, + ) -> None: + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger images. + + Args: + tile_sample_min_height (`int`, *optional*): + The minimum height required for a sample to be separated into tiles across the height dimension. + tile_sample_min_width (`int`, *optional*): + The minimum width required for a sample to be separated into tiles across the width dimension. + tile_sample_stride_height (`int`, *optional*): + The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are + no tiling artifacts produced across the height dimension. + tile_sample_stride_width (`int`, *optional*): + The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling + artifacts produced across the width dimension. + """ + self.use_tiling = True + self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height + self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width + self.tile_sample_min_num_frames = tile_sample_min_num_frames or self.tile_sample_min_num_frames + self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height + self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width + self.tile_sample_stride_num_frames = tile_sample_stride_num_frames or self.tile_sample_stride_num_frames + + def _encode(self, x: torch.Tensor, causal: bool | None = None) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = x.shape + + if self.use_framewise_decoding and num_frames > self.tile_sample_min_num_frames: + return self._temporal_tiled_encode(x, causal=causal) + + if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height): + return self.tiled_encode(x, causal=causal) + + enc = self.encoder(x, causal=causal) + + return enc + + @apply_forward_hook + def encode( + self, x: torch.Tensor, causal: bool | None = None, return_dict: bool = True + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: + """ + Encode a batch of images into latents. + + Args: + x (`torch.Tensor`): Input batch of images. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. + + Returns: + The latent representations of the encoded videos. If `return_dict` is True, a + [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. + """ + if self.use_slicing and x.shape[0] > 1: + encoded_slices = [self._encode(x_slice, causal=causal) for x_slice in x.split(1)] + h = torch.cat(encoded_slices) + else: + h = self._encode(x, causal=causal) + posterior = DiagonalGaussianDistribution(h) + + if not return_dict: + return (posterior,) + return AutoencoderKLOutput(latent_dist=posterior) + + def _decode( + self, + z: torch.Tensor, + temb: torch.Tensor | None = None, + causal: bool | None = None, + return_dict: bool = True, + ) -> DecoderOutput | torch.Tensor: + batch_size, num_channels, num_frames, height, width = z.shape + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio + + if self.use_framewise_decoding and num_frames > tile_latent_min_num_frames: + return self._temporal_tiled_decode(z, temb, causal=causal, return_dict=return_dict) + + if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height): + return self.tiled_decode(z, temb, causal=causal, return_dict=return_dict) + + dec = self.decoder(z, temb, causal=causal) + + if not return_dict: + return (dec,) + + return DecoderOutput(sample=dec) + + @apply_forward_hook + def decode( + self, + z: torch.Tensor, + temb: torch.Tensor | None = None, + causal: bool | None = None, + return_dict: bool = True, + ) -> DecoderOutput | torch.Tensor: + """ + Decode a batch of images. + + Args: + z (`torch.Tensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + """ + if self.use_slicing and z.shape[0] > 1: + if temb is not None: + decoded_slices = [ + self._decode(z_slice, t_slice, causal=causal).sample + for z_slice, t_slice in (z.split(1), temb.split(1)) + ] + else: + decoded_slices = [self._decode(z_slice, causal=causal).sample for z_slice in z.split(1)] + decoded = torch.cat(decoded_slices) + else: + decoded = self._decode(z, temb, causal=causal).sample + + if not return_dict: + return (decoded,) + + return DecoderOutput(sample=decoded) + + def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[3], b.shape[3], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * ( + y / blend_extent + ) + return b + + def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[4], b.shape[4], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * ( + x / blend_extent + ) + return b + + def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) + for x in range(blend_extent): + b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :, x, :, :] * ( + x / blend_extent + ) + return b + + def tiled_encode(self, x: torch.Tensor, causal: bool | None = None) -> torch.Tensor: + r"""Encode a batch of images using a tiled encoder. + + Args: + x (`torch.Tensor`): Input batch of videos. + + Returns: + `torch.Tensor`: + The latent representation of the encoded videos. + """ + batch_size, num_channels, num_frames, height, width = x.shape + latent_height = height // self.spatial_compression_ratio + latent_width = width // self.spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + + blend_height = tile_latent_min_height - tile_latent_stride_height + blend_width = tile_latent_min_width - tile_latent_stride_width + + # Split x into overlapping tiles and encode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, height, self.tile_sample_stride_height): + row = [] + for j in range(0, width, self.tile_sample_stride_width): + time = self.encoder( + x[:, :, :, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width], + causal=causal, + ) + + row.append(time) + rows.append(row) + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_height) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_width) + result_row.append(tile[:, :, :, :tile_latent_stride_height, :tile_latent_stride_width]) + result_rows.append(torch.cat(result_row, dim=4)) + + enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width] + return enc + + def tiled_decode( + self, z: torch.Tensor, temb: torch.Tensor | None, causal: bool | None = None, return_dict: bool = True + ) -> DecoderOutput | torch.Tensor: + r""" + Decode a batch of images using a tiled decoder. + + Args: + z (`torch.Tensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + """ + + batch_size, num_channels, num_frames, height, width = z.shape + sample_height = height * self.spatial_compression_ratio + sample_width = width * self.spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + + blend_height = self.tile_sample_min_height - self.tile_sample_stride_height + blend_width = self.tile_sample_min_width - self.tile_sample_stride_width + + # Split z into overlapping tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, height, tile_latent_stride_height): + row = [] + for j in range(0, width, tile_latent_stride_width): + time = self.decoder( + z[:, :, :, i : i + tile_latent_min_height, j : j + tile_latent_min_width], temb, causal=causal + ) + + row.append(time) + rows.append(row) + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_height) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_width) + result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width]) + result_rows.append(torch.cat(result_row, dim=4)) + + dec = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width] + + if not return_dict: + return (dec,) + + return DecoderOutput(sample=dec) + + def _temporal_tiled_encode(self, x: torch.Tensor, causal: bool | None = None) -> AutoencoderKLOutput: + batch_size, num_channels, num_frames, height, width = x.shape + latent_num_frames = (num_frames - 1) // self.temporal_compression_ratio + 1 + + tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio + tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio + blend_num_frames = tile_latent_min_num_frames - tile_latent_stride_num_frames + + row = [] + for i in range(0, num_frames, self.tile_sample_stride_num_frames): + tile = x[:, :, i : i + self.tile_sample_min_num_frames + 1, :, :] + if self.use_tiling and (height > self.tile_sample_min_height or width > self.tile_sample_min_width): + tile = self.tiled_encode(tile, causal=causal) + else: + tile = self.encoder(tile, causal=causal) + if i > 0: + tile = tile[:, :, 1:, :, :] + row.append(tile) + + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_num_frames) + result_row.append(tile[:, :, :tile_latent_stride_num_frames, :, :]) + else: + result_row.append(tile[:, :, : tile_latent_stride_num_frames + 1, :, :]) + + enc = torch.cat(result_row, dim=2)[:, :, :latent_num_frames] + return enc + + def _temporal_tiled_decode( + self, z: torch.Tensor, temb: torch.Tensor | None, causal: bool | None = None, return_dict: bool = True + ) -> DecoderOutput | torch.Tensor: + batch_size, num_channels, num_frames, height, width = z.shape + num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1 + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio + tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio + blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames + + row = [] + for i in range(0, num_frames, tile_latent_stride_num_frames): + tile = z[:, :, i : i + tile_latent_min_num_frames + 1, :, :] + if self.use_tiling and (tile.shape[-1] > tile_latent_min_width or tile.shape[-2] > tile_latent_min_height): + decoded = self.tiled_decode(tile, temb, causal=causal, return_dict=True).sample + else: + decoded = self.decoder(tile, temb, causal=causal) + if i > 0: + decoded = decoded[:, :, :-1, :, :] + row.append(decoded) + + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_num_frames) + tile = tile[:, :, : self.tile_sample_stride_num_frames, :, :] + result_row.append(tile) + else: + result_row.append(tile[:, :, : self.tile_sample_stride_num_frames + 1, :, :]) + + dec = torch.cat(result_row, dim=2)[:, :, :num_sample_frames] + + if not return_dict: + return (dec,) + return DecoderOutput(sample=dec) + + def forward( + self, + sample: torch.Tensor, + temb: torch.Tensor | None = None, + sample_posterior: bool = False, + encoder_causal: bool | None = None, + decoder_causal: bool | None = None, + return_dict: bool = True, + generator: torch.Generator | None = None, + ) -> torch.Tensor | torch.Tensor: + x = sample + posterior = self.encode(x, causal=encoder_causal).latent_dist + if sample_posterior: + z = posterior.sample(generator=generator) + else: + z = posterior.mode() + dec = self.decode(z, temb, causal=decoder_causal) + if not return_dict: + return (dec.sample,) + return dec diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_ltx2_audio.py b/src/diffusers/models/autoencoders/autoencoder_kl_ltx2_audio.py new file mode 100644 index 000000000000..f9390dab5b74 --- /dev/null +++ b/src/diffusers/models/autoencoders/autoencoder_kl_ltx2_audio.py @@ -0,0 +1,803 @@ +# Copyright 2025 The Lightricks team and The HuggingFace Team. +# All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...utils.accelerate_utils import apply_forward_hook +from ..modeling_outputs import AutoencoderKLOutput +from ..modeling_utils import ModelMixin +from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution + + +LATENT_DOWNSAMPLE_FACTOR = 4 + + +class LTX2AudioCausalConv2d(nn.Module): + """ + A causal 2D convolution that pads asymmetrically along the causal axis. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int], + stride: int = 1, + dilation: int | tuple[int, int] = 1, + groups: int = 1, + bias: bool = True, + causality_axis: str = "height", + ) -> None: + super().__init__() + + self.causality_axis = causality_axis + kernel_size = (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size + dilation = (dilation, dilation) if isinstance(dilation, int) else dilation + + pad_h = (kernel_size[0] - 1) * dilation[0] + pad_w = (kernel_size[1] - 1) * dilation[1] + + if self.causality_axis == "none": + padding = (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2) + elif self.causality_axis in {"width", "width-compatibility"}: + padding = (pad_w, 0, pad_h // 2, pad_h - pad_h // 2) + elif self.causality_axis == "height": + padding = (pad_w // 2, pad_w - pad_w // 2, pad_h, 0) + else: + raise ValueError(f"Invalid causality_axis: {causality_axis}") + + self.padding = padding + self.conv = nn.Conv2d( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=0, + dilation=dilation, + groups=groups, + bias=bias, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = F.pad(x, self.padding) + return self.conv(x) + + +class LTX2AudioPixelNorm(nn.Module): + """ + Per-pixel (per-location) RMS normalization layer. + """ + + def __init__(self, dim: int = 1, eps: float = 1e-8) -> None: + super().__init__() + self.dim = dim + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True) + rms = torch.sqrt(mean_sq + self.eps) + return x / rms + + +class LTX2AudioAttnBlock(nn.Module): + def __init__( + self, + in_channels: int, + norm_type: str = "group", + ) -> None: + super().__init__() + self.in_channels = in_channels + + if norm_type == "group": + self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + elif norm_type == "pixel": + self.norm = LTX2AudioPixelNorm(dim=1, eps=1e-6) + else: + raise ValueError(f"Invalid normalization type: {norm_type}") + self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h_ = self.norm(x) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + batch, channels, height, width = q.shape + q = q.reshape(batch, channels, height * width).permute(0, 2, 1).contiguous() + k = k.reshape(batch, channels, height * width).contiguous() + attn = torch.bmm(q, k) * (int(channels) ** (-0.5)) + attn = torch.nn.functional.softmax(attn, dim=2) + + v = v.reshape(batch, channels, height * width) + attn = attn.permute(0, 2, 1).contiguous() + h_ = torch.bmm(v, attn).reshape(batch, channels, height, width) + + h_ = self.proj_out(h_) + return x + h_ + + +class LTX2AudioResnetBlock(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int | None = None, + conv_shortcut: bool = False, + dropout: float = 0.0, + temb_channels: int = 512, + norm_type: str = "group", + causality_axis: str = "height", + ) -> None: + super().__init__() + self.causality_axis = causality_axis + + if self.causality_axis is not None and self.causality_axis != "none" and norm_type == "group": + raise ValueError("Causal ResnetBlock with GroupNorm is not supported.") + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + self.use_conv_shortcut = conv_shortcut + + if norm_type == "group": + self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + elif norm_type == "pixel": + self.norm1 = LTX2AudioPixelNorm(dim=1, eps=1e-6) + else: + raise ValueError(f"Invalid normalization type: {norm_type}") + self.non_linearity = nn.SiLU() + if causality_axis is not None: + self.conv1 = LTX2AudioCausalConv2d( + in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis + ) + else: + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + if temb_channels > 0: + self.temb_proj = nn.Linear(temb_channels, out_channels) + if norm_type == "group": + self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True) + elif norm_type == "pixel": + self.norm2 = LTX2AudioPixelNorm(dim=1, eps=1e-6) + else: + raise ValueError(f"Invalid normalization type: {norm_type}") + self.dropout = nn.Dropout(dropout) + if causality_axis is not None: + self.conv2 = LTX2AudioCausalConv2d( + out_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis + ) + else: + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + if causality_axis is not None: + self.conv_shortcut = LTX2AudioCausalConv2d( + in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis + ) + else: + self.conv_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + else: + if causality_axis is not None: + self.nin_shortcut = LTX2AudioCausalConv2d( + in_channels, out_channels, kernel_size=1, stride=1, causality_axis=causality_axis + ) + else: + self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor: + h = self.norm1(x) + h = self.non_linearity(h) + h = self.conv1(h) + + if temb is not None: + h = h + self.temb_proj(self.non_linearity(temb))[:, :, None, None] + + h = self.norm2(h) + h = self.non_linearity(h) + h = self.dropout(h) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + x = self.conv_shortcut(x) if self.use_conv_shortcut else self.nin_shortcut(x) + + return x + h + + +class LTX2AudioDownsample(nn.Module): + def __init__(self, in_channels: int, with_conv: bool, causality_axis: str | None = "height") -> None: + super().__init__() + self.with_conv = with_conv + self.causality_axis = causality_axis + + if self.with_conv: + self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.with_conv: + # Padding tuple is in the order: (left, right, top, bottom). + if self.causality_axis == "none": + pad = (0, 1, 0, 1) + elif self.causality_axis == "width": + pad = (2, 0, 0, 1) + elif self.causality_axis == "height": + pad = (0, 1, 2, 0) + elif self.causality_axis == "width-compatibility": + pad = (1, 0, 0, 1) + else: + raise ValueError( + f"Invalid `causality_axis` {self.causality_axis}; supported values are `none`, `width`, `height`," + f" and `width-compatibility`." + ) + + x = F.pad(x, pad, mode="constant", value=0) + x = self.conv(x) + else: + # with_conv=False implies that causality_axis is "none" + x = F.avg_pool2d(x, kernel_size=2, stride=2) + return x + + +class LTX2AudioUpsample(nn.Module): + def __init__(self, in_channels: int, with_conv: bool, causality_axis: str | None = "height") -> None: + super().__init__() + self.with_conv = with_conv + self.causality_axis = causality_axis + if self.with_conv: + if causality_axis is not None: + self.conv = LTX2AudioCausalConv2d( + in_channels, in_channels, kernel_size=3, stride=1, causality_axis=causality_axis + ) + else: + self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") + if self.with_conv: + x = self.conv(x) + if self.causality_axis is None or self.causality_axis == "none": + pass + elif self.causality_axis == "height": + x = x[:, :, 1:, :] + elif self.causality_axis == "width": + x = x[:, :, :, 1:] + elif self.causality_axis == "width-compatibility": + pass + else: + raise ValueError(f"Invalid causality_axis: {self.causality_axis}") + + return x + + +class LTX2AudioAudioPatchifier: + """ + Patchifier for spectrogram/audio latents. + """ + + def __init__( + self, + patch_size: int, + sample_rate: int = 16000, + hop_length: int = 160, + audio_latent_downsample_factor: int = 4, + is_causal: bool = True, + ): + self.hop_length = hop_length + self.sample_rate = sample_rate + self.audio_latent_downsample_factor = audio_latent_downsample_factor + self.is_causal = is_causal + self._patch_size = (1, patch_size, patch_size) + + def patchify(self, audio_latents: torch.Tensor) -> torch.Tensor: + batch, channels, time, freq = audio_latents.shape + return audio_latents.permute(0, 2, 1, 3).reshape(batch, time, channels * freq) + + def unpatchify(self, audio_latents: torch.Tensor, channels: int, mel_bins: int) -> torch.Tensor: + batch, time, _ = audio_latents.shape + return audio_latents.view(batch, time, channels, mel_bins).permute(0, 2, 1, 3) + + @property + def patch_size(self) -> tuple[int, int, int]: + return self._patch_size + + +class LTX2AudioEncoder(nn.Module): + def __init__( + self, + base_channels: int = 128, + output_channels: int = 1, + num_res_blocks: int = 2, + attn_resolutions: tuple[int, ...] | None = None, + in_channels: int = 2, + resolution: int = 256, + latent_channels: int = 8, + ch_mult: tuple[int, ...] = (1, 2, 4), + norm_type: str = "group", + causality_axis: str | None = "width", + dropout: float = 0.0, + mid_block_add_attention: bool = False, + sample_rate: int = 16000, + mel_hop_length: int = 160, + is_causal: bool = True, + mel_bins: int | None = 64, + double_z: bool = True, + ): + super().__init__() + + self.sample_rate = sample_rate + self.mel_hop_length = mel_hop_length + self.is_causal = is_causal + self.mel_bins = mel_bins + + self.base_channels = base_channels + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + self.out_ch = output_channels + self.give_pre_end = False + self.tanh_out = False + self.norm_type = norm_type + self.latent_channels = latent_channels + self.channel_multipliers = ch_mult + self.attn_resolutions = attn_resolutions + self.causality_axis = causality_axis + + base_block_channels = base_channels + base_resolution = resolution + self.z_shape = (1, latent_channels, base_resolution, base_resolution) + + if self.causality_axis is not None: + self.conv_in = LTX2AudioCausalConv2d( + in_channels, base_block_channels, kernel_size=3, stride=1, causality_axis=self.causality_axis + ) + else: + self.conv_in = nn.Conv2d(in_channels, base_block_channels, kernel_size=3, stride=1, padding=1) + + self.down = nn.ModuleList() + block_in = base_block_channels + curr_res = self.resolution + + for level in range(self.num_resolutions): + stage = nn.Module() + stage.block = nn.ModuleList() + stage.attn = nn.ModuleList() + block_out = self.base_channels * self.channel_multipliers[level] + + for _ in range(self.num_res_blocks): + stage.block.append( + LTX2AudioResnetBlock( + in_channels=block_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout, + norm_type=self.norm_type, + causality_axis=self.causality_axis, + ) + ) + block_in = block_out + if self.attn_resolutions: + if curr_res in self.attn_resolutions: + stage.attn.append(LTX2AudioAttnBlock(block_in, norm_type=self.norm_type)) + + if level != self.num_resolutions - 1: + stage.downsample = LTX2AudioDownsample(block_in, True, causality_axis=self.causality_axis) + curr_res = curr_res // 2 + + self.down.append(stage) + + self.mid = nn.Module() + self.mid.block_1 = LTX2AudioResnetBlock( + in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout, + norm_type=self.norm_type, + causality_axis=self.causality_axis, + ) + if mid_block_add_attention: + self.mid.attn_1 = LTX2AudioAttnBlock(block_in, norm_type=self.norm_type) + else: + self.mid.attn_1 = nn.Identity() + self.mid.block_2 = LTX2AudioResnetBlock( + in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout, + norm_type=self.norm_type, + causality_axis=self.causality_axis, + ) + + final_block_channels = block_in + z_channels = 2 * latent_channels if double_z else latent_channels + if self.norm_type == "group": + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=final_block_channels, eps=1e-6, affine=True) + elif self.norm_type == "pixel": + self.norm_out = LTX2AudioPixelNorm(dim=1, eps=1e-6) + else: + raise ValueError(f"Invalid normalization type: {self.norm_type}") + self.non_linearity = nn.SiLU() + + if self.causality_axis is not None: + self.conv_out = LTX2AudioCausalConv2d( + final_block_channels, z_channels, kernel_size=3, stride=1, causality_axis=self.causality_axis + ) + else: + self.conv_out = nn.Conv2d(final_block_channels, z_channels, kernel_size=3, stride=1, padding=1) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # hidden_states expected shape: (batch_size, channels, time, num_mel_bins) + hidden_states = self.conv_in(hidden_states) + + for level in range(self.num_resolutions): + stage = self.down[level] + for block_idx, block in enumerate(stage.block): + hidden_states = block(hidden_states, temb=None) + if stage.attn: + hidden_states = stage.attn[block_idx](hidden_states) + + if level != self.num_resolutions - 1 and hasattr(stage, "downsample"): + hidden_states = stage.downsample(hidden_states) + + hidden_states = self.mid.block_1(hidden_states, temb=None) + hidden_states = self.mid.attn_1(hidden_states) + hidden_states = self.mid.block_2(hidden_states, temb=None) + + hidden_states = self.norm_out(hidden_states) + hidden_states = self.non_linearity(hidden_states) + hidden_states = self.conv_out(hidden_states) + + return hidden_states + + +class LTX2AudioDecoder(nn.Module): + """ + Symmetric decoder that reconstructs audio spectrograms from latent features. + + The decoder mirrors the encoder structure with configurable channel multipliers, attention resolutions, and causal + convolutions. + """ + + def __init__( + self, + base_channels: int = 128, + output_channels: int = 1, + num_res_blocks: int = 2, + attn_resolutions: tuple[int, ...] | None = None, + in_channels: int = 2, + resolution: int = 256, + latent_channels: int = 8, + ch_mult: tuple[int, ...] = (1, 2, 4), + norm_type: str = "group", + causality_axis: str | None = "width", + dropout: float = 0.0, + mid_block_add_attention: bool = False, + sample_rate: int = 16000, + mel_hop_length: int = 160, + is_causal: bool = True, + mel_bins: int | None = 64, + ) -> None: + super().__init__() + + self.sample_rate = sample_rate + self.mel_hop_length = mel_hop_length + self.is_causal = is_causal + self.mel_bins = mel_bins + self.patchifier = LTX2AudioAudioPatchifier( + patch_size=1, + audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR, + sample_rate=sample_rate, + hop_length=mel_hop_length, + is_causal=is_causal, + ) + + self.base_channels = base_channels + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + self.out_ch = output_channels + self.give_pre_end = False + self.tanh_out = False + self.norm_type = norm_type + self.latent_channels = latent_channels + self.channel_multipliers = ch_mult + self.attn_resolutions = attn_resolutions + self.causality_axis = causality_axis + + base_block_channels = base_channels * self.channel_multipliers[-1] + base_resolution = resolution // (2 ** (self.num_resolutions - 1)) + self.z_shape = (1, latent_channels, base_resolution, base_resolution) + + if self.causality_axis is not None: + self.conv_in = LTX2AudioCausalConv2d( + latent_channels, base_block_channels, kernel_size=3, stride=1, causality_axis=self.causality_axis + ) + else: + self.conv_in = nn.Conv2d(latent_channels, base_block_channels, kernel_size=3, stride=1, padding=1) + self.non_linearity = nn.SiLU() + self.mid = nn.Module() + self.mid.block_1 = LTX2AudioResnetBlock( + in_channels=base_block_channels, + out_channels=base_block_channels, + temb_channels=self.temb_ch, + dropout=dropout, + norm_type=self.norm_type, + causality_axis=self.causality_axis, + ) + if mid_block_add_attention: + self.mid.attn_1 = LTX2AudioAttnBlock(base_block_channels, norm_type=self.norm_type) + else: + self.mid.attn_1 = nn.Identity() + self.mid.block_2 = LTX2AudioResnetBlock( + in_channels=base_block_channels, + out_channels=base_block_channels, + temb_channels=self.temb_ch, + dropout=dropout, + norm_type=self.norm_type, + causality_axis=self.causality_axis, + ) + + self.up = nn.ModuleList() + block_in = base_block_channels + curr_res = self.resolution // (2 ** (self.num_resolutions - 1)) + + for level in reversed(range(self.num_resolutions)): + stage = nn.Module() + stage.block = nn.ModuleList() + stage.attn = nn.ModuleList() + block_out = self.base_channels * self.channel_multipliers[level] + + for _ in range(self.num_res_blocks + 1): + stage.block.append( + LTX2AudioResnetBlock( + in_channels=block_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout, + norm_type=self.norm_type, + causality_axis=self.causality_axis, + ) + ) + block_in = block_out + if self.attn_resolutions: + if curr_res in self.attn_resolutions: + stage.attn.append(LTX2AudioAttnBlock(block_in, norm_type=self.norm_type)) + + if level != 0: + stage.upsample = LTX2AudioUpsample(block_in, True, causality_axis=self.causality_axis) + curr_res *= 2 + + self.up.insert(0, stage) + + final_block_channels = block_in + + if self.norm_type == "group": + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=final_block_channels, eps=1e-6, affine=True) + elif self.norm_type == "pixel": + self.norm_out = LTX2AudioPixelNorm(dim=1, eps=1e-6) + else: + raise ValueError(f"Invalid normalization type: {self.norm_type}") + + if self.causality_axis is not None: + self.conv_out = LTX2AudioCausalConv2d( + final_block_channels, output_channels, kernel_size=3, stride=1, causality_axis=self.causality_axis + ) + else: + self.conv_out = nn.Conv2d(final_block_channels, output_channels, kernel_size=3, stride=1, padding=1) + + def forward( + self, + sample: torch.Tensor, + ) -> torch.Tensor: + _, _, frames, mel_bins = sample.shape + + target_frames = frames * LATENT_DOWNSAMPLE_FACTOR + + if self.causality_axis is not None: + target_frames = max(target_frames - (LATENT_DOWNSAMPLE_FACTOR - 1), 1) + + target_channels = self.out_ch + target_mel_bins = self.mel_bins if self.mel_bins is not None else mel_bins + + hidden_features = self.conv_in(sample) + hidden_features = self.mid.block_1(hidden_features, temb=None) + hidden_features = self.mid.attn_1(hidden_features) + hidden_features = self.mid.block_2(hidden_features, temb=None) + + for level in reversed(range(self.num_resolutions)): + stage = self.up[level] + for block_idx, block in enumerate(stage.block): + hidden_features = block(hidden_features, temb=None) + if stage.attn: + hidden_features = stage.attn[block_idx](hidden_features) + + if level != 0 and hasattr(stage, "upsample"): + hidden_features = stage.upsample(hidden_features) + + if self.give_pre_end: + return hidden_features + + hidden = self.norm_out(hidden_features) + hidden = self.non_linearity(hidden) + decoded_output = self.conv_out(hidden) + decoded_output = torch.tanh(decoded_output) if self.tanh_out else decoded_output + + _, _, current_time, current_freq = decoded_output.shape + target_time = target_frames + target_freq = target_mel_bins + + decoded_output = decoded_output[ + :, :target_channels, : min(current_time, target_time), : min(current_freq, target_freq) + ] + + time_padding_needed = target_time - decoded_output.shape[2] + freq_padding_needed = target_freq - decoded_output.shape[3] + + if time_padding_needed > 0 or freq_padding_needed > 0: + padding = ( + 0, + max(freq_padding_needed, 0), + 0, + max(time_padding_needed, 0), + ) + decoded_output = F.pad(decoded_output, padding) + + decoded_output = decoded_output[:, :target_channels, :target_time, :target_freq] + + return decoded_output + + +class AutoencoderKLLTX2Audio(ModelMixin, AutoencoderMixin, ConfigMixin): + r""" + LTX2 audio VAE for encoding and decoding audio latent representations. + """ + + _supports_gradient_checkpointing = False + + @register_to_config + def __init__( + self, + base_channels: int = 128, + output_channels: int = 2, + ch_mult: tuple[int, ...] = (1, 2, 4), + num_res_blocks: int = 2, + attn_resolutions: tuple[int, ...] | None = None, + in_channels: int = 2, + resolution: int = 256, + latent_channels: int = 8, + norm_type: str = "pixel", + causality_axis: str | None = "height", + dropout: float = 0.0, + mid_block_add_attention: bool = False, + sample_rate: int = 16000, + mel_hop_length: int = 160, + is_causal: bool = True, + mel_bins: int | None = 64, + double_z: bool = True, + ) -> None: + super().__init__() + + supported_causality_axes = {"none", "width", "height", "width-compatibility"} + if causality_axis not in supported_causality_axes: + raise ValueError(f"{causality_axis=} is not valid. Supported values: {supported_causality_axes}") + + attn_resolution_set = set(attn_resolutions) if attn_resolutions else attn_resolutions + + self.encoder = LTX2AudioEncoder( + base_channels=base_channels, + output_channels=output_channels, + ch_mult=ch_mult, + num_res_blocks=num_res_blocks, + attn_resolutions=attn_resolution_set, + in_channels=in_channels, + resolution=resolution, + latent_channels=latent_channels, + norm_type=norm_type, + causality_axis=causality_axis, + dropout=dropout, + mid_block_add_attention=mid_block_add_attention, + sample_rate=sample_rate, + mel_hop_length=mel_hop_length, + is_causal=is_causal, + mel_bins=mel_bins, + double_z=double_z, + ) + + self.decoder = LTX2AudioDecoder( + base_channels=base_channels, + output_channels=output_channels, + ch_mult=ch_mult, + num_res_blocks=num_res_blocks, + attn_resolutions=attn_resolution_set, + in_channels=in_channels, + resolution=resolution, + latent_channels=latent_channels, + norm_type=norm_type, + causality_axis=causality_axis, + dropout=dropout, + mid_block_add_attention=mid_block_add_attention, + sample_rate=sample_rate, + mel_hop_length=mel_hop_length, + is_causal=is_causal, + mel_bins=mel_bins, + ) + + # Per-channel statistics for normalizing and denormalizing the latent representation. This statics is computed over + # the entire dataset and stored in model's checkpoint under AudioVAE state_dict + latents_std = torch.ones((base_channels,)) + latents_mean = torch.zeros((base_channels,)) + self.register_buffer("latents_mean", latents_mean, persistent=True) + self.register_buffer("latents_std", latents_std, persistent=True) + + # TODO: calculate programmatically instead of hardcoding + self.temporal_compression_ratio = LATENT_DOWNSAMPLE_FACTOR # 4 + # TODO: confirm whether the mel compression ratio below is correct + self.mel_compression_ratio = LATENT_DOWNSAMPLE_FACTOR + self.use_slicing = False + + def _encode(self, x: torch.Tensor) -> torch.Tensor: + return self.encoder(x) + + @apply_forward_hook + def encode(self, x: torch.Tensor, return_dict: bool = True): + if self.use_slicing and x.shape[0] > 1: + encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)] + h = torch.cat(encoded_slices) + else: + h = self._encode(x) + posterior = DiagonalGaussianDistribution(h) + + if not return_dict: + return (posterior,) + return AutoencoderKLOutput(latent_dist=posterior) + + def _decode(self, z: torch.Tensor) -> torch.Tensor: + return self.decoder(z) + + @apply_forward_hook + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: + if self.use_slicing and z.shape[0] > 1: + decoded_slices = [self._decode(z_slice) for z_slice in z.split(1)] + decoded = torch.cat(decoded_slices) + else: + decoded = self._decode(z) + + if not return_dict: + return (decoded,) + + return DecoderOutput(sample=decoded) + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + return_dict: bool = True, + generator: torch.Generator | None = None, + ) -> DecoderOutput | torch.Tensor: + posterior = self.encode(sample).latent_dist + if sample_posterior: + z = posterior.sample(generator=generator) + else: + z = posterior.mode() + dec = self.decode(z) + if not return_dict: + return (dec.sample,) + return dec diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_magvit.py b/src/diffusers/models/autoencoders/autoencoder_kl_magvit.py index 97ca9d669264..ea0e2cd00d52 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_magvit.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_magvit.py @@ -14,7 +14,6 @@ # limitations under the License. import math -from typing import Optional, Tuple, Union import torch import torch.nn as nn @@ -37,10 +36,10 @@ def __init__( self, in_channels: int, out_channels: int, - kernel_size: Union[int, Tuple[int, ...]] = 3, - stride: Union[int, Tuple[int, ...]] = 1, - padding: Union[int, Tuple[int, ...]] = 1, - dilation: Union[int, Tuple[int, ...]] = 1, + kernel_size: int | tuple[int, ...] = 3, + stride: int | tuple[int, ...] = 1, + padding: int | tuple[int, ...] = 1, + dilation: int | tuple[int, ...] = 1, groups: int = 1, bias: bool = True, padding_mode: str = "zeros", @@ -437,13 +436,13 @@ def __init__( self, in_channels: int = 3, out_channels: int = 8, - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str, ...] = ( "SpatialDownBlock3D", "SpatialTemporalDownBlock3D", "SpatialTemporalDownBlock3D", "SpatialTemporalDownBlock3D", ), - block_out_channels: Tuple[int, ...] = [128, 256, 512, 512], + block_out_channels: tuple[int, ...] = [128, 256, 512, 512], layers_per_block: int = 2, norm_num_groups: int = 32, act_fn: str = "silu", @@ -553,13 +552,13 @@ def __init__( self, in_channels: int = 8, out_channels: int = 3, - up_block_types: Tuple[str, ...] = ( + up_block_types: tuple[str, ...] = ( "SpatialUpBlock3D", "SpatialTemporalUpBlock3D", "SpatialTemporalUpBlock3D", "SpatialTemporalUpBlock3D", ), - block_out_channels: Tuple[int, ...] = [128, 256, 512, 512], + block_out_channels: tuple[int, ...] = [128, 256, 512, 512], layers_per_block: int = 2, norm_num_groups: int = 32, act_fn: str = "silu", @@ -680,14 +679,14 @@ def __init__( in_channels: int = 3, latent_channels: int = 16, out_channels: int = 3, - block_out_channels: Tuple[int, ...] = [128, 256, 512, 512], - down_block_types: Tuple[str, ...] = [ + block_out_channels: tuple[int, ...] = [128, 256, 512, 512], + down_block_types: tuple[str, ...] = [ "SpatialDownBlock3D", "SpatialTemporalDownBlock3D", "SpatialTemporalDownBlock3D", "SpatialTemporalDownBlock3D", ], - up_block_types: Tuple[str, ...] = [ + up_block_types: tuple[str, ...] = [ "SpatialUpBlock3D", "SpatialTemporalUpBlock3D", "SpatialTemporalUpBlock3D", @@ -771,12 +770,12 @@ def _clear_conv_cache(self): def enable_tiling( self, - tile_sample_min_height: Optional[int] = None, - tile_sample_min_width: Optional[int] = None, - tile_sample_min_num_frames: Optional[int] = None, - tile_sample_stride_height: Optional[float] = None, - tile_sample_stride_width: Optional[float] = None, - tile_sample_stride_num_frames: Optional[float] = None, + tile_sample_min_height: int | None = None, + tile_sample_min_width: int | None = None, + tile_sample_min_num_frames: int | None = None, + tile_sample_stride_height: float | None = None, + tile_sample_stride_width: float | None = None, + tile_sample_stride_num_frames: float | None = None, ) -> None: r""" Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to @@ -808,7 +807,7 @@ def enable_tiling( @apply_forward_hook def _encode( self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: """ Encode a batch of images into latents. @@ -838,7 +837,7 @@ def _encode( @apply_forward_hook def encode( self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: """ Encode a batch of images into latents. @@ -863,7 +862,7 @@ def encode( return (posterior,) return AutoencoderKLOutput(latent_dist=posterior) - def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: batch_size, num_channels, num_frames, height, width = z.shape tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio @@ -890,7 +889,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOut return DecoderOutput(sample=dec) @apply_forward_hook - def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: """ Decode a batch of images. @@ -983,7 +982,7 @@ def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> Autoencoder moments = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width] return moments - def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: batch_size, num_channels, num_frames, height, width = z.shape sample_height = height * self.spatial_compression_ratio sample_width = width * self.spatial_compression_ratio @@ -1049,8 +1048,8 @@ def forward( sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[DecoderOutput, torch.Tensor]: + generator: torch.Generator | None = None, + ) -> DecoderOutput | torch.Tensor: r""" Args: sample (`torch.Tensor`): Input sample. diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py b/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py index 7a64ac7de172..a0f831c867b0 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py @@ -14,7 +14,6 @@ # limitations under the License. import functools -from typing import Dict, Optional, Tuple, Union import torch import torch.nn as nn @@ -83,7 +82,7 @@ class MochiResnetBlock3D(nn.Module): def __init__( self, in_channels: int, - out_channels: Optional[int] = None, + out_channels: int | None = None, act_fn: str = "swish", ): super().__init__() @@ -106,7 +105,7 @@ def __init__( def forward( self, inputs: torch.Tensor, - conv_cache: Optional[Dict[str, torch.Tensor]] = None, + conv_cache: dict[str, torch.Tensor] | None = None, ) -> torch.Tensor: new_conv_cache = {} conv_cache = conv_cache or {} @@ -193,7 +192,7 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - conv_cache: Optional[Dict[str, torch.Tensor]] = None, + conv_cache: dict[str, torch.Tensor] | None = None, chunk_size: int = 2**15, ) -> torch.Tensor: r"""Forward method of the `MochiUpBlock3D` class.""" @@ -294,7 +293,7 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - conv_cache: Optional[Dict[str, torch.Tensor]] = None, + conv_cache: dict[str, torch.Tensor] | None = None, ) -> torch.Tensor: r"""Forward method of the `MochiMidBlock3D` class.""" @@ -368,7 +367,7 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - conv_cache: Optional[Dict[str, torch.Tensor]] = None, + conv_cache: dict[str, torch.Tensor] | None = None, ) -> torch.Tensor: r"""Forward method of the `MochiUpBlock3D` class.""" @@ -445,13 +444,13 @@ class MochiEncoder3D(nn.Module): The number of input channels. out_channels (`int`, *optional*): The number of output channels. - block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(128, 256, 512, 768)`): + block_out_channels (`tuple[int, ...]`, *optional*, defaults to `(128, 256, 512, 768)`): The number of output channels for each block. - layers_per_block (`Tuple[int, ...]`, *optional*, defaults to `(3, 3, 4, 6, 3)`): + layers_per_block (`tuple[int, ...]`, *optional*, defaults to `(3, 3, 4, 6, 3)`): The number of resnet blocks for each block. - temporal_expansions (`Tuple[int, ...]`, *optional*, defaults to `(1, 2, 3)`): + temporal_expansions (`tuple[int, ...]`, *optional*, defaults to `(1, 2, 3)`): The temporal expansion factor for each of the up blocks. - spatial_expansions (`Tuple[int, ...]`, *optional*, defaults to `(2, 2, 2)`): + spatial_expansions (`tuple[int, ...]`, *optional*, defaults to `(2, 2, 2)`): The spatial expansion factor for each of the up blocks. non_linearity (`str`, *optional*, defaults to `"swish"`): The non-linearity to use in the decoder. @@ -461,11 +460,11 @@ def __init__( self, in_channels: int, out_channels: int, - block_out_channels: Tuple[int, ...] = (128, 256, 512, 768), - layers_per_block: Tuple[int, ...] = (3, 3, 4, 6, 3), - temporal_expansions: Tuple[int, ...] = (1, 2, 3), - spatial_expansions: Tuple[int, ...] = (2, 2, 2), - add_attention_block: Tuple[bool, ...] = (False, True, True, True, True), + block_out_channels: tuple[int, ...] = (128, 256, 512, 768), + layers_per_block: tuple[int, ...] = (3, 3, 4, 6, 3), + temporal_expansions: tuple[int, ...] = (1, 2, 3), + spatial_expansions: tuple[int, ...] = (2, 2, 2), + add_attention_block: tuple[bool, ...] = (False, True, True, True, True), act_fn: str = "swish", ): super().__init__() @@ -499,9 +498,7 @@ def __init__( self.gradient_checkpointing = False - def forward( - self, hidden_states: torch.Tensor, conv_cache: Optional[Dict[str, torch.Tensor]] = None - ) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, conv_cache: dict[str, torch.Tensor] | None = None) -> torch.Tensor: r"""Forward method of the `MochiEncoder3D` class.""" new_conv_cache = {} @@ -558,13 +555,13 @@ class MochiDecoder3D(nn.Module): The number of input channels. out_channels (`int`, *optional*): The number of output channels. - block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(128, 256, 512, 768)`): + block_out_channels (`tuple[int, ...]`, *optional*, defaults to `(128, 256, 512, 768)`): The number of output channels for each block. - layers_per_block (`Tuple[int, ...]`, *optional*, defaults to `(3, 3, 4, 6, 3)`): + layers_per_block (`tuple[int, ...]`, *optional*, defaults to `(3, 3, 4, 6, 3)`): The number of resnet blocks for each block. - temporal_expansions (`Tuple[int, ...]`, *optional*, defaults to `(1, 2, 3)`): + temporal_expansions (`tuple[int, ...]`, *optional*, defaults to `(1, 2, 3)`): The temporal expansion factor for each of the up blocks. - spatial_expansions (`Tuple[int, ...]`, *optional*, defaults to `(2, 2, 2)`): + spatial_expansions (`tuple[int, ...]`, *optional*, defaults to `(2, 2, 2)`): The spatial expansion factor for each of the up blocks. non_linearity (`str`, *optional*, defaults to `"swish"`): The non-linearity to use in the decoder. @@ -574,10 +571,10 @@ def __init__( self, in_channels: int, # 12 out_channels: int, # 3 - block_out_channels: Tuple[int, ...] = (128, 256, 512, 768), - layers_per_block: Tuple[int, ...] = (3, 3, 4, 6, 3), - temporal_expansions: Tuple[int, ...] = (1, 2, 3), - spatial_expansions: Tuple[int, ...] = (2, 2, 2), + block_out_channels: tuple[int, ...] = (128, 256, 512, 768), + layers_per_block: tuple[int, ...] = (3, 3, 4, 6, 3), + temporal_expansions: tuple[int, ...] = (1, 2, 3), + spatial_expansions: tuple[int, ...] = (2, 2, 2), act_fn: str = "swish", ): super().__init__() @@ -612,9 +609,7 @@ def __init__( self.gradient_checkpointing = False - def forward( - self, hidden_states: torch.Tensor, conv_cache: Optional[Dict[str, torch.Tensor]] = None - ) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, conv_cache: dict[str, torch.Tensor] | None = None) -> torch.Tensor: r"""Forward method of the `MochiDecoder3D` class.""" new_conv_cache = {} @@ -668,8 +663,8 @@ class AutoencoderKLMochi(ModelMixin, AutoencoderMixin, ConfigMixin): Parameters: in_channels (int, *optional*, defaults to 3): Number of channels in the input image. out_channels (int, *optional*, defaults to 3): Number of channels in the output. - block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`): - Tuple of block output channels. + block_out_channels (`tuple[int]`, *optional*, defaults to `(64,)`): + tuple of block output channels. act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. scaling_factor (`float`, *optional*, defaults to `1.15258426`): The component-wise standard deviation of the trained latent space computed using the first batch of the @@ -688,15 +683,15 @@ def __init__( self, in_channels: int = 15, out_channels: int = 3, - encoder_block_out_channels: Tuple[int, ...] = (64, 128, 256, 384), - decoder_block_out_channels: Tuple[int, ...] = (128, 256, 512, 768), + encoder_block_out_channels: tuple[int] = (64, 128, 256, 384), + decoder_block_out_channels: tuple[int] = (128, 256, 512, 768), latent_channels: int = 12, - layers_per_block: Tuple[int, ...] = (3, 3, 4, 6, 3), + layers_per_block: tuple[int, ...] = (3, 3, 4, 6, 3), act_fn: str = "silu", - temporal_expansions: Tuple[int, ...] = (1, 2, 3), - spatial_expansions: Tuple[int, ...] = (2, 2, 2), - add_attention_block: Tuple[bool, ...] = (False, True, True, True, True), - latents_mean: Tuple[float, ...] = ( + temporal_expansions: tuple[int, ...] = (1, 2, 3), + spatial_expansions: tuple[int, ...] = (2, 2, 2), + add_attention_block: tuple[bool, ...] = (False, True, True, True, True), + latents_mean: tuple[float, ...] = ( -0.06730895953510081, -0.038011381506090416, -0.07477820912866141, @@ -710,7 +705,7 @@ def __init__( -0.011931556316503654, -0.0321993391887285, ), - latents_std: Tuple[float, ...] = ( + latents_std: tuple[float, ...] = ( 0.9263795028493863, 0.9248894543193766, 0.9393059390890617, @@ -790,10 +785,10 @@ def __init__( def enable_tiling( self, - tile_sample_min_height: Optional[int] = None, - tile_sample_min_width: Optional[int] = None, - tile_sample_stride_height: Optional[float] = None, - tile_sample_stride_width: Optional[float] = None, + tile_sample_min_height: int | None = None, + tile_sample_min_width: int | None = None, + tile_sample_stride_height: float | None = None, + tile_sample_stride_width: float | None = None, ) -> None: r""" Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to @@ -860,7 +855,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor: @apply_forward_hook def encode( self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: """ Encode a batch of images into latents. @@ -885,7 +880,7 @@ def encode( return (posterior,) return AutoencoderKLOutput(latent_dist=posterior) - def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: batch_size, num_channels, num_frames, height, width = z.shape tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio @@ -915,7 +910,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOut return DecoderOutput(sample=dec) @apply_forward_hook - def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: """ Decode a batch of images. @@ -1013,7 +1008,7 @@ def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width] return enc - def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: r""" Decode a batch of images using a tiled decoder. @@ -1096,8 +1091,8 @@ def forward( sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[torch.Tensor, torch.Tensor]: + generator: torch.Generator | None = None, + ) -> torch.Tensor | torch.Tensor: x = sample posterior = self.encode(x).latent_dist if sample_posterior: diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py index 7f7266146e6b..f2ca0f42a272 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py @@ -18,8 +18,6 @@ # - GitHub: https://github.com/Wan-Video/Wan2.1 # - Paper: https://huggingface.co/papers/2503.20314 -from typing import List, Optional, Tuple, Union - import torch import torch.nn as nn import torch.nn.functional as F @@ -58,9 +56,9 @@ def __init__( self, in_channels: int, out_channels: int, - kernel_size: Union[int, Tuple[int, int, int]], - stride: Union[int, Tuple[int, int, int]] = 1, - padding: Union[int, Tuple[int, int, int]] = 0, + kernel_size: int | tuple[int, int, int], + stride: int | tuple[int, int, int] = 1, + padding: int | tuple[int, int, int] = 0, ) -> None: super().__init__( in_channels=in_channels, @@ -497,7 +495,7 @@ def __init__( out_dim: int, num_res_blocks: int, dropout: float = 0.0, - upsample_mode: Optional[str] = None, + upsample_mode: str | None = None, non_linearity: str = "silu", ): super().__init__() @@ -681,14 +679,14 @@ def __init__( self, base_dim: int = 96, z_dim: int = 16, - dim_mult: Tuple[int, ...] = (1, 2, 4, 4), + dim_mult: list[int] = [1, 2, 4, 4], num_res_blocks: int = 2, - attn_scales: List[float] = [], - temperal_downsample: List[bool] = [False, True, True], + attn_scales: list[float] = [], + temperal_downsample: list[bool] = [False, True, True], dropout: float = 0.0, input_channels: int = 3, - latents_mean: List[float] = [-0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921], - latents_std: List[float] = [2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160], + latents_mean: list[float] = [-0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921], + latents_std: list[float] = [2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160], ) -> None: # fmt: on super().__init__() @@ -738,10 +736,10 @@ def __init__( def enable_tiling( self, - tile_sample_min_height: Optional[int] = None, - tile_sample_min_width: Optional[int] = None, - tile_sample_stride_height: Optional[float] = None, - tile_sample_stride_width: Optional[float] = None, + tile_sample_min_height: int | None = None, + tile_sample_min_width: int | None = None, + tile_sample_stride_height: float | None = None, + tile_sample_stride_width: float | None = None, ) -> None: r""" Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to @@ -809,7 +807,7 @@ def _encode(self, x: torch.Tensor): @apply_forward_hook def encode( self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: r""" Encode a batch of images into latents. @@ -859,7 +857,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True): return DecoderOutput(sample=out) @apply_forward_hook - def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: r""" Decode a batch of images. @@ -965,7 +963,7 @@ def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput: enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width] return enc - def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: r""" Decode a batch of images using a tiled decoder. @@ -1033,8 +1031,8 @@ def forward( sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[DecoderOutput, torch.Tensor]: + generator: torch.Generator | None = None, + ) -> DecoderOutput | torch.Tensor: """ Args: sample (`torch.Tensor`): Input sample. diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py b/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py index 7a307b1eacd8..95d4b0b7b535 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. import itertools -from typing import Optional, Tuple, Union import torch import torch.nn as nn @@ -32,7 +31,7 @@ def __init__( self, in_channels: int = 4, out_channels: int = 3, - block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), + block_out_channels: tuple[int] = (128, 256, 512, 512), layers_per_block: int = 2, ): super().__init__() @@ -146,10 +145,10 @@ class AutoencoderKLTemporalDecoder(ModelMixin, AttentionMixin, AutoencoderMixin, Parameters: in_channels (int, *optional*, defaults to 3): Number of channels in the input image. out_channels (int, *optional*, defaults to 3): Number of channels in the output. - down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`): - Tuple of downsample block types. - block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`): - Tuple of block output channels. + down_block_types (`tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`): + tuple of downsample block types. + block_out_channels (`tuple[int]`, *optional*, defaults to `(64,)`): + tuple of block output channels. layers_per_block: (`int`, *optional*, defaults to 1): Number of layers per block. latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space. sample_size (`int`, *optional*, defaults to `32`): Sample input size. @@ -173,8 +172,8 @@ def __init__( self, in_channels: int = 3, out_channels: int = 3, - down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",), - block_out_channels: Tuple[int, ...] = (64,), + down_block_types: tuple[str] = ("DownEncoderBlock2D",), + block_out_channels: tuple[int] = (64,), layers_per_block: int = 1, latent_channels: int = 4, sample_size: int = 32, @@ -219,7 +218,7 @@ def set_default_attn_processor(self): @apply_forward_hook def encode( self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: """ Encode a batch of images into latents. @@ -249,7 +248,7 @@ def decode( z: torch.Tensor, num_frames: int, return_dict: bool = True, - ) -> Union[DecoderOutput, torch.Tensor]: + ) -> DecoderOutput | torch.Tensor: """ Decode a batch of images. @@ -278,9 +277,9 @@ def forward( sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, num_frames: int = 1, - ) -> Union[DecoderOutput, torch.Tensor]: + ) -> DecoderOutput | torch.Tensor: r""" Args: sample (`torch.Tensor`): Input sample. diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py index 761dff2dc61a..ea5d2efe642f 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Optional, Tuple, Union - import torch import torch.nn as nn import torch.nn.functional as F @@ -149,9 +147,9 @@ def __init__( self, in_channels: int, out_channels: int, - kernel_size: Union[int, Tuple[int, int, int]], - stride: Union[int, Tuple[int, int, int]] = 1, - padding: Union[int, Tuple[int, int, int]] = 0, + kernel_size: int | tuple[int, int, int], + stride: int | tuple[int, int, int] = 1, + padding: int | tuple[int, int, int] = 0, ) -> None: super().__init__( in_channels=in_channels, @@ -730,7 +728,7 @@ def __init__( out_dim: int, num_res_blocks: int, dropout: float = 0.0, - upsample_mode: Optional[str] = None, + upsample_mode: str | None = None, non_linearity: str = "silu", ): super().__init__() @@ -971,14 +969,14 @@ class AutoencoderKLWan(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalMo def __init__( self, base_dim: int = 96, - decoder_base_dim: Optional[int] = None, + decoder_base_dim: int | None = None, z_dim: int = 16, - dim_mult: List[int] = [1, 2, 4, 4], + dim_mult: list[int] = [1, 2, 4, 4], num_res_blocks: int = 2, - attn_scales: List[float] = [], - temperal_downsample: List[bool] = [False, True, True], + attn_scales: list[float] = [], + temperal_downsample: list[bool] = [False, True, True], dropout: float = 0.0, - latents_mean: List[float] = [ + latents_mean: list[float] = [ -0.7571, -0.7089, -0.9113, @@ -996,7 +994,7 @@ def __init__( 0.2503, -0.2921, ], - latents_std: List[float] = [ + latents_std: list[float] = [ 2.8184, 1.4541, 2.3275, @@ -1017,9 +1015,9 @@ def __init__( is_residual: bool = False, in_channels: int = 3, out_channels: int = 3, - patch_size: Optional[int] = None, - scale_factor_temporal: Optional[int] = 4, - scale_factor_spatial: Optional[int] = 8, + patch_size: int | None = None, + scale_factor_temporal: int | None = 4, + scale_factor_spatial: int | None = 8, ) -> None: super().__init__() @@ -1087,10 +1085,10 @@ def __init__( def enable_tiling( self, - tile_sample_min_height: Optional[int] = None, - tile_sample_min_width: Optional[int] = None, - tile_sample_stride_height: Optional[float] = None, - tile_sample_stride_width: Optional[float] = None, + tile_sample_min_height: int | None = None, + tile_sample_min_width: int | None = None, + tile_sample_stride_height: float | None = None, + tile_sample_stride_width: float | None = None, ) -> None: r""" Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to @@ -1155,7 +1153,7 @@ def _encode(self, x: torch.Tensor): @apply_forward_hook def encode( self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: r""" Encode a batch of images into latents. @@ -1211,7 +1209,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True): return DecoderOutput(sample=out) @apply_forward_hook - def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: r""" Decode a batch of images. @@ -1323,7 +1321,7 @@ def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput: enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width] return enc - def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: r""" Decode a batch of images using a tiled decoder. @@ -1406,8 +1404,8 @@ def forward( sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[DecoderOutput, torch.Tensor]: + generator: torch.Generator | None = None, + ) -> DecoderOutput | torch.Tensor: """ Args: sample (`torch.Tensor`): Input sample. diff --git a/src/diffusers/models/autoencoders/autoencoder_oobleck.py b/src/diffusers/models/autoencoders/autoencoder_oobleck.py index d83264559209..239317cffd71 100644 --- a/src/diffusers/models/autoencoders/autoencoder_oobleck.py +++ b/src/diffusers/models/autoencoders/autoencoder_oobleck.py @@ -13,7 +13,6 @@ # limitations under the License. import math from dataclasses import dataclass -from typing import Optional, Tuple, Union import numpy as np import torch @@ -153,7 +152,7 @@ def __init__(self, parameters: torch.Tensor, deterministic: bool = False): self.logvar = torch.log(self.var) self.deterministic = deterministic - def sample(self, generator: Optional[torch.Generator] = None) -> torch.Tensor: + def sample(self, generator: torch.Generator | None = None) -> torch.Tensor: # make sure sample is on the same device as the parameters and has same dtype sample = randn_tensor( self.mean.shape, @@ -303,9 +302,9 @@ class AutoencoderOobleck(ModelMixin, AutoencoderMixin, ConfigMixin): Parameters: encoder_hidden_size (`int`, *optional*, defaults to 128): Intermediate representation dimension for the encoder. - downsampling_ratios (`List[int]`, *optional*, defaults to `[2, 4, 4, 8, 8]`): + downsampling_ratios (`list[int]`, *optional*, defaults to `[2, 4, 4, 8, 8]`): Ratios for downsampling in the encoder. These are used in reverse order for upsampling in the decoder. - channel_multiples (`List[int]`, *optional*, defaults to `[1, 2, 4, 8, 16]`): + channel_multiples (`list[int]`, *optional*, defaults to `[1, 2, 4, 8, 16]`): Multiples used to determine the hidden sizes of the hidden layers. decoder_channels (`int`, *optional*, defaults to 128): Intermediate representation dimension for the decoder. @@ -360,7 +359,7 @@ def __init__( @apply_forward_hook def encode( self, x: torch.Tensor, return_dict: bool = True - ) -> Union[AutoencoderOobleckOutput, Tuple[OobleckDiagonalGaussianDistribution]]: + ) -> AutoencoderOobleckOutput | tuple[OobleckDiagonalGaussianDistribution]: """ Encode a batch of images into latents. @@ -386,7 +385,7 @@ def encode( return AutoencoderOobleckOutput(latent_dist=posterior) - def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[OobleckDecoderOutput, torch.Tensor]: + def _decode(self, z: torch.Tensor, return_dict: bool = True) -> OobleckDecoderOutput | torch.Tensor: dec = self.decoder(z) if not return_dict: @@ -397,7 +396,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[OobleckDec @apply_forward_hook def decode( self, z: torch.FloatTensor, return_dict: bool = True, generator=None - ) -> Union[OobleckDecoderOutput, torch.FloatTensor]: + ) -> OobleckDecoderOutput | torch.FloatTensor: """ Decode a batch of images. @@ -428,8 +427,8 @@ def forward( sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[OobleckDecoderOutput, torch.Tensor]: + generator: torch.Generator | None = None, + ) -> OobleckDecoderOutput | torch.Tensor: r""" Args: sample (`torch.Tensor`): Input sample. diff --git a/src/diffusers/models/autoencoders/autoencoder_rae.py b/src/diffusers/models/autoencoders/autoencoder_rae.py new file mode 100644 index 000000000000..58ea66f8d18d --- /dev/null +++ b/src/diffusers/models/autoencoders/autoencoder_rae.py @@ -0,0 +1,689 @@ +# Copyright 2026 The NYU Vision-X and HuggingFace Teams. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from math import sqrt +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...utils import BaseOutput, logging +from ...utils.accelerate_utils import apply_forward_hook +from ...utils.import_utils import is_transformers_available +from ...utils.torch_utils import randn_tensor + + +if is_transformers_available(): + from transformers import ( + Dinov2WithRegistersConfig, + Dinov2WithRegistersModel, + SiglipVisionConfig, + SiglipVisionModel, + ViTMAEConfig, + ViTMAEModel, + ) + +from ..activations import get_activation +from ..attention import AttentionMixin +from ..attention_processor import Attention +from ..embeddings import get_2d_sincos_pos_embed +from ..modeling_utils import ModelMixin +from .vae import AutoencoderMixin, DecoderOutput, EncoderOutput + + +logger = logging.get_logger(__name__) + + +# --------------------------------------------------------------------------- +# Per-encoder forward functions +# --------------------------------------------------------------------------- +# Each function takes the raw transformers model + images and returns patch +# tokens of shape (B, N, C), stripping CLS / register tokens as needed. + + +def _dinov2_encoder_forward(model: nn.Module, images: torch.Tensor) -> torch.Tensor: + outputs = model(images, output_hidden_states=True) + unused_token_num = 5 # 1 CLS + 4 register tokens + return outputs.last_hidden_state[:, unused_token_num:] + + +def _siglip2_encoder_forward(model: nn.Module, images: torch.Tensor) -> torch.Tensor: + outputs = model(images, output_hidden_states=True, interpolate_pos_encoding=True) + return outputs.last_hidden_state + + +def _mae_encoder_forward(model: nn.Module, images: torch.Tensor, patch_size: int) -> torch.Tensor: + h, w = images.shape[2], images.shape[3] + patch_num = int(h * w // patch_size**2) + if patch_num * patch_size**2 != h * w: + raise ValueError("Image size should be divisible by patch size.") + noise = torch.arange(patch_num).unsqueeze(0).expand(images.shape[0], -1).to(images.device).to(images.dtype) + outputs = model(images, noise, interpolate_pos_encoding=True) + return outputs.last_hidden_state[:, 1:] # remove cls token + + +# --------------------------------------------------------------------------- +# Encoder construction helpers +# --------------------------------------------------------------------------- + + +def _build_encoder( + encoder_type: str, hidden_size: int, patch_size: int, num_hidden_layers: int, head_dim: int = 64 +) -> nn.Module: + """Build a frozen encoder from config (no pretrained download).""" + num_attention_heads = hidden_size // head_dim # all supported encoders use head_dim=64 + + if encoder_type == "dinov2": + config = Dinov2WithRegistersConfig( + hidden_size=hidden_size, + patch_size=patch_size, + image_size=518, + num_attention_heads=num_attention_heads, + num_hidden_layers=num_hidden_layers, + ) + model = Dinov2WithRegistersModel(config) + # RAE strips the final layernorm affine params (identity LN). Remove them from + # the architecture so `from_pretrained` doesn't leave them on the meta device. + model.layernorm.weight = None + model.layernorm.bias = None + elif encoder_type == "siglip2": + config = SiglipVisionConfig( + hidden_size=hidden_size, + patch_size=patch_size, + image_size=256, + num_attention_heads=num_attention_heads, + num_hidden_layers=num_hidden_layers, + ) + model = SiglipVisionModel(config) + # See dinov2 comment above. + model.vision_model.post_layernorm.weight = None + model.vision_model.post_layernorm.bias = None + elif encoder_type == "mae": + config = ViTMAEConfig( + hidden_size=hidden_size, + patch_size=patch_size, + image_size=224, + num_attention_heads=num_attention_heads, + num_hidden_layers=num_hidden_layers, + mask_ratio=0.0, + ) + model = ViTMAEModel(config) + # See dinov2 comment above. + model.layernorm.weight = None + model.layernorm.bias = None + else: + raise ValueError(f"Unknown encoder_type='{encoder_type}'. Available: dinov2, siglip2, mae") + + model.requires_grad_(False) + return model + + +_ENCODER_FORWARD_FNS = { + "dinov2": _dinov2_encoder_forward, + "siglip2": _siglip2_encoder_forward, + "mae": _mae_encoder_forward, +} + + +@dataclass +class RAEDecoderOutput(BaseOutput): + """ + Output of `RAEDecoder`. + + Args: + logits (`torch.Tensor`): + Patch reconstruction logits of shape `(batch_size, num_patches, patch_size**2 * num_channels)`. + """ + + logits: torch.Tensor + + +class ViTMAEIntermediate(nn.Module): + def __init__(self, hidden_size: int, intermediate_size: int, hidden_act: str = "gelu"): + super().__init__() + self.dense = nn.Linear(hidden_size, intermediate_size) + self.intermediate_act_fn = get_activation(hidden_act) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.intermediate_act_fn(hidden_states) + return hidden_states + + +class ViTMAEOutput(nn.Module): + def __init__(self, hidden_size: int, intermediate_size: int, hidden_dropout_prob: float = 0.0): + super().__init__() + self.dense = nn.Linear(intermediate_size, hidden_size) + self.dropout = nn.Dropout(hidden_dropout_prob) + + def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = hidden_states + input_tensor + return hidden_states + + +class ViTMAELayer(nn.Module): + """ + This matches the naming/parameter structure used in RAE-main (ViTMAE decoder block). + """ + + def __init__( + self, + *, + hidden_size: int, + num_attention_heads: int, + intermediate_size: int, + qkv_bias: bool = True, + layer_norm_eps: float = 1e-12, + hidden_dropout_prob: float = 0.0, + attention_probs_dropout_prob: float = 0.0, + hidden_act: str = "gelu", + ): + super().__init__() + if hidden_size % num_attention_heads != 0: + raise ValueError( + f"hidden_size={hidden_size} must be divisible by num_attention_heads={num_attention_heads}" + ) + self.attention = Attention( + query_dim=hidden_size, + heads=num_attention_heads, + dim_head=hidden_size // num_attention_heads, + dropout=attention_probs_dropout_prob, + bias=qkv_bias, + ) + self.intermediate = ViTMAEIntermediate( + hidden_size=hidden_size, intermediate_size=intermediate_size, hidden_act=hidden_act + ) + self.output = ViTMAEOutput( + hidden_size=hidden_size, intermediate_size=intermediate_size, hidden_dropout_prob=hidden_dropout_prob + ) + self.layernorm_before = nn.LayerNorm(hidden_size, eps=layer_norm_eps) + self.layernorm_after = nn.LayerNorm(hidden_size, eps=layer_norm_eps) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + attention_output = self.attention(self.layernorm_before(hidden_states)) + hidden_states = attention_output + hidden_states + + layer_output = self.layernorm_after(hidden_states) + layer_output = self.intermediate(layer_output) + layer_output = self.output(layer_output, hidden_states) + return layer_output + + +class RAEDecoder(nn.Module): + """ + Decoder implementation ported from RAE-main to keep checkpoint compatibility. + + Key attributes (must match checkpoint keys): + - decoder_embed + - decoder_pos_embed + - decoder_layers + - decoder_norm + - decoder_pred + - trainable_cls_token + """ + + def __init__( + self, + hidden_size: int = 768, + decoder_hidden_size: int = 512, + decoder_num_hidden_layers: int = 8, + decoder_num_attention_heads: int = 16, + decoder_intermediate_size: int = 2048, + num_patches: int = 256, + patch_size: int = 16, + num_channels: int = 3, + image_size: int = 256, + qkv_bias: bool = True, + layer_norm_eps: float = 1e-12, + hidden_dropout_prob: float = 0.0, + attention_probs_dropout_prob: float = 0.0, + hidden_act: str = "gelu", + ): + super().__init__() + self.decoder_hidden_size = decoder_hidden_size + self.patch_size = patch_size + self.num_channels = num_channels + self.image_size = image_size + self.num_patches = num_patches + + self.decoder_embed = nn.Linear(hidden_size, decoder_hidden_size, bias=True) + grid_size = int(num_patches**0.5) + pos_embed = get_2d_sincos_pos_embed( + decoder_hidden_size, grid_size, cls_token=True, extra_tokens=1, output_type="pt" + ) + self.register_buffer("decoder_pos_embed", pos_embed.unsqueeze(0).float(), persistent=False) + + self.decoder_layers = nn.ModuleList( + [ + ViTMAELayer( + hidden_size=decoder_hidden_size, + num_attention_heads=decoder_num_attention_heads, + intermediate_size=decoder_intermediate_size, + qkv_bias=qkv_bias, + layer_norm_eps=layer_norm_eps, + hidden_dropout_prob=hidden_dropout_prob, + attention_probs_dropout_prob=attention_probs_dropout_prob, + hidden_act=hidden_act, + ) + for _ in range(decoder_num_hidden_layers) + ] + ) + + self.decoder_norm = nn.LayerNorm(decoder_hidden_size, eps=layer_norm_eps) + self.decoder_pred = nn.Linear(decoder_hidden_size, patch_size**2 * num_channels, bias=True) + self.gradient_checkpointing = False + + self.trainable_cls_token = nn.Parameter(torch.zeros(1, 1, decoder_hidden_size)) + + def interpolate_pos_encoding(self, embeddings: torch.Tensor) -> torch.Tensor: + embeddings_positions = embeddings.shape[1] - 1 + num_positions = self.decoder_pos_embed.shape[1] - 1 + + class_pos_embed = self.decoder_pos_embed[:, 0, :] + patch_pos_embed = self.decoder_pos_embed[:, 1:, :] + dim = self.decoder_pos_embed.shape[-1] + + patch_pos_embed = patch_pos_embed.reshape(1, 1, -1, dim).permute(0, 3, 1, 2) + patch_pos_embed = F.interpolate( + patch_pos_embed, + scale_factor=(1, embeddings_positions / num_positions), + mode="bicubic", + align_corners=False, + ) + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1) + + def interpolate_latent(self, x: torch.Tensor) -> torch.Tensor: + b, l, c = x.shape + if l == self.num_patches: + return x + h = w = int(l**0.5) + x = x.reshape(b, h, w, c).permute(0, 3, 1, 2) + target_size = (int(self.num_patches**0.5), int(self.num_patches**0.5)) + x = F.interpolate(x, size=target_size, mode="bilinear", align_corners=False) + x = x.permute(0, 2, 3, 1).contiguous().view(b, self.num_patches, c) + return x + + def unpatchify(self, patchified_pixel_values: torch.Tensor, original_image_size: tuple[int, int] | None = None): + patch_size, num_channels = self.patch_size, self.num_channels + original_image_size = ( + original_image_size if original_image_size is not None else (self.image_size, self.image_size) + ) + original_height, original_width = original_image_size + num_patches_h = original_height // patch_size + num_patches_w = original_width // patch_size + if num_patches_h * num_patches_w != patchified_pixel_values.shape[1]: + raise ValueError( + f"The number of patches in the patchified pixel values {patchified_pixel_values.shape[1]}, does not match the number of patches on original image {num_patches_h}*{num_patches_w}" + ) + + batch_size = patchified_pixel_values.shape[0] + patchified_pixel_values = patchified_pixel_values.reshape( + batch_size, + num_patches_h, + num_patches_w, + patch_size, + patch_size, + num_channels, + ) + patchified_pixel_values = torch.einsum("nhwpqc->nchpwq", patchified_pixel_values) + pixel_values = patchified_pixel_values.reshape( + batch_size, + num_channels, + num_patches_h * patch_size, + num_patches_w * patch_size, + ) + return pixel_values + + def forward( + self, + hidden_states: torch.Tensor, + *, + interpolate_pos_encoding: bool = False, + drop_cls_token: bool = False, + return_dict: bool = True, + ) -> RAEDecoderOutput | tuple[torch.Tensor]: + x = self.decoder_embed(hidden_states) + if drop_cls_token: + x_ = x[:, 1:, :] + x_ = self.interpolate_latent(x_) + else: + x_ = self.interpolate_latent(x) + + cls_token = self.trainable_cls_token.expand(x_.shape[0], -1, -1) + x = torch.cat([cls_token, x_], dim=1) + + if interpolate_pos_encoding: + if not drop_cls_token: + raise ValueError("interpolate_pos_encoding only supports drop_cls_token=True") + decoder_pos_embed = self.interpolate_pos_encoding(x) + else: + decoder_pos_embed = self.decoder_pos_embed + + hidden_states = x + decoder_pos_embed.to(device=x.device, dtype=x.dtype) + + for layer_module in self.decoder_layers: + hidden_states = layer_module(hidden_states) + + hidden_states = self.decoder_norm(hidden_states) + logits = self.decoder_pred(hidden_states) + logits = logits[:, 1:, :] + + if not return_dict: + return (logits,) + return RAEDecoderOutput(logits=logits) + + +class AutoencoderRAE(ModelMixin, AttentionMixin, AutoencoderMixin, ConfigMixin): + r""" + Representation Autoencoder (RAE) model for encoding images to latents and decoding latents to images. + + This model uses a frozen pretrained encoder (DINOv2, SigLIP2, or MAE) with a trainable ViT decoder to reconstruct + images from learned representations. + + This model inherits from [`ModelMixin`]. Check the superclass documentation for its generic methods implemented for + all models (such as downloading or saving). + + Args: + encoder_type (`str`, *optional*, defaults to `"dinov2"`): + Type of frozen encoder to use. One of `"dinov2"`, `"siglip2"`, or `"mae"`. + encoder_hidden_size (`int`, *optional*, defaults to `768`): + Hidden size of the encoder model. + encoder_patch_size (`int`, *optional*, defaults to `14`): + Patch size of the encoder model. + encoder_num_hidden_layers (`int`, *optional*, defaults to `12`): + Number of hidden layers in the encoder model. + patch_size (`int`, *optional*, defaults to `16`): + Decoder patch size (used for unpatchify and decoder head). + encoder_input_size (`int`, *optional*, defaults to `224`): + Input size expected by the encoder. + image_size (`int`, *optional*): + Decoder output image size. If `None`, it is derived from encoder token count and `patch_size` like + RAE-main: `image_size = patch_size * sqrt(num_patches)`, where `num_patches = (encoder_input_size // + encoder_patch_size) ** 2`. + num_channels (`int`, *optional*, defaults to `3`): + Number of input/output channels. + encoder_norm_mean (`list`, *optional*, defaults to `[0.485, 0.456, 0.406]`): + Channel-wise mean for encoder input normalization (ImageNet defaults). + encoder_norm_std (`list`, *optional*, defaults to `[0.229, 0.224, 0.225]`): + Channel-wise std for encoder input normalization (ImageNet defaults). + latents_mean (`list` or `tuple`, *optional*): + Optional mean for latent normalization. Tensor inputs are accepted and converted to config-serializable + lists. + latents_std (`list` or `tuple`, *optional*): + Optional standard deviation for latent normalization. Tensor inputs are accepted and converted to + config-serializable lists. + noise_tau (`float`, *optional*, defaults to `0.0`): + Noise level for training (adds noise to latents during training). + reshape_to_2d (`bool`, *optional*, defaults to `True`): + Whether to reshape latents to 2D (B, C, H, W) format. + use_encoder_loss (`bool`, *optional*, defaults to `False`): + Whether to use encoder hidden states in the loss (for advanced training). + """ + + # NOTE: gradient checkpointing is not wired up for this model yet. + _supports_gradient_checkpointing = False + _no_split_modules = ["ViTMAELayer"] + _keys_to_ignore_on_load_unexpected = ["decoder.decoder_pos_embed"] + + @register_to_config + def __init__( + self, + encoder_type: str = "dinov2", + encoder_hidden_size: int = 768, + encoder_patch_size: int = 14, + encoder_num_hidden_layers: int = 12, + decoder_hidden_size: int = 512, + decoder_num_hidden_layers: int = 8, + decoder_num_attention_heads: int = 16, + decoder_intermediate_size: int = 2048, + patch_size: int = 16, + encoder_input_size: int = 224, + image_size: int | None = None, + num_channels: int = 3, + encoder_norm_mean: list | None = None, + encoder_norm_std: list | None = None, + latents_mean: list | tuple | torch.Tensor | None = None, + latents_std: list | tuple | torch.Tensor | None = None, + noise_tau: float = 0.0, + reshape_to_2d: bool = True, + use_encoder_loss: bool = False, + scaling_factor: float = 1.0, + ): + super().__init__() + + if encoder_type not in _ENCODER_FORWARD_FNS: + raise ValueError( + f"Unknown encoder_type='{encoder_type}'. Available: {sorted(_ENCODER_FORWARD_FNS.keys())}" + ) + + def _to_config_compatible(value: Any) -> Any: + if isinstance(value, torch.Tensor): + return value.detach().cpu().tolist() + if isinstance(value, tuple): + return [_to_config_compatible(v) for v in value] + if isinstance(value, list): + return [_to_config_compatible(v) for v in value] + return value + + def _as_optional_tensor(value: torch.Tensor | list | tuple | None) -> torch.Tensor | None: + if value is None: + return None + if isinstance(value, torch.Tensor): + return value.detach().clone() + return torch.tensor(value, dtype=torch.float32) + + latents_std_tensor = _as_optional_tensor(latents_std) + + # Ensure config values are JSON-serializable (list/None), even if caller passes torch.Tensors. + self.register_to_config( + latents_mean=_to_config_compatible(latents_mean), + latents_std=_to_config_compatible(latents_std), + ) + + self.encoder_input_size = encoder_input_size + self.noise_tau = float(noise_tau) + self.reshape_to_2d = bool(reshape_to_2d) + self.use_encoder_loss = bool(use_encoder_loss) + + # Validate early, before building the (potentially large) encoder/decoder. + encoder_patch_size = int(encoder_patch_size) + if self.encoder_input_size % encoder_patch_size != 0: + raise ValueError( + f"encoder_input_size={self.encoder_input_size} must be divisible by encoder_patch_size={encoder_patch_size}." + ) + decoder_patch_size = int(patch_size) + if decoder_patch_size <= 0: + raise ValueError("patch_size must be a positive integer (this is decoder_patch_size).") + + # Frozen representation encoder (built from config, no downloads) + self.encoder: nn.Module = _build_encoder( + encoder_type=encoder_type, + hidden_size=encoder_hidden_size, + patch_size=encoder_patch_size, + num_hidden_layers=encoder_num_hidden_layers, + ) + self._encoder_forward_fn = _ENCODER_FORWARD_FNS[encoder_type] + num_patches = (self.encoder_input_size // encoder_patch_size) ** 2 + + grid = int(sqrt(num_patches)) + if grid * grid != num_patches: + raise ValueError(f"Computed num_patches={num_patches} must be a perfect square.") + + derived_image_size = decoder_patch_size * grid + if image_size is None: + image_size = derived_image_size + else: + image_size = int(image_size) + if image_size != derived_image_size: + raise ValueError( + f"image_size={image_size} must equal decoder_patch_size*sqrt(num_patches)={derived_image_size} " + f"for patch_size={decoder_patch_size} and computed num_patches={num_patches}." + ) + + # Encoder input normalization stats (ImageNet defaults) + if encoder_norm_mean is None: + encoder_norm_mean = [0.485, 0.456, 0.406] + if encoder_norm_std is None: + encoder_norm_std = [0.229, 0.224, 0.225] + encoder_mean_tensor = torch.tensor(encoder_norm_mean, dtype=torch.float32).view(1, 3, 1, 1) + encoder_std_tensor = torch.tensor(encoder_norm_std, dtype=torch.float32).view(1, 3, 1, 1) + + self.register_buffer("encoder_mean", encoder_mean_tensor, persistent=True) + self.register_buffer("encoder_std", encoder_std_tensor, persistent=True) + + # Latent normalization buffers (defaults are no-ops; actual values come from checkpoint) + latents_mean_tensor = _as_optional_tensor(latents_mean) + if latents_mean_tensor is None: + latents_mean_tensor = torch.zeros(1) + self.register_buffer("_latents_mean", latents_mean_tensor, persistent=True) + + if latents_std_tensor is None: + latents_std_tensor = torch.ones(1) + self.register_buffer("_latents_std", latents_std_tensor, persistent=True) + + # ViT-MAE style decoder + self.decoder = RAEDecoder( + hidden_size=int(encoder_hidden_size), + decoder_hidden_size=int(decoder_hidden_size), + decoder_num_hidden_layers=int(decoder_num_hidden_layers), + decoder_num_attention_heads=int(decoder_num_attention_heads), + decoder_intermediate_size=int(decoder_intermediate_size), + num_patches=int(num_patches), + patch_size=int(decoder_patch_size), + num_channels=int(num_channels), + image_size=int(image_size), + ) + self.num_patches = int(num_patches) + self.decoder_patch_size = int(decoder_patch_size) + self.decoder_image_size = int(image_size) + + # Slicing support (batch dimension) similar to other diffusers autoencoders + self.use_slicing = False + + def _noising(self, x: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor: + # Per-sample random sigma in [0, noise_tau] + noise_sigma = self.noise_tau * torch.rand( + (x.size(0),) + (1,) * (x.ndim - 1), device=x.device, dtype=x.dtype, generator=generator + ) + return x + noise_sigma * randn_tensor(x.shape, generator=generator, device=x.device, dtype=x.dtype) + + def _resize_and_normalize(self, x: torch.Tensor) -> torch.Tensor: + _, _, h, w = x.shape + if h != self.encoder_input_size or w != self.encoder_input_size: + x = F.interpolate( + x, size=(self.encoder_input_size, self.encoder_input_size), mode="bicubic", align_corners=False + ) + mean = self.encoder_mean.to(device=x.device, dtype=x.dtype) + std = self.encoder_std.to(device=x.device, dtype=x.dtype) + return (x - mean) / std + + def _denormalize_image(self, x: torch.Tensor) -> torch.Tensor: + mean = self.encoder_mean.to(device=x.device, dtype=x.dtype) + std = self.encoder_std.to(device=x.device, dtype=x.dtype) + return x * std + mean + + def _normalize_latents(self, z: torch.Tensor) -> torch.Tensor: + latents_mean = self._latents_mean.to(device=z.device, dtype=z.dtype) + latents_std = self._latents_std.to(device=z.device, dtype=z.dtype) + return (z - latents_mean) / (latents_std + 1e-5) + + def _denormalize_latents(self, z: torch.Tensor) -> torch.Tensor: + latents_mean = self._latents_mean.to(device=z.device, dtype=z.dtype) + latents_std = self._latents_std.to(device=z.device, dtype=z.dtype) + return z * (latents_std + 1e-5) + latents_mean + + def _encode(self, x: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor: + x = self._resize_and_normalize(x) + + if self.config.encoder_type == "mae": + tokens = self._encoder_forward_fn(self.encoder, x, self.config.encoder_patch_size) + else: + tokens = self._encoder_forward_fn(self.encoder, x) # (B, N, C) + + if self.training and self.noise_tau > 0: + tokens = self._noising(tokens, generator=generator) + + if self.reshape_to_2d: + b, n, c = tokens.shape + side = int(sqrt(n)) + if side * side != n: + raise ValueError(f"Token length n={n} is not a perfect square; cannot reshape to 2D.") + z = tokens.transpose(1, 2).contiguous().view(b, c, side, side) # (B, C, h, w) + else: + z = tokens + + z = self._normalize_latents(z) + + # Follow diffusers convention: optionally scale latents for diffusion + if self.config.scaling_factor != 1.0: + z = z * self.config.scaling_factor + + return z + + @apply_forward_hook + def encode( + self, x: torch.Tensor, return_dict: bool = True, generator: torch.Generator | None = None + ) -> EncoderOutput | tuple[torch.Tensor]: + if self.use_slicing and x.shape[0] > 1: + latents = torch.cat([self._encode(x_slice, generator=generator) for x_slice in x.split(1)], dim=0) + else: + latents = self._encode(x, generator=generator) + + if not return_dict: + return (latents,) + return EncoderOutput(latent=latents) + + def _decode(self, z: torch.Tensor) -> torch.Tensor: + # Undo scaling factor if applied at encode time + if self.config.scaling_factor != 1.0: + z = z / self.config.scaling_factor + + z = self._denormalize_latents(z) + + if self.reshape_to_2d: + b, c, h, w = z.shape + tokens = z.view(b, c, h * w).transpose(1, 2).contiguous() # (B, N, C) + else: + tokens = z + + logits = self.decoder(tokens, return_dict=True).logits + x_rec = self.decoder.unpatchify(logits) + x_rec = self._denormalize_image(x_rec) + return x_rec.to(device=z.device) + + @apply_forward_hook + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]: + if self.use_slicing and z.shape[0] > 1: + decoded = torch.cat([self._decode(z_slice) for z_slice in z.split(1)], dim=0) + else: + decoded = self._decode(z) + + if not return_dict: + return (decoded,) + return DecoderOutput(sample=decoded) + + def forward( + self, sample: torch.Tensor, return_dict: bool = True, generator: torch.Generator | None = None + ) -> DecoderOutput | tuple[torch.Tensor]: + latents = self.encode(sample, return_dict=False, generator=generator)[0] + decoded = self.decode(latents, return_dict=False)[0] + if not return_dict: + return (decoded,) + return DecoderOutput(sample=decoded) diff --git a/src/diffusers/models/autoencoders/autoencoder_tiny.py b/src/diffusers/models/autoencoders/autoencoder_tiny.py index b9ac713d7392..14bc57ca9022 100644 --- a/src/diffusers/models/autoencoders/autoencoder_tiny.py +++ b/src/diffusers/models/autoencoders/autoencoder_tiny.py @@ -14,7 +14,6 @@ from dataclasses import dataclass -from typing import Optional, Tuple, Union import torch @@ -50,11 +49,11 @@ class AutoencoderTiny(ModelMixin, AutoencoderMixin, ConfigMixin): Parameters: in_channels (`int`, *optional*, defaults to 3): Number of channels in the input image. out_channels (`int`, *optional*, defaults to 3): Number of channels in the output. - encoder_block_out_channels (`Tuple[int]`, *optional*, defaults to `(64, 64, 64, 64)`): - Tuple of integers representing the number of output channels for each encoder block. The length of the + encoder_block_out_channels (`tuple[int]`, *optional*, defaults to `(64, 64, 64, 64)`): + tuple of integers representing the number of output channels for each encoder block. The length of the tuple should be equal to the number of encoder blocks. - decoder_block_out_channels (`Tuple[int]`, *optional*, defaults to `(64, 64, 64, 64)`): - Tuple of integers representing the number of output channels for each decoder block. The length of the + decoder_block_out_channels (`tuple[int]`, *optional*, defaults to `(64, 64, 64, 64)`): + tuple of integers representing the number of output channels for each decoder block. The length of the tuple should be equal to the number of decoder blocks. act_fn (`str`, *optional*, defaults to `"relu"`): Activation function to be used throughout the model. @@ -64,12 +63,12 @@ class AutoencoderTiny(ModelMixin, AutoencoderMixin, ConfigMixin): upsampling_scaling_factor (`int`, *optional*, defaults to 2): Scaling factor for upsampling in the decoder. It determines the size of the output image during the upsampling process. - num_encoder_blocks (`Tuple[int]`, *optional*, defaults to `(1, 3, 3, 3)`): - Tuple of integers representing the number of encoder blocks at each stage of the encoding process. The + num_encoder_blocks (`tuple[int]`, *optional*, defaults to `(1, 3, 3, 3)`): + tuple of integers representing the number of encoder blocks at each stage of the encoding process. The length of the tuple should be equal to the number of stages in the encoder. Each stage has a different number of encoder blocks. - num_decoder_blocks (`Tuple[int]`, *optional*, defaults to `(3, 3, 3, 1)`): - Tuple of integers representing the number of decoder blocks at each stage of the decoding process. The + num_decoder_blocks (`tuple[int]`, *optional*, defaults to `(3, 3, 3, 1)`): + tuple of integers representing the number of decoder blocks at each stage of the decoding process. The length of the tuple should be equal to the number of stages in the decoder. Each stage has a different number of decoder blocks. latent_magnitude (`float`, *optional*, defaults to 3.0): @@ -99,14 +98,14 @@ def __init__( self, in_channels: int = 3, out_channels: int = 3, - encoder_block_out_channels: Tuple[int, ...] = (64, 64, 64, 64), - decoder_block_out_channels: Tuple[int, ...] = (64, 64, 64, 64), + encoder_block_out_channels: tuple[int, ...] = (64, 64, 64, 64), + decoder_block_out_channels: tuple[int, ...] = (64, 64, 64, 64), act_fn: str = "relu", upsample_fn: str = "nearest", latent_channels: int = 4, upsampling_scaling_factor: int = 2, - num_encoder_blocks: Tuple[int, ...] = (1, 3, 3, 3), - num_decoder_blocks: Tuple[int, ...] = (3, 3, 3, 1), + num_encoder_blocks: tuple[int, ...] = (1, 3, 3, 3), + num_decoder_blocks: tuple[int, ...] = (3, 3, 3, 1), latent_magnitude: int = 3, latent_shift: float = 0.5, force_upcast: bool = False, @@ -258,7 +257,7 @@ def _tiled_decode(self, x: torch.Tensor) -> torch.Tensor: return out @apply_forward_hook - def encode(self, x: torch.Tensor, return_dict: bool = True) -> Union[AutoencoderTinyOutput, Tuple[torch.Tensor]]: + def encode(self, x: torch.Tensor, return_dict: bool = True) -> AutoencoderTinyOutput | tuple[torch.Tensor]: if self.use_slicing and x.shape[0] > 1: output = [ self._tiled_encode(x_slice) if self.use_tiling else self.encoder(x_slice) for x_slice in x.split(1) @@ -274,8 +273,8 @@ def encode(self, x: torch.Tensor, return_dict: bool = True) -> Union[Autoencoder @apply_forward_hook def decode( - self, x: torch.Tensor, generator: Optional[torch.Generator] = None, return_dict: bool = True - ) -> Union[DecoderOutput, Tuple[torch.Tensor]]: + self, x: torch.Tensor, generator: torch.Generator | None = None, return_dict: bool = True + ) -> DecoderOutput | tuple[torch.Tensor]: if self.use_slicing and x.shape[0] > 1: output = [ self._tiled_decode(x_slice) if self.use_tiling else self.decoder(x_slice) for x_slice in x.split(1) @@ -293,7 +292,7 @@ def forward( self, sample: torch.Tensor, return_dict: bool = True, - ) -> Union[DecoderOutput, Tuple[torch.Tensor]]: + ) -> DecoderOutput | tuple[torch.Tensor]: r""" Args: sample (`torch.Tensor`): Input sample. diff --git a/src/diffusers/models/autoencoders/consistency_decoder_vae.py b/src/diffusers/models/autoencoders/consistency_decoder_vae.py index db9404f4ac70..2d53b745e0fb 100644 --- a/src/diffusers/models/autoencoders/consistency_decoder_vae.py +++ b/src/diffusers/models/autoencoders/consistency_decoder_vae.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. from dataclasses import dataclass -from typing import Optional, Tuple, Union import torch import torch.nn.functional as F @@ -77,9 +76,9 @@ def __init__( latent_channels: int = 4, sample_size: int = 32, encoder_act_fn: str = "silu", - encoder_block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), + encoder_block_out_channels: tuple[int, ...] = (128, 256, 512, 512), encoder_double_z: bool = True, - encoder_down_block_types: Tuple[str, ...] = ( + encoder_down_block_types: tuple[str, ...] = ( "DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D", @@ -90,8 +89,8 @@ def __init__( encoder_norm_num_groups: int = 32, encoder_out_channels: int = 4, decoder_add_attention: bool = False, - decoder_block_out_channels: Tuple[int, ...] = (320, 640, 1024, 1024), - decoder_down_block_types: Tuple[str, ...] = ( + decoder_block_out_channels: tuple[int, ...] = (320, 640, 1024, 1024), + decoder_down_block_types: tuple[str, ...] = ( "ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D", @@ -106,7 +105,7 @@ def __init__( decoder_out_channels: int = 6, decoder_resnet_time_scale_shift: str = "scale_shift", decoder_time_embedding_type: str = "learned", - decoder_up_block_types: Tuple[str, ...] = ( + decoder_up_block_types: tuple[str, ...] = ( "ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D", @@ -186,7 +185,7 @@ def set_default_attn_processor(self): @apply_forward_hook def encode( self, x: torch.Tensor, return_dict: bool = True - ) -> Union[ConsistencyDecoderVAEOutput, Tuple[DiagonalGaussianDistribution]]: + ) -> ConsistencyDecoderVAEOutput | tuple[DiagonalGaussianDistribution]: """ Encode a batch of images into latents. @@ -222,21 +221,21 @@ def encode( def decode( self, z: torch.Tensor, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, return_dict: bool = True, num_inference_steps: int = 2, - ) -> Union[DecoderOutput, Tuple[torch.Tensor]]: + ) -> DecoderOutput | tuple[torch.Tensor]: """ Decodes the input latent vector `z` using the consistency decoder VAE model. Args: z (torch.Tensor): The input latent vector. - generator (Optional[torch.Generator]): The random number generator. Default is None. + generator (torch.Generator | None): The random number generator. Default is None. return_dict (bool): Whether to return the output as a dictionary. Default is True. num_inference_steps (int): The number of inference steps. Default is 2. Returns: - Union[DecoderOutput, Tuple[torch.Tensor]]: The decoded output. + DecoderOutput | tuple[torch.Tensor]: The decoded output. """ z = (z * self.config.scaling_factor - self.means) / self.stds @@ -279,7 +278,7 @@ def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch. b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent) return b - def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> Union[ConsistencyDecoderVAEOutput, Tuple]: + def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> ConsistencyDecoderVAEOutput | tuple: r"""Encode a batch of images using a tiled encoder. When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several @@ -339,8 +338,8 @@ def forward( sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True, - generator: Optional[torch.Generator] = None, - ) -> Union[DecoderOutput, Tuple[torch.Tensor]]: + generator: torch.Generator | None = None, + ) -> DecoderOutput | tuple[torch.Tensor]: r""" Args: sample (`torch.Tensor`): Input sample. diff --git a/src/diffusers/models/autoencoders/vae.py b/src/diffusers/models/autoencoders/vae.py index 9c6031a988f9..042cb5c10021 100644 --- a/src/diffusers/models/autoencoders/vae.py +++ b/src/diffusers/models/autoencoders/vae.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. from dataclasses import dataclass -from typing import Optional, Tuple import numpy as np import torch @@ -54,7 +53,7 @@ class DecoderOutput(BaseOutput): """ sample: torch.Tensor - commit_loss: Optional[torch.FloatTensor] = None + commit_loss: torch.FloatTensor | None = None class Encoder(nn.Module): @@ -66,10 +65,10 @@ class Encoder(nn.Module): The number of input channels. out_channels (`int`, *optional*, defaults to 3): The number of output channels. - down_block_types (`Tuple[str, ...]`, *optional*, defaults to `("DownEncoderBlock2D",)`): + down_block_types (`tuple[str, ...]`, *optional*, defaults to `("DownEncoderBlock2D",)`): The types of down blocks to use. See `~diffusers.models.unet_2d_blocks.get_down_block` for available options. - block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`): + block_out_channels (`tuple[int, ...]`, *optional*, defaults to `(64,)`): The number of output channels for each block. layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block. @@ -85,8 +84,8 @@ def __init__( self, in_channels: int = 3, out_channels: int = 3, - down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",), - block_out_channels: Tuple[int, ...] = (64,), + down_block_types: tuple[str, ...] = ("DownEncoderBlock2D",), + block_out_channels: tuple[int, ...] = (64,), layers_per_block: int = 2, norm_num_groups: int = 32, act_fn: str = "silu", @@ -187,9 +186,9 @@ class Decoder(nn.Module): The number of input channels. out_channels (`int`, *optional*, defaults to 3): The number of output channels. - up_block_types (`Tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`): + up_block_types (`tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`): The types of up blocks to use. See `~diffusers.models.unet_2d_blocks.get_up_block` for available options. - block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`): + block_out_channels (`tuple[int, ...]`, *optional*, defaults to `(64,)`): The number of output channels for each block. layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block. @@ -205,8 +204,8 @@ def __init__( self, in_channels: int = 3, out_channels: int = 3, - up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",), - block_out_channels: Tuple[int, ...] = (64,), + up_block_types: tuple[str, ...] = ("UpDecoderBlock2D",), + block_out_channels: tuple[int, ...] = (64,), layers_per_block: int = 2, norm_num_groups: int = 32, act_fn: str = "silu", @@ -280,7 +279,7 @@ def __init__( def forward( self, sample: torch.Tensor, - latent_embeds: Optional[torch.Tensor] = None, + latent_embeds: torch.Tensor | None = None, ) -> torch.Tensor: r"""The forward method of the `Decoder` class.""" @@ -402,9 +401,9 @@ class MaskConditionDecoder(nn.Module): The number of input channels. out_channels (`int`, *optional*, defaults to 3): The number of output channels. - up_block_types (`Tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`): + up_block_types (`tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`): The types of up blocks to use. See `~diffusers.models.unet_2d_blocks.get_up_block` for available options. - block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`): + block_out_channels (`tuple[int, ...]`, *optional*, defaults to `(64,)`): The number of output channels for each block. layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block. @@ -420,8 +419,8 @@ def __init__( self, in_channels: int = 3, out_channels: int = 3, - up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",), - block_out_channels: Tuple[int, ...] = (64,), + up_block_types: tuple[str, ...] = ("UpDecoderBlock2D",), + block_out_channels: tuple[int, ...] = (64,), layers_per_block: int = 2, norm_num_groups: int = 32, act_fn: str = "silu", @@ -500,9 +499,9 @@ def __init__( def forward( self, z: torch.Tensor, - image: Optional[torch.Tensor] = None, - mask: Optional[torch.Tensor] = None, - latent_embeds: Optional[torch.Tensor] = None, + image: torch.Tensor | None = None, + mask: torch.Tensor | None = None, + latent_embeds: torch.Tensor | None = None, ) -> torch.Tensor: r"""The forward method of the `MaskConditionDecoder` class.""" sample = z @@ -633,7 +632,7 @@ def unmap_to_all(self, inds: torch.LongTensor) -> torch.LongTensor: back = torch.gather(used[None, :][inds.shape[0] * [0], :], 1, inds) return back.reshape(ishape) - def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, Tuple]: + def forward(self, z: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, tuple]: # reshape z -> (batch, height, width, channel) and flatten z = z.permute(0, 2, 3, 1).contiguous() z_flattened = z.view(-1, self.vq_embed_dim) @@ -667,7 +666,7 @@ def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, Tuple]: return z_q, loss, (perplexity, min_encodings, min_encoding_indices) - def get_codebook_entry(self, indices: torch.LongTensor, shape: Tuple[int, ...]) -> torch.Tensor: + def get_codebook_entry(self, indices: torch.LongTensor, shape: tuple[int, ...]) -> torch.Tensor: # shape specifying (batch, height, width, channel) if self.remap is not None: indices = indices.reshape(shape[0], -1) # add batch axis @@ -698,7 +697,7 @@ def __init__(self, parameters: torch.Tensor, deterministic: bool = False): self.mean, device=self.parameters.device, dtype=self.parameters.dtype ) - def sample(self, generator: Optional[torch.Generator] = None) -> torch.Tensor: + def sample(self, generator: torch.Generator | None = None) -> torch.Tensor: # make sure sample is on the same device as the parameters and has same dtype sample = randn_tensor( self.mean.shape, @@ -728,7 +727,7 @@ def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor: dim=[1, 2, 3], ) - def nll(self, sample: torch.Tensor, dims: Tuple[int, ...] = [1, 2, 3]) -> torch.Tensor: + def nll(self, sample: torch.Tensor, dims: tuple[int, ...] = [1, 2, 3]) -> torch.Tensor: if self.deterministic: return torch.Tensor([0.0]) logtwopi = np.log(2.0 * np.pi) @@ -745,7 +744,7 @@ class IdentityDistribution(object): def __init__(self, parameters: torch.Tensor): self.parameters = parameters - def sample(self, generator: Optional[torch.Generator] = None) -> torch.Tensor: + def sample(self, generator: torch.Generator | None = None) -> torch.Tensor: return self.parameters def mode(self) -> torch.Tensor: @@ -761,10 +760,10 @@ class EncoderTiny(nn.Module): The number of input channels. out_channels (`int`): The number of output channels. - num_blocks (`Tuple[int, ...]`): + num_blocks (`tuple[int, ...]`): Each value of the tuple represents a Conv2d layer followed by `value` number of `AutoencoderTinyBlock`'s to use. - block_out_channels (`Tuple[int, ...]`): + block_out_channels (`tuple[int, ...]`): The number of output channels for each block. act_fn (`str`): The activation function to use. See `~diffusers.models.activations.get_activation` for available options. @@ -774,8 +773,8 @@ def __init__( self, in_channels: int, out_channels: int, - num_blocks: Tuple[int, ...], - block_out_channels: Tuple[int, ...], + num_blocks: tuple[int, ...], + block_out_channels: tuple[int, ...], act_fn: str, ): super().__init__() @@ -827,10 +826,10 @@ class DecoderTiny(nn.Module): The number of input channels. out_channels (`int`): The number of output channels. - num_blocks (`Tuple[int, ...]`): + num_blocks (`tuple[int, ...]`): Each value of the tuple represents a Conv2d layer followed by `value` number of `AutoencoderTinyBlock`'s to use. - block_out_channels (`Tuple[int, ...]`): + block_out_channels (`tuple[int, ...]`): The number of output channels for each block. upsampling_scaling_factor (`int`): The scaling factor to use for upsampling. @@ -842,8 +841,8 @@ def __init__( self, in_channels: int, out_channels: int, - num_blocks: Tuple[int, ...], - block_out_channels: Tuple[int, ...], + num_blocks: tuple[int, ...], + block_out_channels: tuple[int, ...], upsampling_scaling_factor: int, act_fn: str, upsample_fn: str, diff --git a/src/diffusers/models/autoencoders/vq_model.py b/src/diffusers/models/autoencoders/vq_model.py index 82436473dfc6..9214fb7faad6 100644 --- a/src/diffusers/models/autoencoders/vq_model.py +++ b/src/diffusers/models/autoencoders/vq_model.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. from dataclasses import dataclass -from typing import Optional, Tuple, Union import torch import torch.nn as nn @@ -48,12 +47,12 @@ class VQModel(ModelMixin, AutoencoderMixin, ConfigMixin): Parameters: in_channels (int, *optional*, defaults to 3): Number of channels in the input image. out_channels (int, *optional*, defaults to 3): Number of channels in the output. - down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`): - Tuple of downsample block types. - up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`): - Tuple of upsample block types. - block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`): - Tuple of block output channels. + down_block_types (`tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`): + tuple of downsample block types. + up_block_types (`tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`): + tuple of upsample block types. + block_out_channels (`tuple[int]`, *optional*, defaults to `(64,)`): + tuple of block output channels. layers_per_block (`int`, *optional*, defaults to `1`): Number of layers per block. act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. latent_channels (`int`, *optional*, defaults to `3`): Number of channels in the latent space. @@ -80,16 +79,16 @@ def __init__( self, in_channels: int = 3, out_channels: int = 3, - down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",), - up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",), - block_out_channels: Tuple[int, ...] = (64,), + down_block_types: tuple[str, ...] = ("DownEncoderBlock2D",), + up_block_types: tuple[str, ...] = ("UpDecoderBlock2D",), + block_out_channels: tuple[int, ...] = (64,), layers_per_block: int = 1, act_fn: str = "silu", latent_channels: int = 3, sample_size: int = 32, num_vq_embeddings: int = 256, norm_num_groups: int = 32, - vq_embed_dim: Optional[int] = None, + vq_embed_dim: int | None = None, scaling_factor: float = 0.18215, norm_type: str = "group", # group, spatial mid_block_add_attention=True, @@ -143,7 +142,7 @@ def encode(self, x: torch.Tensor, return_dict: bool = True) -> VQEncoderOutput: @apply_forward_hook def decode( self, h: torch.Tensor, force_not_quantize: bool = False, return_dict: bool = True, shape=None - ) -> Union[DecoderOutput, torch.Tensor]: + ) -> DecoderOutput | torch.Tensor: # also go through quantization layer if not force_not_quantize: quant, commit_loss, _ = self.quantize(h) @@ -161,9 +160,7 @@ def decode( return DecoderOutput(sample=dec, commit_loss=commit_loss) - def forward( - self, sample: torch.Tensor, return_dict: bool = True - ) -> Union[DecoderOutput, Tuple[torch.Tensor, ...]]: + def forward(self, sample: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor, ...]: r""" The [`VQModel`] forward method. diff --git a/src/diffusers/models/cache_utils.py b/src/diffusers/models/cache_utils.py index f4ad1af278f5..5f9587a1b4de 100644 --- a/src/diffusers/models/cache_utils.py +++ b/src/diffusers/models/cache_utils.py @@ -41,9 +41,11 @@ def enable_cache(self, config) -> None: Enable caching techniques on the model. Args: - config (`Union[PyramidAttentionBroadcastConfig]`): + config (`PyramidAttentionBroadcastConfig | FasterCacheConfig | FirstBlockCacheConfig`): The configuration for applying the caching technique. Currently supported caching techniques are: - [`~hooks.PyramidAttentionBroadcastConfig`] + - [`~hooks.FasterCacheConfig`] + - [`~hooks.FirstBlockCacheConfig`] Example: @@ -66,10 +68,12 @@ def enable_cache(self, config) -> None: from ..hooks import ( FasterCacheConfig, FirstBlockCacheConfig, + MagCacheConfig, PyramidAttentionBroadcastConfig, TaylorSeerCacheConfig, apply_faster_cache, apply_first_block_cache, + apply_mag_cache, apply_pyramid_attention_broadcast, apply_taylorseer_cache, ) @@ -83,6 +87,8 @@ def enable_cache(self, config) -> None: apply_faster_cache(self, config) elif isinstance(config, FirstBlockCacheConfig): apply_first_block_cache(self, config) + elif isinstance(config, MagCacheConfig): + apply_mag_cache(self, config) elif isinstance(config, PyramidAttentionBroadcastConfig): apply_pyramid_attention_broadcast(self, config) elif isinstance(config, TaylorSeerCacheConfig): @@ -97,11 +103,13 @@ def disable_cache(self) -> None: FasterCacheConfig, FirstBlockCacheConfig, HookRegistry, + MagCacheConfig, PyramidAttentionBroadcastConfig, TaylorSeerCacheConfig, ) from ..hooks.faster_cache import _FASTER_CACHE_BLOCK_HOOK, _FASTER_CACHE_DENOISER_HOOK from ..hooks.first_block_cache import _FBC_BLOCK_HOOK, _FBC_LEADER_BLOCK_HOOK + from ..hooks.mag_cache import _MAG_CACHE_BLOCK_HOOK, _MAG_CACHE_LEADER_BLOCK_HOOK from ..hooks.pyramid_attention_broadcast import _PYRAMID_ATTENTION_BROADCAST_HOOK from ..hooks.taylorseer_cache import _TAYLORSEER_CACHE_HOOK @@ -116,6 +124,9 @@ def disable_cache(self) -> None: elif isinstance(self._cache_config, FirstBlockCacheConfig): registry.remove_hook(_FBC_LEADER_BLOCK_HOOK, recurse=True) registry.remove_hook(_FBC_BLOCK_HOOK, recurse=True) + elif isinstance(self._cache_config, MagCacheConfig): + registry.remove_hook(_MAG_CACHE_LEADER_BLOCK_HOOK, recurse=True) + registry.remove_hook(_MAG_CACHE_BLOCK_HOOK, recurse=True) elif isinstance(self._cache_config, PyramidAttentionBroadcastConfig): registry.remove_hook(_PYRAMID_ATTENTION_BROADCAST_HOOK, recurse=True) elif isinstance(self._cache_config, TaylorSeerCacheConfig): diff --git a/src/diffusers/models/controlnet.py b/src/diffusers/models/controlnet.py deleted file mode 100644 index c18bd8751dcb..000000000000 --- a/src/diffusers/models/controlnet.py +++ /dev/null @@ -1,115 +0,0 @@ -# Copyright 2025 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -from typing import Optional, Tuple, Union - -from ..utils import deprecate -from .controlnets.controlnet import ( # noqa - ControlNetConditioningEmbedding, - ControlNetModel, - ControlNetOutput, - zero_module, -) - - -class ControlNetOutput(ControlNetOutput): - def __init__(self, *args, **kwargs): - deprecation_message = "Importing `ControlNetOutput` from `diffusers.models.controlnet` is deprecated and this will be removed in a future version. Please use `from diffusers.models.controlnets.controlnet import ControlNetOutput`, instead." - deprecate("diffusers.models.controlnet.ControlNetOutput", "0.34", deprecation_message) - super().__init__(*args, **kwargs) - - -class ControlNetModel(ControlNetModel): - def __init__( - self, - in_channels: int = 4, - conditioning_channels: int = 3, - flip_sin_to_cos: bool = True, - freq_shift: int = 0, - down_block_types: Tuple[str, ...] = ( - "CrossAttnDownBlock2D", - "CrossAttnDownBlock2D", - "CrossAttnDownBlock2D", - "DownBlock2D", - ), - mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn", - only_cross_attention: Union[bool, Tuple[bool]] = False, - block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280), - layers_per_block: int = 2, - downsample_padding: int = 1, - mid_block_scale_factor: float = 1, - act_fn: str = "silu", - norm_num_groups: Optional[int] = 32, - norm_eps: float = 1e-5, - cross_attention_dim: int = 1280, - transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1, - encoder_hid_dim: Optional[int] = None, - encoder_hid_dim_type: Optional[str] = None, - attention_head_dim: Union[int, Tuple[int, ...]] = 8, - num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None, - use_linear_projection: bool = False, - class_embed_type: Optional[str] = None, - addition_embed_type: Optional[str] = None, - addition_time_embed_dim: Optional[int] = None, - num_class_embeds: Optional[int] = None, - upcast_attention: bool = False, - resnet_time_scale_shift: str = "default", - projection_class_embeddings_input_dim: Optional[int] = None, - controlnet_conditioning_channel_order: str = "rgb", - conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256), - global_pool_conditions: bool = False, - addition_embed_type_num_heads: int = 64, - ): - deprecation_message = "Importing `ControlNetModel` from `diffusers.models.controlnet` is deprecated and this will be removed in a future version. Please use `from diffusers.models.controlnets.controlnet import ControlNetModel`, instead." - deprecate("diffusers.models.controlnet.ControlNetModel", "0.34", deprecation_message) - super().__init__( - in_channels=in_channels, - conditioning_channels=conditioning_channels, - flip_sin_to_cos=flip_sin_to_cos, - freq_shift=freq_shift, - down_block_types=down_block_types, - mid_block_type=mid_block_type, - only_cross_attention=only_cross_attention, - block_out_channels=block_out_channels, - layers_per_block=layers_per_block, - downsample_padding=downsample_padding, - mid_block_scale_factor=mid_block_scale_factor, - act_fn=act_fn, - norm_num_groups=norm_num_groups, - norm_eps=norm_eps, - cross_attention_dim=cross_attention_dim, - transformer_layers_per_block=transformer_layers_per_block, - encoder_hid_dim=encoder_hid_dim, - encoder_hid_dim_type=encoder_hid_dim_type, - attention_head_dim=attention_head_dim, - num_attention_heads=num_attention_heads, - use_linear_projection=use_linear_projection, - class_embed_type=class_embed_type, - addition_embed_type=addition_embed_type, - addition_time_embed_dim=addition_time_embed_dim, - num_class_embeds=num_class_embeds, - upcast_attention=upcast_attention, - resnet_time_scale_shift=resnet_time_scale_shift, - projection_class_embeddings_input_dim=projection_class_embeddings_input_dim, - controlnet_conditioning_channel_order=controlnet_conditioning_channel_order, - conditioning_embedding_out_channels=conditioning_embedding_out_channels, - global_pool_conditions=global_pool_conditions, - addition_embed_type_num_heads=addition_embed_type_num_heads, - ) - - -class ControlNetConditioningEmbedding(ControlNetConditioningEmbedding): - def __init__(self, *args, **kwargs): - deprecation_message = "Importing `ControlNetConditioningEmbedding` from `diffusers.models.controlnet` is deprecated and this will be removed in a future version. Please use `from diffusers.models.controlnets.controlnet import ControlNetConditioningEmbedding`, instead." - deprecate("diffusers.models.controlnet.ControlNetConditioningEmbedding", "0.34", deprecation_message) - super().__init__(*args, **kwargs) diff --git a/src/diffusers/models/controlnet_flux.py b/src/diffusers/models/controlnet_flux.py deleted file mode 100644 index e82748436d86..000000000000 --- a/src/diffusers/models/controlnet_flux.py +++ /dev/null @@ -1,70 +0,0 @@ -# Copyright 2025 Black Forest Labs, The HuggingFace Team and The InstantX Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - - -from typing import List - -from ..utils import deprecate, logging -from .controlnets.controlnet_flux import FluxControlNetModel, FluxControlNetOutput, FluxMultiControlNetModel - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class FluxControlNetOutput(FluxControlNetOutput): - def __init__(self, *args, **kwargs): - deprecation_message = "Importing `FluxControlNetOutput` from `diffusers.models.controlnet_flux` is deprecated and this will be removed in a future version. Please use `from diffusers.models.controlnets.controlnet_flux import FluxControlNetOutput`, instead." - deprecate("diffusers.models.controlnet_flux.FluxControlNetOutput", "0.34", deprecation_message) - super().__init__(*args, **kwargs) - - -class FluxControlNetModel(FluxControlNetModel): - def __init__( - self, - patch_size: int = 1, - in_channels: int = 64, - num_layers: int = 19, - num_single_layers: int = 38, - attention_head_dim: int = 128, - num_attention_heads: int = 24, - joint_attention_dim: int = 4096, - pooled_projection_dim: int = 768, - guidance_embeds: bool = False, - axes_dims_rope: List[int] = [16, 56, 56], - num_mode: int = None, - conditioning_embedding_channels: int = None, - ): - deprecation_message = "Importing `FluxControlNetModel` from `diffusers.models.controlnet_flux` is deprecated and this will be removed in a future version. Please use `from diffusers.models.controlnets.controlnet_flux import FluxControlNetModel`, instead." - deprecate("diffusers.models.controlnet_flux.FluxControlNetModel", "0.34", deprecation_message) - super().__init__( - patch_size=patch_size, - in_channels=in_channels, - num_layers=num_layers, - num_single_layers=num_single_layers, - attention_head_dim=attention_head_dim, - num_attention_heads=num_attention_heads, - joint_attention_dim=joint_attention_dim, - pooled_projection_dim=pooled_projection_dim, - guidance_embeds=guidance_embeds, - axes_dims_rope=axes_dims_rope, - num_mode=num_mode, - conditioning_embedding_channels=conditioning_embedding_channels, - ) - - -class FluxMultiControlNetModel(FluxMultiControlNetModel): - def __init__(self, *args, **kwargs): - deprecation_message = "Importing `FluxMultiControlNetModel` from `diffusers.models.controlnet_flux` is deprecated and this will be removed in a future version. Please use `from diffusers.models.controlnets.controlnet_flux import FluxMultiControlNetModel`, instead." - deprecate("diffusers.models.controlnet_flux.FluxMultiControlNetModel", "0.34", deprecation_message) - super().__init__(*args, **kwargs) diff --git a/src/diffusers/models/controlnet_sd3.py b/src/diffusers/models/controlnet_sd3.py deleted file mode 100644 index d239ad4eb3e8..000000000000 --- a/src/diffusers/models/controlnet_sd3.py +++ /dev/null @@ -1,68 +0,0 @@ -# Copyright 2025 Stability AI, The HuggingFace Team and The InstantX Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - - -from ..utils import deprecate, logging -from .controlnets.controlnet_sd3 import SD3ControlNetModel, SD3ControlNetOutput, SD3MultiControlNetModel - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class SD3ControlNetOutput(SD3ControlNetOutput): - def __init__(self, *args, **kwargs): - deprecation_message = "Importing `SD3ControlNetOutput` from `diffusers.models.controlnet_sd3` is deprecated and this will be removed in a future version. Please use `from diffusers.models.controlnets.controlnet_sd3 import SD3ControlNetOutput`, instead." - deprecate("diffusers.models.controlnet_sd3.SD3ControlNetOutput", "0.34", deprecation_message) - super().__init__(*args, **kwargs) - - -class SD3ControlNetModel(SD3ControlNetModel): - def __init__( - self, - sample_size: int = 128, - patch_size: int = 2, - in_channels: int = 16, - num_layers: int = 18, - attention_head_dim: int = 64, - num_attention_heads: int = 18, - joint_attention_dim: int = 4096, - caption_projection_dim: int = 1152, - pooled_projection_dim: int = 2048, - out_channels: int = 16, - pos_embed_max_size: int = 96, - extra_conditioning_channels: int = 0, - ): - deprecation_message = "Importing `SD3ControlNetModel` from `diffusers.models.controlnet_sd3` is deprecated and this will be removed in a future version. Please use `from diffusers.models.controlnets.controlnet_sd3 import SD3ControlNetModel`, instead." - deprecate("diffusers.models.controlnet_sd3.SD3ControlNetModel", "0.34", deprecation_message) - super().__init__( - sample_size=sample_size, - patch_size=patch_size, - in_channels=in_channels, - num_layers=num_layers, - attention_head_dim=attention_head_dim, - num_attention_heads=num_attention_heads, - joint_attention_dim=joint_attention_dim, - caption_projection_dim=caption_projection_dim, - pooled_projection_dim=pooled_projection_dim, - out_channels=out_channels, - pos_embed_max_size=pos_embed_max_size, - extra_conditioning_channels=extra_conditioning_channels, - ) - - -class SD3MultiControlNetModel(SD3MultiControlNetModel): - def __init__(self, *args, **kwargs): - deprecation_message = "Importing `SD3MultiControlNetModel` from `diffusers.models.controlnet_sd3` is deprecated and this will be removed in a future version. Please use `from diffusers.models.controlnets.controlnet_sd3 import SD3MultiControlNetModel`, instead." - deprecate("diffusers.models.controlnet_sd3.SD3MultiControlNetModel", "0.34", deprecation_message) - super().__init__(*args, **kwargs) diff --git a/src/diffusers/models/controlnet_sparsectrl.py b/src/diffusers/models/controlnet_sparsectrl.py deleted file mode 100644 index 5c67af4fe9c1..000000000000 --- a/src/diffusers/models/controlnet_sparsectrl.py +++ /dev/null @@ -1,116 +0,0 @@ -# Copyright 2025 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - - -from typing import Optional, Tuple, Union - -from ..utils import deprecate, logging -from .controlnets.controlnet_sparsectrl import ( # noqa - SparseControlNetConditioningEmbedding, - SparseControlNetModel, - SparseControlNetOutput, - zero_module, -) - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class SparseControlNetOutput(SparseControlNetOutput): - def __init__(self, *args, **kwargs): - deprecation_message = "Importing `SparseControlNetOutput` from `diffusers.models.controlnet_sparsectrl` is deprecated and this will be removed in a future version. Please use `from diffusers.models.controlnets.controlnet_sparsectrl import SparseControlNetOutput`, instead." - deprecate("diffusers.models.controlnet_sparsectrl.SparseControlNetOutput", "0.34", deprecation_message) - super().__init__(*args, **kwargs) - - -class SparseControlNetConditioningEmbedding(SparseControlNetConditioningEmbedding): - def __init__(self, *args, **kwargs): - deprecation_message = "Importing `SparseControlNetConditioningEmbedding` from `diffusers.models.controlnet_sparsectrl` is deprecated and this will be removed in a future version. Please use `from diffusers.models.controlnets.controlnet_sparsectrl import SparseControlNetConditioningEmbedding`, instead." - deprecate( - "diffusers.models.controlnet_sparsectrl.SparseControlNetConditioningEmbedding", "0.34", deprecation_message - ) - super().__init__(*args, **kwargs) - - -class SparseControlNetModel(SparseControlNetModel): - def __init__( - self, - in_channels: int = 4, - conditioning_channels: int = 4, - flip_sin_to_cos: bool = True, - freq_shift: int = 0, - down_block_types: Tuple[str, ...] = ( - "CrossAttnDownBlockMotion", - "CrossAttnDownBlockMotion", - "CrossAttnDownBlockMotion", - "DownBlockMotion", - ), - only_cross_attention: Union[bool, Tuple[bool]] = False, - block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280), - layers_per_block: int = 2, - downsample_padding: int = 1, - mid_block_scale_factor: float = 1, - act_fn: str = "silu", - norm_num_groups: Optional[int] = 32, - norm_eps: float = 1e-5, - cross_attention_dim: int = 768, - transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1, - transformer_layers_per_mid_block: Optional[Union[int, Tuple[int]]] = None, - temporal_transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1, - attention_head_dim: Union[int, Tuple[int, ...]] = 8, - num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None, - use_linear_projection: bool = False, - upcast_attention: bool = False, - resnet_time_scale_shift: str = "default", - conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256), - global_pool_conditions: bool = False, - controlnet_conditioning_channel_order: str = "rgb", - motion_max_seq_length: int = 32, - motion_num_attention_heads: int = 8, - concat_conditioning_mask: bool = True, - use_simplified_condition_embedding: bool = True, - ): - deprecation_message = "Importing `SparseControlNetModel` from `diffusers.models.controlnet_sparsectrl` is deprecated and this will be removed in a future version. Please use `from diffusers.models.controlnets.controlnet_sparsectrl import SparseControlNetModel`, instead." - deprecate("diffusers.models.controlnet_sparsectrl.SparseControlNetModel", "0.34", deprecation_message) - super().__init__( - in_channels=in_channels, - conditioning_channels=conditioning_channels, - flip_sin_to_cos=flip_sin_to_cos, - freq_shift=freq_shift, - down_block_types=down_block_types, - only_cross_attention=only_cross_attention, - block_out_channels=block_out_channels, - layers_per_block=layers_per_block, - downsample_padding=downsample_padding, - mid_block_scale_factor=mid_block_scale_factor, - act_fn=act_fn, - norm_num_groups=norm_num_groups, - norm_eps=norm_eps, - cross_attention_dim=cross_attention_dim, - transformer_layers_per_block=transformer_layers_per_block, - transformer_layers_per_mid_block=transformer_layers_per_mid_block, - temporal_transformer_layers_per_block=temporal_transformer_layers_per_block, - attention_head_dim=attention_head_dim, - num_attention_heads=num_attention_heads, - use_linear_projection=use_linear_projection, - upcast_attention=upcast_attention, - resnet_time_scale_shift=resnet_time_scale_shift, - conditioning_embedding_out_channels=conditioning_embedding_out_channels, - global_pool_conditions=global_pool_conditions, - controlnet_conditioning_channel_order=controlnet_conditioning_channel_order, - motion_max_seq_length=motion_max_seq_length, - motion_num_attention_heads=motion_num_attention_heads, - concat_conditioning_mask=concat_conditioning_mask, - use_simplified_condition_embedding=use_simplified_condition_embedding, - ) diff --git a/src/diffusers/models/controlnets/__init__.py b/src/diffusers/models/controlnets/__init__.py index fee7f231e899..853a2207f903 100644 --- a/src/diffusers/models/controlnets/__init__.py +++ b/src/diffusers/models/controlnets/__init__.py @@ -3,6 +3,7 @@ if is_torch_available(): from .controlnet import ControlNetModel, ControlNetOutput + from .controlnet_cosmos import CosmosControlNetModel from .controlnet_flux import FluxControlNetModel, FluxControlNetOutput, FluxMultiControlNetModel from .controlnet_hunyuan import ( HunyuanControlNetOutput, diff --git a/src/diffusers/models/controlnets/controlnet.py b/src/diffusers/models/controlnets/controlnet.py index 0b5b9fa3efba..8c2ff2fdd123 100644 --- a/src/diffusers/models/controlnets/controlnet.py +++ b/src/diffusers/models/controlnets/controlnet.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch from torch import nn @@ -21,7 +21,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin from ...loaders.single_file_model import FromOriginalModelMixin -from ...utils import BaseOutput, logging +from ...utils import BaseOutput, apply_lora_scale, logging from ..attention import AttentionMixin from ..attention_processor import ( ADDED_KV_ATTENTION_PROCESSORS, @@ -58,7 +58,7 @@ class ControlNetOutput(BaseOutput): Output can be used to condition the original UNet's middle block activation. """ - down_block_res_samples: Tuple[torch.Tensor] + down_block_res_samples: tuple[torch.Tensor] mid_block_res_sample: torch.Tensor @@ -76,7 +76,7 @@ def __init__( self, conditioning_embedding_channels: int, conditioning_channels: int = 3, - block_out_channels: Tuple[int, ...] = (16, 32, 96, 256), + block_out_channels: tuple[int, ...] = (16, 32, 96, 256), ): super().__init__() @@ -120,7 +120,7 @@ class ControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, FromOriginalModel The frequency shift to apply to the time embedding. down_block_types (`tuple[str]`, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`): The tuple of downsample blocks to use. - only_cross_attention (`Union[bool, Tuple[bool]]`, defaults to `False`): + only_cross_attention (`bool | tuple[bool]`, defaults to `False`): block_out_channels (`tuple[int]`, defaults to `(320, 640, 1280, 1280)`): The tuple of output channels for each block. layers_per_block (`int`, defaults to 2): @@ -138,7 +138,7 @@ class ControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, FromOriginalModel The epsilon to use for the normalization. cross_attention_dim (`int`, defaults to 1280): The dimension of the cross attention features. - transformer_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 1): + transformer_layers_per_block (`int` or `tuple[int]`, *optional*, defaults to 1): The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`], [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`]. @@ -148,7 +148,7 @@ class ControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, FromOriginalModel encoder_hid_dim_type (`str`, *optional*, defaults to `None`): If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`. - attention_head_dim (`Union[int, Tuple[int]]`, defaults to 8): + attention_head_dim (`int | tuple[int]`, defaults to 8): The dimension of the attention heads. use_linear_projection (`bool`, defaults to `False`): class_embed_type (`str`, *optional*, defaults to `None`): @@ -185,37 +185,37 @@ def __init__( conditioning_channels: int = 3, flip_sin_to_cos: bool = True, freq_shift: int = 0, - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str, ...] = ( "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D", ), - mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn", - only_cross_attention: Union[bool, Tuple[bool]] = False, - block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280), + mid_block_type: str | None = "UNetMidBlock2DCrossAttn", + only_cross_attention: bool | tuple[bool] = False, + block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280), layers_per_block: int = 2, downsample_padding: int = 1, mid_block_scale_factor: float = 1, act_fn: str = "silu", - norm_num_groups: Optional[int] = 32, + norm_num_groups: int | None = 32, norm_eps: float = 1e-5, cross_attention_dim: int = 1280, - transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1, - encoder_hid_dim: Optional[int] = None, - encoder_hid_dim_type: Optional[str] = None, - attention_head_dim: Union[int, Tuple[int, ...]] = 8, - num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None, + transformer_layers_per_block: int | tuple[int, ...] = 1, + encoder_hid_dim: int | None = None, + encoder_hid_dim_type: str | None = None, + attention_head_dim: int | tuple[int, ...] = 8, + num_attention_heads: int | tuple[int, ...] | None = None, use_linear_projection: bool = False, - class_embed_type: Optional[str] = None, - addition_embed_type: Optional[str] = None, - addition_time_embed_dim: Optional[int] = None, - num_class_embeds: Optional[int] = None, + class_embed_type: str | None = None, + addition_embed_type: str | None = None, + addition_time_embed_dim: int | None = None, + num_class_embeds: int | None = None, upcast_attention: bool = False, resnet_time_scale_shift: str = "default", - projection_class_embeddings_input_dim: Optional[int] = None, + projection_class_embeddings_input_dim: int | None = None, controlnet_conditioning_channel_order: str = "rgb", - conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256), + conditioning_embedding_out_channels: tuple[int, ...] | None = (16, 32, 96, 256), global_pool_conditions: bool = False, addition_embed_type_num_heads: int = 64, ): @@ -445,7 +445,7 @@ def from_unet( cls, unet: UNet2DConditionModel, controlnet_conditioning_channel_order: str = "rgb", - conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256), + conditioning_embedding_out_channels: tuple[int, ...] | None = (16, 32, 96, 256), load_weights_from_unet: bool = True, conditioning_channels: int = 3, ): @@ -533,7 +533,7 @@ def set_default_attn_processor(self): self.set_attn_processor(processor) # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attention_slice - def set_attention_slice(self, slice_size: Union[str, int, List[int]]) -> None: + def set_attention_slice(self, slice_size: str | int | list[int]) -> None: r""" Enable sliced attention computation. @@ -587,7 +587,7 @@ def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): # Recursively walk through all the children. # Any children which exposes the set_attention_slice method # gets the message - def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]): + def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: list[int]): if hasattr(module, "set_attention_slice"): module.set_attention_slice(slice_size.pop()) @@ -598,28 +598,29 @@ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[i for module in self.children(): fn_recursive_set_attention_slice(module, reversed_slice_size) + @apply_lora_scale("cross_attention_kwargs") def forward( self, sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], + timestep: torch.Tensor | float | int, encoder_hidden_states: torch.Tensor, controlnet_cond: torch.Tensor, conditioning_scale: float = 1.0, - class_labels: Optional[torch.Tensor] = None, - timestep_cond: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + class_labels: torch.Tensor | None = None, + timestep_cond: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + added_cond_kwargs: dict[str, torch.Tensor] | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, guess_mode: bool = False, return_dict: bool = True, - ) -> Union[ControlNetOutput, Tuple[Tuple[torch.Tensor, ...], torch.Tensor]]: + ) -> ControlNetOutput | tuple[tuple[torch.Tensor, ...], torch.Tensor]: """ The [`ControlNetModel`] forward method. Args: sample (`torch.Tensor`): The noisy input tensor. - timestep (`Union[torch.Tensor, float, int]`): + timestep (`torch.Tensor | float | int`): The number of timesteps to denoise an input. encoder_hidden_states (`torch.Tensor`): The encoder hidden states. diff --git a/src/diffusers/models/controlnets/controlnet_cosmos.py b/src/diffusers/models/controlnets/controlnet_cosmos.py new file mode 100644 index 000000000000..e39f8dfb568a --- /dev/null +++ b/src/diffusers/models/controlnets/controlnet_cosmos.py @@ -0,0 +1,317 @@ +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn as nn + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin +from ...utils import BaseOutput, is_torchvision_available, logging +from ..modeling_utils import ModelMixin +from ..transformers.transformer_cosmos import ( + CosmosEmbedding, + CosmosLearnablePositionalEmbed, + CosmosPatchEmbed, + CosmosRotaryPosEmbed, + CosmosTransformerBlock, +) + + +if is_torchvision_available(): + from torchvision import transforms + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class CosmosControlNetOutput(BaseOutput): + """ + Output of [`CosmosControlNetModel`]. + + Args: + control_block_samples (`list[torch.Tensor]`): + List of control block activations to be injected into transformer blocks. + """ + + control_block_samples: List[torch.Tensor] + + +class CosmosControlNetModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): + r""" + ControlNet for Cosmos Transfer2.5. + + This model duplicates the shared embedding modules from the transformer (patch_embed, time_embed, + learnable_pos_embed, img_context_proj) to enable proper CPU offloading. The forward() method computes everything + internally from raw inputs. + """ + + _supports_gradient_checkpointing = True + _skip_layerwise_casting_patterns = ["patch_embed", "patch_embed_base", "time_embed"] + _no_split_modules = ["CosmosTransformerBlock"] + _keep_in_fp32_modules = ["learnable_pos_embed"] + + @register_to_config + def __init__( + self, + n_controlnet_blocks: int = 4, + in_channels: int = 130, + latent_channels: int = 18, # base latent channels (latents + condition_mask) + padding_mask + model_channels: int = 2048, + num_attention_heads: int = 32, + attention_head_dim: int = 128, + mlp_ratio: float = 4.0, + text_embed_dim: int = 1024, + adaln_lora_dim: int = 256, + patch_size: Tuple[int, int, int] = (1, 2, 2), + max_size: Tuple[int, int, int] = (128, 240, 240), + rope_scale: Tuple[float, float, float] = (2.0, 1.0, 1.0), + extra_pos_embed_type: str | None = None, + img_context_dim_in: int | None = None, + img_context_dim_out: int = 2048, + use_crossattn_projection: bool = False, + crossattn_proj_in_channels: int = 1024, + encoder_hidden_states_channels: int = 1024, + ): + super().__init__() + + self.patch_embed = CosmosPatchEmbed(in_channels, model_channels, patch_size, bias=False) + + self.patch_embed_base = CosmosPatchEmbed(latent_channels, model_channels, patch_size, bias=False) + self.time_embed = CosmosEmbedding(model_channels, model_channels) + + self.learnable_pos_embed = None + if extra_pos_embed_type == "learnable": + self.learnable_pos_embed = CosmosLearnablePositionalEmbed( + hidden_size=model_channels, + max_size=max_size, + patch_size=patch_size, + ) + + self.img_context_proj = None + if img_context_dim_in is not None and img_context_dim_in > 0: + self.img_context_proj = nn.Sequential( + nn.Linear(img_context_dim_in, img_context_dim_out, bias=True), + nn.GELU(), + ) + + # Cross-attention projection for text embeddings (same as transformer) + self.crossattn_proj = None + if use_crossattn_projection: + self.crossattn_proj = nn.Sequential( + nn.Linear(crossattn_proj_in_channels, encoder_hidden_states_channels, bias=True), + nn.GELU(), + ) + + # RoPE for both control and base latents + self.rope = CosmosRotaryPosEmbed( + hidden_size=attention_head_dim, max_size=max_size, patch_size=patch_size, rope_scale=rope_scale + ) + + self.control_blocks = nn.ModuleList( + [ + CosmosTransformerBlock( + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + cross_attention_dim=text_embed_dim, + mlp_ratio=mlp_ratio, + adaln_lora_dim=adaln_lora_dim, + qk_norm="rms_norm", + out_bias=False, + img_context=img_context_dim_in is not None and img_context_dim_in > 0, + before_proj=(block_idx == 0), + after_proj=True, + ) + for block_idx in range(n_controlnet_blocks) + ] + ) + + self.gradient_checkpointing = False + + def _expand_conditioning_scale(self, conditioning_scale: float | list[float]) -> List[float]: + if isinstance(conditioning_scale, list): + scales = conditioning_scale + else: + scales = [conditioning_scale] * len(self.control_blocks) + + if len(scales) < len(self.control_blocks): + logger.warning( + "Received %d control scales, but control network defines %d blocks. " + "Scales will be trimmed or repeated to match.", + len(scales), + len(self.control_blocks), + ) + scales = (scales * len(self.control_blocks))[: len(self.control_blocks)] + return scales + + def forward( + self, + controls_latents: torch.Tensor, + latents: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: Union[Optional[torch.Tensor], Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]], + condition_mask: torch.Tensor, + conditioning_scale: float | list[float] = 1.0, + padding_mask: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + fps: int | None = None, + return_dict: bool = True, + ) -> Union[CosmosControlNetOutput, Tuple[List[torch.Tensor]]]: + """ + Forward pass for the ControlNet. + + Args: + controls_latents: Control signal latents [B, C, T, H, W] + latents: Base latents from the noising process [B, C, T, H, W] + timestep: Diffusion timestep tensor + encoder_hidden_states: Tuple of (text_context, img_context) or text_context + condition_mask: Conditioning mask [B, 1, T, H, W] + conditioning_scale: Scale factor(s) for control outputs + padding_mask: Padding mask [B, 1, H, W] or None + attention_mask: Optional attention mask or None + fps: Frames per second for RoPE or None + return_dict: Whether to return a CosmosControlNetOutput or a tuple + + Returns: + CosmosControlNetOutput or tuple of control tensors + """ + B, C, T, H, W = controls_latents.shape + + # 1. Prepare control latents + control_hidden_states = controls_latents + vace_in_channels = self.config.in_channels - 1 + if control_hidden_states.shape[1] < vace_in_channels - 1: + pad_C = vace_in_channels - 1 - control_hidden_states.shape[1] + control_hidden_states = torch.cat( + [ + control_hidden_states, + torch.zeros( + (B, pad_C, T, H, W), dtype=control_hidden_states.dtype, device=control_hidden_states.device + ), + ], + dim=1, + ) + + if condition_mask is not None: + control_hidden_states = torch.cat([control_hidden_states, condition_mask], dim=1) + else: + control_hidden_states = torch.cat( + [control_hidden_states, torch.zeros_like(controls_latents[:, :1])], dim=1 + ) + + padding_mask_resized = transforms.functional.resize( + padding_mask, list(control_hidden_states.shape[-2:]), interpolation=transforms.InterpolationMode.NEAREST + ) + control_hidden_states = torch.cat( + [control_hidden_states, padding_mask_resized.unsqueeze(2).repeat(B, 1, T, 1, 1)], dim=1 + ) + + # 2. Prepare base latents (same processing as transformer.forward) + base_hidden_states = latents + if condition_mask is not None: + base_hidden_states = torch.cat([base_hidden_states, condition_mask], dim=1) + + base_padding_mask = transforms.functional.resize( + padding_mask, list(base_hidden_states.shape[-2:]), interpolation=transforms.InterpolationMode.NEAREST + ) + base_hidden_states = torch.cat( + [base_hidden_states, base_padding_mask.unsqueeze(2).repeat(B, 1, T, 1, 1)], dim=1 + ) + + # 3. Generate positional embeddings (shared for both) + image_rotary_emb = self.rope(control_hidden_states, fps=fps) + extra_pos_emb = self.learnable_pos_embed(control_hidden_states) if self.learnable_pos_embed else None + + # 4. Patchify control latents + control_hidden_states = self.patch_embed(control_hidden_states) + control_hidden_states = control_hidden_states.flatten(1, 3) + + # 5. Patchify base latents + p_t, p_h, p_w = self.config.patch_size + post_patch_num_frames = T // p_t + post_patch_height = H // p_h + post_patch_width = W // p_w + + base_hidden_states = self.patch_embed_base(base_hidden_states) + base_hidden_states = base_hidden_states.flatten(1, 3) + + # 6. Time embeddings + if timestep.ndim == 1: + temb, embedded_timestep = self.time_embed(base_hidden_states, timestep) + elif timestep.ndim == 5: + batch_size, _, num_frames, _, _ = latents.shape + assert timestep.shape == (batch_size, 1, num_frames, 1, 1), ( + f"Expected timestep to have shape [B, 1, T, 1, 1], but got {timestep.shape}" + ) + timestep_flat = timestep.flatten() + temb, embedded_timestep = self.time_embed(base_hidden_states, timestep_flat) + temb, embedded_timestep = ( + x.view(batch_size, post_patch_num_frames, 1, 1, -1) + .expand(-1, -1, post_patch_height, post_patch_width, -1) + .flatten(1, 3) + for x in (temb, embedded_timestep) + ) + else: + raise ValueError(f"Expected timestep to have shape [B, 1, T, 1, 1] or [T], but got {timestep.shape}") + + # 7. Process encoder hidden states + if isinstance(encoder_hidden_states, tuple): + text_context, img_context = encoder_hidden_states + else: + text_context = encoder_hidden_states + img_context = None + + # Apply cross-attention projection to text context + if self.crossattn_proj is not None: + text_context = self.crossattn_proj(text_context) + + # Apply cross-attention projection to image context (if provided) + if img_context is not None and self.img_context_proj is not None: + img_context = self.img_context_proj(img_context) + + # Combine text and image context into a single tuple + if self.config.img_context_dim_in is not None and self.config.img_context_dim_in > 0: + processed_encoder_hidden_states = (text_context, img_context) + else: + processed_encoder_hidden_states = text_context + + # 8. Prepare attention mask + if attention_mask is not None: + attention_mask = attention_mask.unsqueeze(1).unsqueeze(1) # [B, 1, 1, S] + + # 9. Run control blocks + scales = self._expand_conditioning_scale(conditioning_scale) + result = [] + for block_idx, (block, scale) in enumerate(zip(self.control_blocks, scales)): + if torch.is_grad_enabled() and self.gradient_checkpointing: + control_hidden_states, control_proj = self._gradient_checkpointing_func( + block, + control_hidden_states, + processed_encoder_hidden_states, + embedded_timestep, + temb, + image_rotary_emb, + extra_pos_emb, + attention_mask, + None, # controlnet_residual + base_hidden_states, + block_idx, + ) + else: + control_hidden_states, control_proj = block( + hidden_states=control_hidden_states, + encoder_hidden_states=processed_encoder_hidden_states, + embedded_timestep=embedded_timestep, + temb=temb, + image_rotary_emb=image_rotary_emb, + extra_pos_emb=extra_pos_emb, + attention_mask=attention_mask, + controlnet_residual=None, + latents=base_hidden_states, + block_idx=block_idx, + ) + result.append(control_proj * scale) + + if not return_dict: + return (result,) + + return CosmosControlNetOutput(control_block_samples=result) diff --git a/src/diffusers/models/controlnets/controlnet_flax.py b/src/diffusers/models/controlnets/controlnet_flax.py index f7a8b98fa2f0..5bc1a446338e 100644 --- a/src/diffusers/models/controlnets/controlnet_flax.py +++ b/src/diffusers/models/controlnets/controlnet_flax.py @@ -11,8 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Tuple, Union - import flax import flax.linen as nn import jax @@ -49,7 +47,7 @@ class FlaxControlNetOutput(BaseOutput): class FlaxControlNetConditioningEmbedding(nn.Module): conditioning_embedding_channels: int - block_out_channels: Tuple[int, ...] = (16, 32, 96, 256) + block_out_channels: tuple[int, ...] = (16, 32, 96, 256) dtype: jnp.dtype = jnp.float32 def setup(self) -> None: @@ -132,15 +130,15 @@ class FlaxControlNetModel(nn.Module, FlaxModelMixin, ConfigMixin): The size of the input sample. in_channels (`int`, *optional*, defaults to 4): The number of channels in the input sample. - down_block_types (`Tuple[str]`, *optional*, defaults to `("FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxDownBlock2D")`): + down_block_types (`tuple[str]`, *optional*, defaults to `("FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxDownBlock2D")`): The tuple of downsample blocks to use. - block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): + block_out_channels (`tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): The tuple of output channels for each block. layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block. - attention_head_dim (`int` or `Tuple[int]`, *optional*, defaults to 8): + attention_head_dim (`int` or `tuple[int]`, *optional*, defaults to 8): The dimension of the attention heads. - num_attention_heads (`int` or `Tuple[int]`, *optional*): + num_attention_heads (`int` or `tuple[int]`, *optional*): The number of attention heads. cross_attention_dim (`int`, *optional*, defaults to 768): The dimension of the cross attention features. @@ -157,17 +155,17 @@ class FlaxControlNetModel(nn.Module, FlaxModelMixin, ConfigMixin): sample_size: int = 32 in_channels: int = 4 - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str, ...] = ( "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D", ) - only_cross_attention: Union[bool, Tuple[bool, ...]] = False - block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280) + only_cross_attention: bool | tuple[bool, ...] = False + block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280) layers_per_block: int = 2 - attention_head_dim: Union[int, Tuple[int, ...]] = 8 - num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None + attention_head_dim: int | tuple[int, ...] = 8 + num_attention_heads: int | tuple[int, ...] | None = None cross_attention_dim: int = 1280 dropout: float = 0.0 use_linear_projection: bool = False @@ -175,7 +173,7 @@ class FlaxControlNetModel(nn.Module, FlaxModelMixin, ConfigMixin): flip_sin_to_cos: bool = True freq_shift: int = 0 controlnet_conditioning_channel_order: str = "rgb" - conditioning_embedding_out_channels: Tuple[int, ...] = (16, 32, 96, 256) + conditioning_embedding_out_channels: tuple[int, ...] = (16, 32, 96, 256) def init_weights(self, rng: jax.Array) -> FrozenDict: # init input tensors @@ -327,13 +325,13 @@ def setup(self) -> None: def __call__( self, sample: jnp.ndarray, - timesteps: Union[jnp.ndarray, float, int], + timesteps: jnp.ndarray | float | int, encoder_hidden_states: jnp.ndarray, controlnet_cond: jnp.ndarray, conditioning_scale: float = 1.0, return_dict: bool = True, train: bool = False, - ) -> Union[FlaxControlNetOutput, Tuple[Tuple[jnp.ndarray, ...], jnp.ndarray]]: + ) -> FlaxControlNetOutput | tuple[tuple[jnp.ndarray, ...], jnp.ndarray]: r""" Args: sample (`jnp.ndarray`): (batch, channel, height, width) noisy inputs tensor diff --git a/src/diffusers/models/controlnets/controlnet_flux.py b/src/diffusers/models/controlnets/controlnet_flux.py index 639a8ad7390a..56482c299c05 100644 --- a/src/diffusers/models/controlnets/controlnet_flux.py +++ b/src/diffusers/models/controlnets/controlnet_flux.py @@ -13,14 +13,18 @@ # limitations under the License. from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, BaseOutput, logging, scale_lora_layers, unscale_lora_layers +from ...utils import ( + BaseOutput, + apply_lora_scale, + logging, +) from ..attention import AttentionMixin from ..controlnets.controlnet import ControlNetConditioningEmbedding, zero_module from ..embeddings import CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings, FluxPosEmbed @@ -34,8 +38,8 @@ @dataclass class FluxControlNetOutput(BaseOutput): - controlnet_block_samples: Tuple[torch.Tensor] - controlnet_single_block_samples: Tuple[torch.Tensor] + controlnet_block_samples: tuple[torch.Tensor] + controlnet_single_block_samples: tuple[torch.Tensor] class FluxControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, PeftAdapterMixin): @@ -53,7 +57,7 @@ def __init__( joint_attention_dim: int = 4096, pooled_projection_dim: int = 768, guidance_embeds: bool = False, - axes_dims_rope: List[int] = [16, 56, 56], + axes_dims_rope: list[int] = [16, 56, 56], num_mode: int = None, conditioning_embedding_channels: int = None, ): @@ -150,6 +154,7 @@ def from_transformer( return controlnet + @apply_lora_scale("joint_attention_kwargs") def forward( self, hidden_states: torch.Tensor, @@ -162,9 +167,9 @@ def forward( img_ids: torch.Tensor = None, txt_ids: torch.Tensor = None, guidance: torch.Tensor = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, + joint_attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - ) -> Union[torch.FloatTensor, Transformer2DModelOutput]: + ) -> torch.FloatTensor | Transformer2DModelOutput: """ The [`FluxTransformer2DModel`] forward method. @@ -197,20 +202,6 @@ def forward( If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a `tuple` where the first element is the sample tensor. """ - if joint_attention_kwargs is not None: - joint_attention_kwargs = joint_attention_kwargs.copy() - lora_scale = joint_attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective." - ) hidden_states = self.x_embedder(hidden_states) if self.input_hint_block is not None: @@ -323,10 +314,6 @@ def forward( None if len(controlnet_single_block_samples) == 0 else controlnet_single_block_samples ) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (controlnet_block_samples, controlnet_single_block_samples) @@ -344,7 +331,7 @@ class FluxMultiControlNetModel(ModelMixin): compatible with `FluxControlNetModel`. Args: - controlnets (`List[FluxControlNetModel]`): + controlnets (`list[FluxControlNetModel]`): Provides additional conditioning to the unet during the denoising process. You must set multiple `FluxControlNetModel` as a list. """ @@ -356,18 +343,18 @@ def __init__(self, controlnets): def forward( self, hidden_states: torch.FloatTensor, - controlnet_cond: List[torch.tensor], - controlnet_mode: List[torch.tensor], - conditioning_scale: List[float], + controlnet_cond: list[torch.tensor], + controlnet_mode: list[torch.tensor], + conditioning_scale: list[float], encoder_hidden_states: torch.Tensor = None, pooled_projections: torch.Tensor = None, timestep: torch.LongTensor = None, img_ids: torch.Tensor = None, txt_ids: torch.Tensor = None, guidance: torch.Tensor = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, + joint_attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - ) -> Union[FluxControlNetOutput, Tuple]: + ) -> FluxControlNetOutput | tuple: # ControlNet-Union with multiple conditions # only load one ControlNet for saving memories if len(self.nets) == 1: diff --git a/src/diffusers/models/controlnets/controlnet_hunyuan.py b/src/diffusers/models/controlnets/controlnet_hunyuan.py index d17d5692aa40..6ef92d78dd6e 100644 --- a/src/diffusers/models/controlnets/controlnet_hunyuan.py +++ b/src/diffusers/models/controlnets/controlnet_hunyuan.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. from dataclasses import dataclass -from typing import Dict, Optional, Union import torch from torch import nn @@ -27,7 +26,7 @@ ) from ..modeling_utils import ModelMixin from ..transformers.hunyuan_transformer_2d import HunyuanDiTBlock -from .controlnet import Tuple, zero_module +from .controlnet import zero_module logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -35,7 +34,7 @@ @dataclass class HunyuanControlNetOutput(BaseOutput): - controlnet_block_samples: Tuple[torch.Tensor] + controlnet_block_samples: tuple[torch.Tensor] class HunyuanDiT2DControlNetModel(ModelMixin, ConfigMixin): @@ -45,8 +44,8 @@ def __init__( conditioning_channels: int = 3, num_attention_heads: int = 16, attention_head_dim: int = 88, - in_channels: Optional[int] = None, - patch_size: Optional[int] = None, + in_channels: int | None = None, + patch_size: int | None = None, activation_fn: str = "gelu-approximate", sample_size=32, hidden_size=1152, @@ -116,7 +115,7 @@ def __init__( self.controlnet_blocks.append(controlnet_block) @property - def attn_processors(self) -> Dict[str, AttentionProcessor]: + def attn_processors(self) -> dict[str, AttentionProcessor]: r""" Returns: `dict` of attention processors: A dictionary containing all attention processors used in the model with @@ -125,7 +124,7 @@ def attn_processors(self) -> Dict[str, AttentionProcessor]: # set recursively processors = {} - def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]): + def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: dict[str, AttentionProcessor]): if hasattr(module, "get_processor"): processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True) @@ -139,7 +138,7 @@ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: return processors - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): + def set_attn_processor(self, processor: AttentionProcessor | dict[str, AttentionProcessor]): r""" Sets the attention processor to use to compute attention. @@ -317,7 +316,7 @@ class HunyuanDiT2DMultiControlNetModel(ModelMixin): designed to be compatible with `HunyuanDiT2DControlNetModel`. Args: - controlnets (`List[HunyuanDiT2DControlNetModel]`): + controlnets (`list[HunyuanDiT2DControlNetModel]`): Provides additional conditioning to the unet during the denoising process. You must set multiple `HunyuanDiT2DControlNetModel` as a list. """ diff --git a/src/diffusers/models/controlnets/controlnet_qwenimage.py b/src/diffusers/models/controlnets/controlnet_qwenimage.py index 86971271788f..cfe7c159ad89 100644 --- a/src/diffusers/models/controlnets/controlnet_qwenimage.py +++ b/src/diffusers/models/controlnets/controlnet_qwenimage.py @@ -13,14 +13,19 @@ # limitations under the License. from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, BaseOutput, logging, scale_lora_layers, unscale_lora_layers +from ...utils import ( + BaseOutput, + apply_lora_scale, + deprecate, + logging, +) from ..attention import AttentionMixin from ..cache_utils import CacheMixin from ..controlnets.controlnet import zero_module @@ -31,6 +36,7 @@ QwenImageTransformerBlock, QwenTimestepProjEmbeddings, RMSNorm, + compute_text_seq_len_from_mask, ) @@ -39,7 +45,7 @@ @dataclass class QwenImageControlNetOutput(BaseOutput): - controlnet_block_samples: Tuple[torch.Tensor] + controlnet_block_samples: tuple[torch.Tensor] class QwenImageControlNetModel( @@ -52,12 +58,12 @@ def __init__( self, patch_size: int = 2, in_channels: int = 64, - out_channels: Optional[int] = 16, + out_channels: int | None = 16, num_layers: int = 60, attention_head_dim: int = 128, num_attention_heads: int = 24, joint_attention_dim: int = 3584, - axes_dims_rope: Tuple[int, int, int] = (16, 56, 56), + axes_dims_rope: tuple[int, int, int] = (16, 56, 56), extra_condition_channels: int = 0, # for controlnet-inpainting ): super().__init__() @@ -122,6 +128,7 @@ def from_transformer( return controlnet + @apply_lora_scale("joint_attention_kwargs") def forward( self, hidden_states: torch.Tensor, @@ -130,13 +137,13 @@ def forward( encoder_hidden_states: torch.Tensor = None, encoder_hidden_states_mask: torch.Tensor = None, timestep: torch.LongTensor = None, - img_shapes: Optional[List[Tuple[int, int, int]]] = None, - txt_seq_lens: Optional[List[int]] = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, + img_shapes: list[tuple[int, int, int]] | None = None, + txt_seq_lens: list[int] | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - ) -> Union[torch.FloatTensor, Transformer2DModelOutput]: + ) -> torch.FloatTensor | Transformer2DModelOutput: """ - The [`FluxTransformer2DModel`] forward method. + The [`QwenImageControlNetModel`] forward method. Args: hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`): @@ -147,38 +154,39 @@ def forward( The scale factor for ControlNet outputs. encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`): Conditional embeddings (embeddings computed from the input conditions such as prompts) to use. - pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected - from the embeddings of input conditions. + encoder_hidden_states_mask (`torch.Tensor` of shape `(batch_size, text_sequence_length)`, *optional*): + Mask for the encoder hidden states. Expected to have 1.0 for valid tokens and 0.0 for padding tokens. + Used in the attention processor to prevent attending to padding tokens. The mask can have any pattern + (not just contiguous valid tokens followed by padding) since it's applied element-wise in attention. timestep ( `torch.LongTensor`): Used to indicate denoising step. - block_controlnet_hidden_states: (`list` of `torch.Tensor`): - A list of tensors that if specified are added to the residuals of transformer blocks. + img_shapes (`list[tuple[int, int, int]]`, *optional*): + Image shapes for RoPE computation. + txt_seq_lens (`list[int]`, *optional*): + **Deprecated**. Not needed anymore, we use `encoder_hidden_states` instead to infer text sequence + length. joint_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain - tuple. + Whether or not to return a [`~models.controlnet.ControlNetOutput`] instead of a plain tuple. Returns: - If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a - `tuple` where the first element is the sample tensor. + If `return_dict` is True, a [`~models.controlnet.ControlNetOutput`] is returned, otherwise a `tuple` where + the first element is the controlnet block samples. """ - if joint_attention_kwargs is not None: - joint_attention_kwargs = joint_attention_kwargs.copy() - lora_scale = joint_attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 + # Handle deprecated txt_seq_lens parameter + if txt_seq_lens is not None: + deprecate( + "txt_seq_lens", + "0.39.0", + "Passing `txt_seq_lens` to `QwenImageControlNetModel.forward()` is deprecated and will be removed in " + "version 0.39.0. The text sequence length is now automatically inferred from `encoder_hidden_states` " + "and `encoder_hidden_states_mask`.", + standard_warn=False, + ) - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective." - ) hidden_states = self.img_in(hidden_states) # add @@ -186,32 +194,47 @@ def forward( temb = self.time_text_embed(timestep, hidden_states) - image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device) + # Use the encoder_hidden_states sequence length for RoPE computation and normalize mask + text_seq_len, _, encoder_hidden_states_mask = compute_text_seq_len_from_mask( + encoder_hidden_states, encoder_hidden_states_mask + ) + + image_rotary_emb = self.pos_embed(img_shapes, max_txt_seq_len=text_seq_len, device=hidden_states.device) timestep = timestep.to(hidden_states.dtype) encoder_hidden_states = self.txt_norm(encoder_hidden_states) encoder_hidden_states = self.txt_in(encoder_hidden_states) + # Construct joint attention mask once to avoid reconstructing in every block + block_attention_kwargs = joint_attention_kwargs.copy() if joint_attention_kwargs is not None else {} + if encoder_hidden_states_mask is not None: + # Build joint mask: [text_mask, all_ones_for_image] + batch_size, image_seq_len = hidden_states.shape[:2] + image_mask = torch.ones((batch_size, image_seq_len), dtype=torch.bool, device=hidden_states.device) + joint_attention_mask = torch.cat([encoder_hidden_states_mask, image_mask], dim=1) + block_attention_kwargs["attention_mask"] = joint_attention_mask + block_samples = () - for index_block, block in enumerate(self.transformer_blocks): + for block in self.transformer_blocks: if torch.is_grad_enabled() and self.gradient_checkpointing: encoder_hidden_states, hidden_states = self._gradient_checkpointing_func( block, hidden_states, encoder_hidden_states, - encoder_hidden_states_mask, + None, # Don't pass encoder_hidden_states_mask (using attention_mask instead) temb, image_rotary_emb, + block_attention_kwargs, ) else: encoder_hidden_states, hidden_states = block( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, - encoder_hidden_states_mask=encoder_hidden_states_mask, + encoder_hidden_states_mask=None, # Don't pass (using attention_mask instead) temb=temb, image_rotary_emb=image_rotary_emb, - joint_attention_kwargs=joint_attention_kwargs, + joint_attention_kwargs=block_attention_kwargs, ) block_samples = block_samples + (hidden_states,) @@ -225,10 +248,6 @@ def forward( controlnet_block_samples = [sample * conditioning_scale for sample in controlnet_block_samples] controlnet_block_samples = None if len(controlnet_block_samples) == 0 else controlnet_block_samples - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return controlnet_block_samples @@ -245,7 +264,7 @@ class QwenImageMultiControlNetModel(ModelMixin, ConfigMixin, PeftAdapterMixin, F to be compatible with `QwenImageControlNetModel`. Args: - controlnets (`List[QwenImageControlNetModel]`): + controlnets (`list[QwenImageControlNetModel]`): Provides additional conditioning to the unet during the denoising process. You must set multiple `QwenImageControlNetModel` as a list. """ @@ -257,16 +276,25 @@ def __init__(self, controlnets): def forward( self, hidden_states: torch.FloatTensor, - controlnet_cond: List[torch.tensor], - conditioning_scale: List[float], + controlnet_cond: list[torch.tensor], + conditioning_scale: list[float], encoder_hidden_states: torch.Tensor = None, encoder_hidden_states_mask: torch.Tensor = None, timestep: torch.LongTensor = None, - img_shapes: Optional[List[Tuple[int, int, int]]] = None, - txt_seq_lens: Optional[List[int]] = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, + img_shapes: list[tuple[int, int, int]] | None = None, + txt_seq_lens: list[int] | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - ) -> Union[QwenImageControlNetOutput, Tuple]: + ) -> QwenImageControlNetOutput | tuple: + if txt_seq_lens is not None: + deprecate( + "txt_seq_lens", + "0.39.0", + "Passing `txt_seq_lens` to `QwenImageMultiControlNetModel.forward()` is deprecated and will be " + "removed in version 0.39.0. The text sequence length is now automatically inferred from " + "`encoder_hidden_states` and `encoder_hidden_states_mask`.", + standard_warn=False, + ) # ControlNet-Union with multiple conditions # only load one ControlNet for saving memories if len(self.nets) == 1: @@ -281,7 +309,6 @@ def forward( encoder_hidden_states_mask=encoder_hidden_states_mask, timestep=timestep, img_shapes=img_shapes, - txt_seq_lens=txt_seq_lens, joint_attention_kwargs=joint_attention_kwargs, return_dict=return_dict, ) diff --git a/src/diffusers/models/controlnets/controlnet_sana.py b/src/diffusers/models/controlnets/controlnet_sana.py index c71a8b326635..29e5591fa284 100644 --- a/src/diffusers/models/controlnets/controlnet_sana.py +++ b/src/diffusers/models/controlnets/controlnet_sana.py @@ -13,14 +13,14 @@ # limitations under the License. from dataclasses import dataclass -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any import torch from torch import nn from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, BaseOutput, logging, scale_lora_layers, unscale_lora_layers +from ...utils import BaseOutput, apply_lora_scale, logging from ..attention import AttentionMixin from ..embeddings import PatchEmbed, PixArtAlphaTextProjection from ..modeling_outputs import Transformer2DModelOutput @@ -35,7 +35,7 @@ @dataclass class SanaControlNetOutput(BaseOutput): - controlnet_block_samples: Tuple[torch.Tensor] + controlnet_block_samples: tuple[torch.Tensor] class SanaControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, PeftAdapterMixin): @@ -47,13 +47,13 @@ class SanaControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, PeftAdapterMi def __init__( self, in_channels: int = 32, - out_channels: Optional[int] = 32, + out_channels: int | None = 32, num_attention_heads: int = 70, attention_head_dim: int = 32, num_layers: int = 7, - num_cross_attention_heads: Optional[int] = 20, - cross_attention_head_dim: Optional[int] = 112, - cross_attention_dim: Optional[int] = 2240, + num_cross_attention_heads: int | None = 20, + cross_attention_head_dim: int | None = 112, + cross_attention_dim: int | None = 2240, caption_channels: int = 2304, mlp_ratio: float = 2.5, dropout: float = 0.0, @@ -62,7 +62,7 @@ def __init__( patch_size: int = 1, norm_elementwise_affine: bool = False, norm_eps: float = 1e-6, - interpolation_scale: Optional[int] = None, + interpolation_scale: int | None = None, ) -> None: super().__init__() @@ -117,6 +117,7 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, @@ -124,26 +125,11 @@ def forward( timestep: torch.LongTensor, controlnet_cond: torch.Tensor, conditioning_scale: float = 1.0, - encoder_attention_mask: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_attention_mask: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - ) -> Union[Tuple[torch.Tensor, ...], Transformer2DModelOutput]: - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) - + ) -> tuple[torch.Tensor, ...] | Transformer2DModelOutput: # ensure attention_mask is a bias, and give it a singleton query_tokens dimension. # we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward. # we can tell by counting dims; if ndim == 2: it's a mask rather than a bias. @@ -218,10 +204,6 @@ def forward( block_res_sample = controlnet_block(block_res_sample) controlnet_block_res_samples = controlnet_block_res_samples + (block_res_sample,) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - controlnet_block_res_samples = [sample * conditioning_scale for sample in controlnet_block_res_samples] if not return_dict: diff --git a/src/diffusers/models/controlnets/controlnet_sd3.py b/src/diffusers/models/controlnets/controlnet_sd3.py index 08b86ff344eb..b8cb97adb41a 100644 --- a/src/diffusers/models/controlnets/controlnet_sd3.py +++ b/src/diffusers/models/controlnets/controlnet_sd3.py @@ -14,14 +14,14 @@ from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, logging from ..attention import AttentionMixin, JointTransformerBlock from ..attention_processor import Attention, FusedJointAttnProcessor2_0 from ..embeddings import CombinedTimestepTextProjEmbeddings, PatchEmbed @@ -36,7 +36,7 @@ @dataclass class SD3ControlNetOutput(BaseOutput): - controlnet_block_samples: Tuple[torch.Tensor] + controlnet_block_samples: tuple[torch.Tensor] class SD3ControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin): @@ -69,7 +69,7 @@ class SD3ControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, PeftAdapterMix The maximum latent height/width of positional embeddings. extra_conditioning_channels (`int`, defaults to `0`): The number of extra channels to use for conditioning for patch embedding. - dual_attention_layers (`Tuple[int, ...]`, defaults to `()`): + dual_attention_layers (`tuple[int, ...]`, defaults to `()`): The number of dual-stream transformer blocks to use. qk_norm (`str`, *optional*, defaults to `None`): The normalization to use for query and key in the attention layer. If `None`, no normalization is used. @@ -99,9 +99,9 @@ def __init__( out_channels: int = 16, pos_embed_max_size: int = 96, extra_conditioning_channels: int = 0, - dual_attention_layers: Tuple[int, ...] = (), - qk_norm: Optional[str] = None, - pos_embed_type: Optional[str] = "sincos", + dual_attention_layers: tuple[int, ...] = (), + qk_norm: str | None = None, + pos_embed_type: str | None = "sincos", use_pos_embed: bool = True, force_zeros_for_pooled_projection: bool = True, ): @@ -175,7 +175,7 @@ def __init__( self.gradient_checkpointing = False # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking - def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None: + def enable_forward_chunking(self, chunk_size: int | None = None, dim: int = 0) -> None: """ Sets the attention processor to use [feed forward chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers). @@ -269,6 +269,7 @@ def from_transformer( return controlnet + @apply_lora_scale("joint_attention_kwargs") def forward( self, hidden_states: torch.Tensor, @@ -277,9 +278,9 @@ def forward( encoder_hidden_states: torch.Tensor = None, pooled_projections: torch.Tensor = None, timestep: torch.LongTensor = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, + joint_attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - ) -> Union[torch.Tensor, Transformer2DModelOutput]: + ) -> torch.Tensor | Transformer2DModelOutput: """ The [`SD3Transformer2DModel`] forward method. @@ -308,21 +309,6 @@ def forward( If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a `tuple` where the first element is the sample tensor. """ - if joint_attention_kwargs is not None: - joint_attention_kwargs = joint_attention_kwargs.copy() - lora_scale = joint_attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective." - ) - if self.pos_embed is not None and hidden_states.ndim != 4: raise ValueError("hidden_states must be 4D when pos_embed is used") @@ -382,10 +368,6 @@ def forward( # 6. scaling controlnet_block_res_samples = [sample * conditioning_scale for sample in controlnet_block_res_samples] - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (controlnet_block_res_samples,) @@ -400,7 +382,7 @@ class SD3MultiControlNetModel(ModelMixin): compatible with `SD3ControlNetModel`. Args: - controlnets (`List[SD3ControlNetModel]`): + controlnets (`list[SD3ControlNetModel]`): Provides additional conditioning to the unet during the denoising process. You must set multiple `SD3ControlNetModel` as a list. """ @@ -412,14 +394,14 @@ def __init__(self, controlnets): def forward( self, hidden_states: torch.Tensor, - controlnet_cond: List[torch.tensor], - conditioning_scale: List[float], + controlnet_cond: list[torch.tensor], + conditioning_scale: list[float], pooled_projections: torch.Tensor, encoder_hidden_states: torch.Tensor = None, timestep: torch.LongTensor = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, + joint_attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - ) -> Union[SD3ControlNetOutput, Tuple]: + ) -> SD3ControlNetOutput | tuple: for i, (image, scale, controlnet) in enumerate(zip(controlnet_cond, conditioning_scale, self.nets)): block_samples = controlnet( hidden_states=hidden_states, diff --git a/src/diffusers/models/controlnets/controlnet_sparsectrl.py b/src/diffusers/models/controlnets/controlnet_sparsectrl.py index 8e7faf2d44b0..7da627fe6dd4 100644 --- a/src/diffusers/models/controlnets/controlnet_sparsectrl.py +++ b/src/diffusers/models/controlnets/controlnet_sparsectrl.py @@ -13,7 +13,7 @@ # limitations under the License. from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch from torch import nn @@ -55,7 +55,7 @@ class SparseControlNetOutput(BaseOutput): Output can be used to condition the original UNet's middle block activation. """ - down_block_res_samples: Tuple[torch.Tensor] + down_block_res_samples: tuple[torch.Tensor] mid_block_res_sample: torch.Tensor @@ -64,7 +64,7 @@ def __init__( self, conditioning_embedding_channels: int, conditioning_channels: int = 3, - block_out_channels: Tuple[int, ...] = (16, 32, 96, 256), + block_out_channels: tuple[int, ...] = (16, 32, 96, 256), ): super().__init__() @@ -110,7 +110,7 @@ class SparseControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, FromOrigina The frequency shift to apply to the time embedding. down_block_types (`tuple[str]`, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`): The tuple of downsample blocks to use. - only_cross_attention (`Union[bool, Tuple[bool]]`, defaults to `False`): + only_cross_attention (`bool | tuple[bool]`, defaults to `False`): block_out_channels (`tuple[int]`, defaults to `(320, 640, 1280, 1280)`): The tuple of output channels for each block. layers_per_block (`int`, defaults to 2): @@ -128,28 +128,28 @@ class SparseControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, FromOrigina The epsilon to use for the normalization. cross_attention_dim (`int`, defaults to 1280): The dimension of the cross attention features. - transformer_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 1): + transformer_layers_per_block (`int` or `tuple[int]`, *optional*, defaults to 1): The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`], [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`]. - transformer_layers_per_mid_block (`int` or `Tuple[int]`, *optional*, defaults to 1): + transformer_layers_per_mid_block (`int` or `tuple[int]`, *optional*, defaults to 1): The number of transformer layers to use in each layer in the middle block. - attention_head_dim (`int` or `Tuple[int]`, defaults to 8): + attention_head_dim (`int` or `tuple[int]`, defaults to 8): The dimension of the attention heads. - num_attention_heads (`int` or `Tuple[int]`, *optional*): + num_attention_heads (`int` or `tuple[int]`, *optional*): The number of heads to use for multi-head attention. use_linear_projection (`bool`, defaults to `False`): upcast_attention (`bool`, defaults to `False`): resnet_time_scale_shift (`str`, defaults to `"default"`): Time scale shift config for ResNet blocks (see `ResnetBlock2D`). Choose from `default` or `scale_shift`. - conditioning_embedding_out_channels (`Tuple[int]`, defaults to `(16, 32, 96, 256)`): + conditioning_embedding_out_channels (`tuple[int]`, defaults to `(16, 32, 96, 256)`): The tuple of output channel for each block in the `conditioning_embedding` layer. global_pool_conditions (`bool`, defaults to `False`): TODO(Patrick) - unused parameter controlnet_conditioning_channel_order (`str`, defaults to `rgb`): motion_max_seq_length (`int`, defaults to `32`): The maximum sequence length to use in the motion module. - motion_num_attention_heads (`int` or `Tuple[int]`, defaults to `8`): + motion_num_attention_heads (`int` or `tuple[int]`, defaults to `8`): The number of heads to use in each attention layer of the motion module. concat_conditioning_mask (`bool`, defaults to `True`): use_simplified_condition_embedding (`bool`, defaults to `True`): @@ -164,30 +164,30 @@ def __init__( conditioning_channels: int = 4, flip_sin_to_cos: bool = True, freq_shift: int = 0, - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str, ...] = ( "CrossAttnDownBlockMotion", "CrossAttnDownBlockMotion", "CrossAttnDownBlockMotion", "DownBlockMotion", ), - only_cross_attention: Union[bool, Tuple[bool]] = False, - block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280), + only_cross_attention: bool | tuple[bool] = False, + block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280), layers_per_block: int = 2, downsample_padding: int = 1, mid_block_scale_factor: float = 1, act_fn: str = "silu", - norm_num_groups: Optional[int] = 32, + norm_num_groups: int | None = 32, norm_eps: float = 1e-5, cross_attention_dim: int = 768, - transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1, - transformer_layers_per_mid_block: Optional[Union[int, Tuple[int]]] = None, - temporal_transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1, - attention_head_dim: Union[int, Tuple[int, ...]] = 8, - num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None, + transformer_layers_per_block: int | tuple[int, ...] = 1, + transformer_layers_per_mid_block: int | tuple[int] | None = None, + temporal_transformer_layers_per_block: int | tuple[int, ...] = 1, + attention_head_dim: int | tuple[int, ...] = 8, + num_attention_heads: int | tuple[int, ...] | None = None, use_linear_projection: bool = False, upcast_attention: bool = False, resnet_time_scale_shift: str = "default", - conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256), + conditioning_embedding_out_channels: tuple[int, ...] | None = (16, 32, 96, 256), global_pool_conditions: bool = False, controlnet_conditioning_channel_order: str = "rgb", motion_max_seq_length: int = 32, @@ -389,7 +389,7 @@ def from_unet( cls, unet: UNet2DConditionModel, controlnet_conditioning_channel_order: str = "rgb", - conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256), + conditioning_embedding_out_channels: tuple[int, ...] | None = (16, 32, 96, 256), load_weights_from_unet: bool = True, conditioning_channels: int = 3, ) -> "SparseControlNetModel": @@ -465,7 +465,7 @@ def set_default_attn_processor(self): self.set_attn_processor(processor) # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attention_slice - def set_attention_slice(self, slice_size: Union[str, int, List[int]]) -> None: + def set_attention_slice(self, slice_size: str | int | list[int]) -> None: r""" Enable sliced attention computation. @@ -519,7 +519,7 @@ def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): # Recursively walk through all the children. # Any children which exposes the set_attention_slice method # gets the message - def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]): + def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: list[int]): if hasattr(module, "set_attention_slice"): module.set_attention_slice(slice_size.pop()) @@ -533,24 +533,24 @@ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[i def forward( self, sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], + timestep: torch.Tensor | float | int, encoder_hidden_states: torch.Tensor, controlnet_cond: torch.Tensor, conditioning_scale: float = 1.0, - timestep_cond: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - conditioning_mask: Optional[torch.Tensor] = None, + timestep_cond: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + conditioning_mask: torch.Tensor | None = None, guess_mode: bool = False, return_dict: bool = True, - ) -> Union[SparseControlNetOutput, Tuple[Tuple[torch.Tensor, ...], torch.Tensor]]: + ) -> SparseControlNetOutput | tuple[tuple[torch.Tensor, ...], torch.Tensor]: """ The [`SparseControlNetModel`] forward method. Args: sample (`torch.Tensor`): The noisy input tensor. - timestep (`Union[torch.Tensor, float, int]`): + timestep (`torch.Tensor | float | int`): The number of timesteps to denoise an input. encoder_hidden_states (`torch.Tensor`): The encoder hidden states. diff --git a/src/diffusers/models/controlnets/controlnet_union.py b/src/diffusers/models/controlnets/controlnet_union.py index b4ee6536ca2f..8e434ba1f250 100644 --- a/src/diffusers/models/controlnets/controlnet_union.py +++ b/src/diffusers/models/controlnets/controlnet_union.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch from torch import nn @@ -94,7 +94,7 @@ class ControlNetUnionModel(ModelMixin, AttentionMixin, ConfigMixin, FromOriginal The frequency shift to apply to the time embedding. down_block_types (`tuple[str]`, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`): The tuple of downsample blocks to use. - only_cross_attention (`Union[bool, Tuple[bool]]`, defaults to `False`): + only_cross_attention (`bool | tuple[bool]`, defaults to `False`): block_out_channels (`tuple[int]`, defaults to `(320, 640, 1280, 1280)`): The tuple of output channels for each block. layers_per_block (`int`, defaults to 2): @@ -112,7 +112,7 @@ class ControlNetUnionModel(ModelMixin, AttentionMixin, ConfigMixin, FromOriginal The epsilon to use for the normalization. cross_attention_dim (`int`, defaults to 1280): The dimension of the cross attention features. - transformer_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 1): + transformer_layers_per_block (`int` or `tuple[int]`, *optional*, defaults to 1): The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`], [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`]. @@ -122,7 +122,7 @@ class ControlNetUnionModel(ModelMixin, AttentionMixin, ConfigMixin, FromOriginal encoder_hid_dim_type (`str`, *optional*, defaults to `None`): If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`. - attention_head_dim (`Union[int, Tuple[int]]`, defaults to 8): + attention_head_dim (`int | tuple[int]`, defaults to 8): The dimension of the attention heads. use_linear_projection (`bool`, defaults to `False`): class_embed_type (`str`, *optional*, defaults to `None`): @@ -156,36 +156,36 @@ def __init__( conditioning_channels: int = 3, flip_sin_to_cos: bool = True, freq_shift: int = 0, - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str, ...] = ( "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D", ), - only_cross_attention: Union[bool, Tuple[bool]] = False, - block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280), + only_cross_attention: bool | tuple[bool] = False, + block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280), layers_per_block: int = 2, downsample_padding: int = 1, mid_block_scale_factor: float = 1, act_fn: str = "silu", - norm_num_groups: Optional[int] = 32, + norm_num_groups: int | None = 32, norm_eps: float = 1e-5, cross_attention_dim: int = 1280, - transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1, - encoder_hid_dim: Optional[int] = None, - encoder_hid_dim_type: Optional[str] = None, - attention_head_dim: Union[int, Tuple[int, ...]] = 8, - num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None, + transformer_layers_per_block: int | tuple[int, ...] = 1, + encoder_hid_dim: int | None = None, + encoder_hid_dim_type: str | None = None, + attention_head_dim: int | tuple[int, ...] = 8, + num_attention_heads: int | tuple[int, ...] | None = None, use_linear_projection: bool = False, - class_embed_type: Optional[str] = None, - addition_embed_type: Optional[str] = None, - addition_time_embed_dim: Optional[int] = None, - num_class_embeds: Optional[int] = None, + class_embed_type: str | None = None, + addition_embed_type: str | None = None, + addition_time_embed_dim: int | None = None, + num_class_embeds: int | None = None, upcast_attention: bool = False, resnet_time_scale_shift: str = "default", - projection_class_embeddings_input_dim: Optional[int] = None, + projection_class_embeddings_input_dim: int | None = None, controlnet_conditioning_channel_order: str = "rgb", - conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (48, 96, 192, 384), + conditioning_embedding_out_channels: tuple[int, ...] | None = (48, 96, 192, 384), global_pool_conditions: bool = False, addition_embed_type_num_heads: int = 64, num_control_type: int = 6, @@ -390,7 +390,7 @@ def from_unet( cls, unet: UNet2DConditionModel, controlnet_conditioning_channel_order: str = "rgb", - conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256), + conditioning_embedding_out_channels: tuple[int, ...] | None = (16, 32, 96, 256), load_weights_from_unet: bool = True, ): r""" @@ -472,7 +472,7 @@ def set_default_attn_processor(self): self.set_attn_processor(processor) # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attention_slice - def set_attention_slice(self, slice_size: Union[str, int, List[int]]) -> None: + def set_attention_slice(self, slice_size: str | int | list[int]) -> None: r""" Enable sliced attention computation. @@ -526,7 +526,7 @@ def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): # Recursively walk through all the children. # Any children which exposes the set_attention_slice method # gets the message - def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]): + def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: list[int]): if hasattr(module, "set_attention_slice"): module.set_attention_slice(slice_size.pop()) @@ -540,37 +540,37 @@ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[i def forward( self, sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], + timestep: torch.Tensor | float | int, encoder_hidden_states: torch.Tensor, - controlnet_cond: List[torch.Tensor], + controlnet_cond: list[torch.Tensor], control_type: torch.Tensor, - control_type_idx: List[int], - conditioning_scale: Union[float, List[float]] = 1.0, - class_labels: Optional[torch.Tensor] = None, - timestep_cond: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + control_type_idx: list[int], + conditioning_scale: float | list[float] = 1.0, + class_labels: torch.Tensor | None = None, + timestep_cond: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + added_cond_kwargs: dict[str, torch.Tensor] | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, from_multi: bool = False, guess_mode: bool = False, return_dict: bool = True, - ) -> Union[ControlNetOutput, Tuple[Tuple[torch.Tensor, ...], torch.Tensor]]: + ) -> ControlNetOutput | tuple[tuple[torch.Tensor, ...], torch.Tensor]: """ The [`ControlNetUnionModel`] forward method. Args: sample (`torch.Tensor`): The noisy input tensor. - timestep (`Union[torch.Tensor, float, int]`): + timestep (`torch.Tensor | float | int`): The number of timesteps to denoise an input. encoder_hidden_states (`torch.Tensor`): The encoder hidden states. - controlnet_cond (`List[torch.Tensor]`): + controlnet_cond (`list[torch.Tensor]`): The conditional input tensors. control_type (`torch.Tensor`): A tensor of shape `(batch, num_control_type)` with values `0` or `1` depending on whether the control type is used. - control_type_idx (`List[int]`): + control_type_idx (`list[int]`): The indices of `control_type`. conditioning_scale (`float`, defaults to `1.0`): The scale factor for ControlNet outputs. diff --git a/src/diffusers/models/controlnets/controlnet_xs.py b/src/diffusers/models/controlnets/controlnet_xs.py index 119492b0fac4..6221d4878de9 100644 --- a/src/diffusers/models/controlnets/controlnet_xs.py +++ b/src/diffusers/models/controlnets/controlnet_xs.py @@ -13,7 +13,7 @@ # limitations under the License. from dataclasses import dataclass from math import gcd -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch from torch import Tensor, nn @@ -71,8 +71,8 @@ def __init__( resnets: nn.ModuleList, base_to_ctrl: nn.ModuleList, ctrl_to_base: nn.ModuleList, - attentions: Optional[nn.ModuleList] = None, - downsampler: Optional[nn.Conv2d] = None, + attentions: nn.ModuleList | None = None, + downsampler: nn.Conv2d | None = None, ): super().__init__() self.resnets = resnets @@ -107,14 +107,14 @@ def get_down_block_adapter( ctrl_in_channels: int, ctrl_out_channels: int, temb_channels: int, - max_norm_num_groups: Optional[int] = 32, + max_norm_num_groups: int | None = 32, has_crossattn=True, - transformer_layers_per_block: Optional[Union[int, Tuple[int]]] = 1, - num_attention_heads: Optional[int] = 1, - cross_attention_dim: Optional[int] = 1024, + transformer_layers_per_block: int | tuple[int] | None = 1, + num_attention_heads: int | None = 1, + cross_attention_dim: int | None = 1024, add_downsample: bool = True, - upcast_attention: Optional[bool] = False, - use_linear_projection: Optional[bool] = True, + upcast_attention: bool | None = False, + use_linear_projection: bool | None = True, ): num_layers = 2 # only support sd + sdxl @@ -195,11 +195,11 @@ def get_down_block_adapter( def get_mid_block_adapter( base_channels: int, ctrl_channels: int, - temb_channels: Optional[int] = None, - max_norm_num_groups: Optional[int] = 32, + temb_channels: int | None = None, + max_norm_num_groups: int | None = 32, transformer_layers_per_block: int = 1, - num_attention_heads: Optional[int] = 1, - cross_attention_dim: Optional[int] = 1024, + num_attention_heads: int | None = 1, + cross_attention_dim: int | None = 1024, upcast_attention: bool = False, use_linear_projection: bool = True, ): @@ -230,7 +230,7 @@ def get_mid_block_adapter( def get_up_block_adapter( out_channels: int, prev_output_channel: int, - ctrl_skip_channels: List[int], + ctrl_skip_channels: list[int], ): ctrl_to_base = [] num_layers = 3 # only support sd + sdxl @@ -278,7 +278,7 @@ class ControlNetXSAdapter(ModelMixin, AttentionMixin, ConfigMixin): The tuple of downsample blocks to use. sample_size (`int`, defaults to 96): Height and width of input/output sample. - transformer_layers_per_block (`Union[int, Tuple[int]]`, defaults to 1): + transformer_layers_per_block (`int | tuple[int]`, defaults to 1): The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`]. upcast_attention (`bool`, defaults to `True`): @@ -293,21 +293,21 @@ def __init__( self, conditioning_channels: int = 3, conditioning_channel_order: str = "rgb", - conditioning_embedding_out_channels: Tuple[int, ...] = (16, 32, 96, 256), + conditioning_embedding_out_channels: tuple[int] = (16, 32, 96, 256), time_embedding_mix: float = 1.0, learn_time_embedding: bool = False, - num_attention_heads: Union[int, Tuple[int]] = 4, - block_out_channels: Tuple[int, ...] = (4, 8, 16, 16), - base_block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280), + num_attention_heads: int | tuple[int] = 4, + block_out_channels: tuple[int] = (4, 8, 16, 16), + base_block_out_channels: tuple[int] = (320, 640, 1280, 1280), cross_attention_dim: int = 1024, - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str] = ( "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D", ), - sample_size: Optional[int] = 96, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, + sample_size: int | None = 96, + transformer_layers_per_block: int | tuple[int] = 1, upcast_attention: bool = True, max_norm_num_groups: int = 32, use_linear_projection: bool = True, @@ -429,14 +429,14 @@ def __init__( def from_unet( cls, unet: UNet2DConditionModel, - size_ratio: Optional[float] = None, - block_out_channels: Optional[List[int]] = None, - num_attention_heads: Optional[List[int]] = None, + size_ratio: float | None = None, + block_out_channels: list[int] | None = None, + num_attention_heads: list[int] | None = None, learn_time_embedding: bool = False, time_embedding_mix: int = 1.0, conditioning_channels: int = 3, conditioning_channel_order: str = "rgb", - conditioning_embedding_out_channels: Tuple[int, ...] = (16, 32, 96, 256), + conditioning_embedding_out_channels: tuple[int] = (16, 32, 96, 256), ): r""" Instantiate a [`ControlNetXSAdapter`] from a [`UNet2DConditionModel`]. @@ -447,9 +447,9 @@ def from_unet( size_ratio (float, *optional*, defaults to `None`): When given, block_out_channels is set to a fraction of the base model's block_out_channels. Either this or `block_out_channels` must be given. - block_out_channels (`List[int]`, *optional*, defaults to `None`): + block_out_channels (`list[int]`, *optional*, defaults to `None`): Down blocks output channels in control model. Either this or `size_ratio` must be given. - num_attention_heads (`List[int]`, *optional*, defaults to `None`): + num_attention_heads (`list[int]`, *optional*, defaults to `None`): The dimension of the attention heads. The naming seems a bit confusing and it is, see https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 for why. learn_time_embedding (`bool`, defaults to `False`): @@ -461,7 +461,7 @@ def from_unet( Number of channels of conditioning input (e.g. an image) conditioning_channel_order (`str`, defaults to `"rgb"`): The channel order of conditional image. Will convert to `rgb` if it's `bgr`. - conditioning_embedding_out_channels (`Tuple[int]`, defaults to `(16, 32, 96, 256)`): + conditioning_embedding_out_channels (`tuple[int]`, defaults to `(16, 32, 96, 256)`): The tuple of output channel for each block in the `controlnet_cond_embedding` layer. """ @@ -528,38 +528,33 @@ class UNetControlNetXSModel(ModelMixin, AttentionMixin, ConfigMixin): def __init__( self, # unet configs - sample_size: Optional[int] = 96, - down_block_types: Tuple[str, ...] = ( + sample_size: int | None = 96, + down_block_types: tuple[str] = ( "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D", ), - up_block_types: Tuple[str, ...] = ( - "UpBlock2D", - "CrossAttnUpBlock2D", - "CrossAttnUpBlock2D", - "CrossAttnUpBlock2D", - ), - block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280), - norm_num_groups: Optional[int] = 32, - cross_attention_dim: Union[int, Tuple[int]] = 1024, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, - num_attention_heads: Union[int, Tuple[int]] = 8, - addition_embed_type: Optional[str] = None, - addition_time_embed_dim: Optional[int] = None, + up_block_types: tuple[str] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"), + block_out_channels: tuple[int] = (320, 640, 1280, 1280), + norm_num_groups: int | None = 32, + cross_attention_dim: int | tuple[int] = 1024, + transformer_layers_per_block: int | tuple[int] = 1, + num_attention_heads: int | tuple[int] = 8, + addition_embed_type: str | None = None, + addition_time_embed_dim: int | None = None, upcast_attention: bool = True, use_linear_projection: bool = True, - time_cond_proj_dim: Optional[int] = None, - projection_class_embeddings_input_dim: Optional[int] = None, + time_cond_proj_dim: int | None = None, + projection_class_embeddings_input_dim: int | None = None, # additional controlnet configs time_embedding_mix: float = 1.0, ctrl_conditioning_channels: int = 3, - ctrl_conditioning_embedding_out_channels: Tuple[int, ...] = (16, 32, 96, 256), + ctrl_conditioning_embedding_out_channels: tuple[int] = (16, 32, 96, 256), ctrl_conditioning_channel_order: str = "rgb", ctrl_learn_time_embedding: bool = False, - ctrl_block_out_channels: Tuple[int, ...] = (4, 8, 16, 16), - ctrl_num_attention_heads: Union[int, Tuple[int]] = 4, + ctrl_block_out_channels: tuple[int] = (4, 8, 16, 16), + ctrl_num_attention_heads: int | tuple[int] = 4, ctrl_max_norm_num_groups: int = 32, ): super().__init__() @@ -724,11 +719,11 @@ def __init__( def from_unet( cls, unet: UNet2DConditionModel, - controlnet: Optional[ControlNetXSAdapter] = None, - size_ratio: Optional[float] = None, - ctrl_block_out_channels: Optional[List[float]] = None, - time_embedding_mix: Optional[float] = None, - ctrl_optional_kwargs: Optional[Dict] = None, + controlnet: ControlNetXSAdapter | None = None, + size_ratio: float | None = None, + ctrl_block_out_channels: list[float] | None = None, + time_embedding_mix: float | None = None, + ctrl_optional_kwargs: dict | None = None, ): r""" Instantiate a [`UNetControlNetXSModel`] from a [`UNet2DConditionModel`] and an optional [`ControlNetXSAdapter`] @@ -742,7 +737,7 @@ def from_unet( adapter will be created. size_ratio (float, *optional*, defaults to `None`): Used to construct the controlnet if none is given. See [`ControlNetXSAdapter.from_unet`] for details. - ctrl_block_out_channels (`List[int]`, *optional*, defaults to `None`): + ctrl_block_out_channels (`list[int]`, *optional*, defaults to `None`): Used to construct the controlnet if none is given. See [`ControlNetXSAdapter.from_unet`] for details, where this parameter is called `block_out_channels`. time_embedding_mix (`float`, *optional*, defaults to None): @@ -953,25 +948,25 @@ def unfuse_qkv_projections(self): def forward( self, sample: Tensor, - timestep: Union[torch.Tensor, float, int], + timestep: torch.Tensor | float | int, encoder_hidden_states: torch.Tensor, - controlnet_cond: Optional[torch.Tensor] = None, - conditioning_scale: Optional[float] = 1.0, - class_labels: Optional[torch.Tensor] = None, - timestep_cond: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, + controlnet_cond: torch.Tensor | None = None, + conditioning_scale: float | None = 1.0, + class_labels: torch.Tensor | None = None, + timestep_cond: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + added_cond_kwargs: dict[str, torch.Tensor] | None = None, return_dict: bool = True, apply_control: bool = True, - ) -> Union[ControlNetXSOutput, Tuple]: + ) -> ControlNetXSOutput | tuple: """ The [`ControlNetXSModel`] forward method. Args: sample (`Tensor`): The noisy input tensor. - timestep (`Union[torch.Tensor, float, int]`): + timestep (`torch.Tensor | float | int`): The number of timesteps to denoise an input. encoder_hidden_states (`torch.Tensor`): The encoder hidden states. @@ -1166,13 +1161,13 @@ def __init__( norm_num_groups: int = 32, ctrl_max_norm_num_groups: int = 32, has_crossattn=True, - transformer_layers_per_block: Optional[Union[int, Tuple[int]]] = 1, - base_num_attention_heads: Optional[int] = 1, - ctrl_num_attention_heads: Optional[int] = 1, - cross_attention_dim: Optional[int] = 1024, + transformer_layers_per_block: int | tuple[int] | None = 1, + base_num_attention_heads: int | None = 1, + ctrl_num_attention_heads: int | None = 1, + cross_attention_dim: int | None = 1024, add_downsample: bool = True, - upcast_attention: Optional[bool] = False, - use_linear_projection: Optional[bool] = True, + upcast_attention: bool | None = False, + use_linear_projection: bool | None = True, ): super().__init__() base_resnets = [] @@ -1361,14 +1356,14 @@ def forward( self, hidden_states_base: Tensor, temb: Tensor, - encoder_hidden_states: Optional[Tensor] = None, - hidden_states_ctrl: Optional[Tensor] = None, - conditioning_scale: Optional[float] = 1.0, - attention_mask: Optional[Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - encoder_attention_mask: Optional[Tensor] = None, + encoder_hidden_states: Tensor | None = None, + hidden_states_ctrl: Tensor | None = None, + conditioning_scale: float | None = 1.0, + attention_mask: Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + encoder_attention_mask: Tensor | None = None, apply_control: bool = True, - ) -> Tuple[Tensor, Tensor, Tuple[Tensor, ...], Tuple[Tensor, ...]]: + ) -> tuple[Tensor, Tensor, tuple[Tensor, ...], tuple[Tensor, ...]]: if cross_attention_kwargs is not None: if cross_attention_kwargs.get("scale", None) is not None: logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.") @@ -1455,15 +1450,15 @@ def __init__( self, base_channels: int, ctrl_channels: int, - temb_channels: Optional[int] = None, + temb_channels: int | None = None, norm_num_groups: int = 32, ctrl_max_norm_num_groups: int = 32, transformer_layers_per_block: int = 1, - base_num_attention_heads: Optional[int] = 1, - ctrl_num_attention_heads: Optional[int] = 1, - cross_attention_dim: Optional[int] = 1024, + base_num_attention_heads: int | None = 1, + ctrl_num_attention_heads: int | None = 1, + cross_attention_dim: int | None = 1024, upcast_attention: bool = False, - use_linear_projection: Optional[bool] = True, + use_linear_projection: bool | None = True, ): super().__init__() @@ -1568,13 +1563,13 @@ def forward( hidden_states_base: Tensor, temb: Tensor, encoder_hidden_states: Tensor, - hidden_states_ctrl: Optional[Tensor] = None, - conditioning_scale: Optional[float] = 1.0, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - attention_mask: Optional[Tensor] = None, - encoder_attention_mask: Optional[Tensor] = None, + hidden_states_ctrl: Tensor | None = None, + conditioning_scale: float | None = 1.0, + cross_attention_kwargs: dict[str, Any] | None = None, + attention_mask: Tensor | None = None, + encoder_attention_mask: Tensor | None = None, apply_control: bool = True, - ) -> Tuple[Tensor, Tensor]: + ) -> tuple[Tensor, Tensor]: if cross_attention_kwargs is not None: if cross_attention_kwargs.get("scale", None) is not None: logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.") @@ -1606,17 +1601,17 @@ def __init__( in_channels: int, out_channels: int, prev_output_channel: int, - ctrl_skip_channels: List[int], + ctrl_skip_channels: list[int], temb_channels: int, norm_num_groups: int = 32, - resolution_idx: Optional[int] = None, + resolution_idx: int | None = None, has_crossattn=True, transformer_layers_per_block: int = 1, num_attention_heads: int = 1, cross_attention_dim: int = 1024, add_upsample: bool = True, upcast_attention: bool = False, - use_linear_projection: Optional[bool] = True, + use_linear_projection: bool | None = True, ): super().__init__() resnets = [] @@ -1751,15 +1746,15 @@ def freeze_base_params(self) -> None: def forward( self, hidden_states: Tensor, - res_hidden_states_tuple_base: Tuple[Tensor, ...], - res_hidden_states_tuple_ctrl: Tuple[Tensor, ...], + res_hidden_states_tuple_base: tuple[Tensor, ...], + res_hidden_states_tuple_ctrl: tuple[Tensor, ...], temb: Tensor, - encoder_hidden_states: Optional[Tensor] = None, - conditioning_scale: Optional[float] = 1.0, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - attention_mask: Optional[Tensor] = None, - upsample_size: Optional[int] = None, - encoder_attention_mask: Optional[Tensor] = None, + encoder_hidden_states: Tensor | None = None, + conditioning_scale: float | None = 1.0, + cross_attention_kwargs: dict[str, Any] | None = None, + attention_mask: Tensor | None = None, + upsample_size: int | None = None, + encoder_attention_mask: Tensor | None = None, apply_control: bool = True, ) -> Tensor: if cross_attention_kwargs is not None: diff --git a/src/diffusers/models/controlnets/controlnet_z_image.py b/src/diffusers/models/controlnets/controlnet_z_image.py index 3f79ec925419..85fa0d365547 100644 --- a/src/diffusers/models/controlnets/controlnet_z_image.py +++ b/src/diffusers/models/controlnets/controlnet_z_image.py @@ -13,7 +13,7 @@ # limitations under the License. import math -from typing import List, Literal, Optional, Tuple +from typing import Literal import torch import torch.nn as nn @@ -94,9 +94,9 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - freqs_cis: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + freqs_cis: torch.Tensor | None = None, ) -> torch.Tensor: query = attn.to_q(hidden_states) key = attn.to_k(hidden_states) @@ -234,10 +234,10 @@ def forward( x: torch.Tensor, attn_mask: torch.Tensor, freqs_cis: torch.Tensor, - adaln_input: Optional[torch.Tensor] = None, - noise_mask: Optional[torch.Tensor] = None, - adaln_noisy: Optional[torch.Tensor] = None, - adaln_clean: Optional[torch.Tensor] = None, + adaln_input: torch.Tensor | None = None, + noise_mask: torch.Tensor | None = None, + adaln_noisy: torch.Tensor | None = None, + adaln_clean: torch.Tensor | None = None, ): if self.modulation: seq_len = x.shape[1] @@ -291,8 +291,8 @@ class RopeEmbedder: def __init__( self, theta: float = 256.0, - axes_dims: List[int] = (16, 56, 56), - axes_lens: List[int] = (64, 128, 128), + axes_dims: list[int] = (16, 56, 56), + axes_lens: list[int] = (64, 128, 128), ): self.theta = theta self.axes_dims = axes_dims @@ -301,7 +301,7 @@ def __init__( self.freqs_cis = None @staticmethod - def precompute_freqs_cis(dim: List[int], end: List[int], theta: float = 256.0): + def precompute_freqs_cis(dim: list[int], end: list[int], theta: float = 256.0): with torch.device("cpu"): freqs_cis = [] for i, (d, e) in enumerate(zip(dim, end)): @@ -389,7 +389,7 @@ def forward( x: torch.Tensor, attn_mask: torch.Tensor, freqs_cis: torch.Tensor, - adaln_input: Optional[torch.Tensor] = None, + adaln_input: torch.Tensor | None = None, ): # Control if self.block_id == 0: @@ -435,10 +435,10 @@ class ZImageControlNetModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOrigi @register_to_config def __init__( self, - control_layers_places: List[int] = None, - control_refiner_layers_places: List[int] = None, + control_layers_places: list[int] = None, + control_refiner_layers_places: list[int] = None, control_in_dim=None, - add_control_noise_refiner: Optional[Literal["control_layers", "control_noise_refiner"]] = None, + add_control_noise_refiner: Literal["control_layers", "control_noise_refiner"] | None = None, all_patch_size=(2,), all_f_patch_size=(1,), dim=3840, @@ -505,15 +505,15 @@ def __init__( ] ) - self.t_scale: Optional[float] = None - self.t_embedder: Optional[TimestepEmbedder] = None - self.all_x_embedder: Optional[nn.ModuleDict] = None - self.cap_embedder: Optional[nn.Sequential] = None - self.rope_embedder: Optional[RopeEmbedder] = None - self.noise_refiner: Optional[nn.ModuleList] = None - self.context_refiner: Optional[nn.ModuleList] = None - self.x_pad_token: Optional[nn.Parameter] = None - self.cap_pad_token: Optional[nn.Parameter] = None + self.t_scale: float | None = None + self.t_embedder: TimestepEmbedder | None = None + self.all_x_embedder: nn.ModuleDict | None = None + self.cap_embedder: nn.Sequential | None = None + self.rope_embedder: RopeEmbedder | None = None + self.noise_refiner: nn.ModuleList | None = None + self.context_refiner: nn.ModuleList | None = None + self.x_pad_token: nn.Parameter | None = None + self.cap_pad_token: nn.Parameter | None = None @classmethod def from_transformer(cls, controlnet, transformer): @@ -551,10 +551,10 @@ def _patchify_image(self, image: torch.Tensor, patch_size: int, f_patch_size: in def _pad_with_ids( self, feat: torch.Tensor, - pos_grid_size: Tuple, - pos_start: Tuple, + pos_grid_size: tuple, + pos_start: tuple, device: torch.device, - noise_mask_val: Optional[int] = None, + noise_mask_val: int | None = None, ): """Pad feature to SEQ_MULTI_OF, create position IDs and pad mask.""" ori_len = len(feat) @@ -587,7 +587,7 @@ def _pad_with_ids( # Copied from diffusers.models.transformers.transformer_z_image.ZImageTransformer2DModel.patchify_and_embed def patchify_and_embed( - self, all_image: List[torch.Tensor], all_cap_feats: List[torch.Tensor], patch_size: int, f_patch_size: int + self, all_image: list[torch.Tensor], all_cap_feats: list[torch.Tensor], patch_size: int, f_patch_size: int ): """Patchify for basic mode: single image per batch item.""" device = all_image[0].device @@ -625,7 +625,7 @@ def patchify_and_embed( def patchify( self, - all_image: List[torch.Tensor], + all_image: list[torch.Tensor], patch_size: int, f_patch_size: int, ): @@ -653,10 +653,10 @@ def patchify( def forward( self, - x: List[torch.Tensor], + x: list[torch.Tensor], t, - cap_feats: List[torch.Tensor], - control_context: List[torch.Tensor], + cap_feats: list[torch.Tensor], + control_context: list[torch.Tensor], conditioning_scale: float = 1.0, patch_size=2, f_patch_size=1, diff --git a/src/diffusers/models/controlnets/multicontrolnet.py b/src/diffusers/models/controlnets/multicontrolnet.py index 87a952294997..705c59c0f925 100644 --- a/src/diffusers/models/controlnets/multicontrolnet.py +++ b/src/diffusers/models/controlnets/multicontrolnet.py @@ -1,5 +1,5 @@ import os -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from torch import nn @@ -20,30 +20,30 @@ class MultiControlNetModel(ModelMixin): compatible with `ControlNetModel`. Args: - controlnets (`List[ControlNetModel]`): + controlnets (`list[ControlNetModel]`): Provides additional conditioning to the unet during the denoising process. You must set multiple `ControlNetModel` as a list. """ - def __init__(self, controlnets: Union[List[ControlNetModel], Tuple[ControlNetModel]]): + def __init__(self, controlnets: list[ControlNetModel] | tuple[ControlNetModel]): super().__init__() self.nets = nn.ModuleList(controlnets) def forward( self, sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], + timestep: torch.Tensor | float | int, encoder_hidden_states: torch.Tensor, - controlnet_cond: List[torch.tensor], - conditioning_scale: List[float], - class_labels: Optional[torch.Tensor] = None, - timestep_cond: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + controlnet_cond: list[torch.tensor], + conditioning_scale: list[float], + class_labels: torch.Tensor | None = None, + timestep_cond: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + added_cond_kwargs: dict[str, torch.Tensor] | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, guess_mode: bool = False, return_dict: bool = True, - ) -> Union[ControlNetOutput, Tuple]: + ) -> ControlNetOutput | tuple: for i, (image, scale, controlnet) in enumerate(zip(controlnet_cond, conditioning_scale, self.nets)): down_samples, mid_sample = controlnet( sample=sample, @@ -74,11 +74,11 @@ def forward( def save_pretrained( self, - save_directory: Union[str, os.PathLike], + save_directory: str | os.PathLike, is_main_process: bool = True, save_function: Callable = None, safe_serialization: bool = True, - variant: Optional[str] = None, + variant: str | None = None, ): """ Save a model and its configuration file to a directory, so that it can be re-loaded using the @@ -111,7 +111,7 @@ def save_pretrained( ) @classmethod - def from_pretrained(cls, pretrained_model_path: Optional[Union[str, os.PathLike]], **kwargs): + def from_pretrained(cls, pretrained_model_path: str | os.PathLike | None, **kwargs): r""" Instantiate a pretrained MultiControlNet model from multiple pre-trained controlnet models. @@ -134,7 +134,7 @@ def from_pretrained(cls, pretrained_model_path: Optional[Union[str, os.PathLike] Override the default `torch.dtype` and load the model under this dtype. output_loading_info(`bool`, *optional*, defaults to `False`): Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. - device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*): + device_map (`str` or `dict[str, int | str | torch.device]`, *optional*): A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer name, once a given module name is inside, every submodule of it will be sent to the same device. diff --git a/src/diffusers/models/controlnets/multicontrolnet_union.py b/src/diffusers/models/controlnets/multicontrolnet_union.py index d5506dc186e3..98552f99623a 100644 --- a/src/diffusers/models/controlnets/multicontrolnet_union.py +++ b/src/diffusers/models/controlnets/multicontrolnet_union.py @@ -1,5 +1,5 @@ import os -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from torch import nn @@ -21,32 +21,32 @@ class MultiControlNetUnionModel(ModelMixin): be compatible with `ControlNetUnionModel`. Args: - controlnets (`List[ControlNetUnionModel]`): + controlnets (`list[ControlNetUnionModel]`): Provides additional conditioning to the unet during the denoising process. You must set multiple `ControlNetUnionModel` as a list. """ - def __init__(self, controlnets: Union[List[ControlNetUnionModel], Tuple[ControlNetUnionModel]]): + def __init__(self, controlnets: list[ControlNetUnionModel] | tuple[ControlNetUnionModel]): super().__init__() self.nets = nn.ModuleList(controlnets) def forward( self, sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], + timestep: torch.Tensor | float | int, encoder_hidden_states: torch.Tensor, - controlnet_cond: List[torch.tensor], - control_type: List[torch.Tensor], - control_type_idx: List[List[int]], - conditioning_scale: List[float], - class_labels: Optional[torch.Tensor] = None, - timestep_cond: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + controlnet_cond: list[torch.tensor], + control_type: list[torch.Tensor], + control_type_idx: list[list[int]], + conditioning_scale: list[float], + class_labels: torch.Tensor | None = None, + timestep_cond: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + added_cond_kwargs: dict[str, torch.Tensor] | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, guess_mode: bool = False, return_dict: bool = True, - ) -> Union[ControlNetOutput, Tuple]: + ) -> ControlNetOutput | tuple: down_block_res_samples, mid_block_res_sample = None, None for i, (image, ctype, ctype_idx, scale, controlnet) in enumerate( zip(controlnet_cond, control_type, control_type_idx, conditioning_scale, self.nets) @@ -86,11 +86,11 @@ def forward( # Copied from diffusers.models.controlnets.multicontrolnet.MultiControlNetModel.save_pretrained with ControlNet->ControlNetUnion def save_pretrained( self, - save_directory: Union[str, os.PathLike], + save_directory: str | os.PathLike, is_main_process: bool = True, save_function: Callable = None, safe_serialization: bool = True, - variant: Optional[str] = None, + variant: str | None = None, ): """ Save a model and its configuration file to a directory, so that it can be re-loaded using the @@ -124,7 +124,7 @@ def save_pretrained( @classmethod # Copied from diffusers.models.controlnets.multicontrolnet.MultiControlNetModel.from_pretrained with ControlNet->ControlNetUnion - def from_pretrained(cls, pretrained_model_path: Optional[Union[str, os.PathLike]], **kwargs): + def from_pretrained(cls, pretrained_model_path: str | os.PathLike | None, **kwargs): r""" Instantiate a pretrained MultiControlNetUnion model from multiple pre-trained controlnet models. @@ -147,7 +147,7 @@ def from_pretrained(cls, pretrained_model_path: Optional[Union[str, os.PathLike] Override the default `torch.dtype` and load the model under this dtype. output_loading_info(`bool`, *optional*, defaults to `False`): Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. - device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*): + device_map (`str` or `dict[str, int | str | torch.device]`, *optional*): A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer name, once a given module name is inside, every submodule of it will be sent to the same device. diff --git a/src/diffusers/models/downsampling.py b/src/diffusers/models/downsampling.py index 505816422b2a..871c0ed7ddf7 100644 --- a/src/diffusers/models/downsampling.py +++ b/src/diffusers/models/downsampling.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Tuple - import torch import torch.nn as nn import torch.nn.functional as F @@ -43,7 +41,7 @@ def __init__( self, channels: int, use_conv: bool = False, - out_channels: Optional[int] = None, + out_channels: int | None = None, padding: int = 1, name: str = "conv", ): @@ -86,7 +84,7 @@ def __init__( self, channels: int, use_conv: bool = False, - out_channels: Optional[int] = None, + out_channels: int | None = None, padding: int = 1, name: str = "conv", kernel_size=3, @@ -165,10 +163,10 @@ class FirDownsample2D(nn.Module): def __init__( self, - channels: Optional[int] = None, - out_channels: Optional[int] = None, + channels: int | None = None, + out_channels: int | None = None, use_conv: bool = False, - fir_kernel: Tuple[int, int, int, int] = (1, 3, 3, 1), + fir_kernel: tuple[int, int, int, int] = (1, 3, 3, 1), ): super().__init__() out_channels = out_channels if out_channels else channels @@ -181,8 +179,8 @@ def __init__( def _downsample_2d( self, hidden_states: torch.Tensor, - weight: Optional[torch.Tensor] = None, - kernel: Optional[torch.Tensor] = None, + weight: torch.Tensor | None = None, + kernel: torch.Tensor | None = None, factor: int = 2, gain: float = 1, ) -> torch.Tensor: @@ -355,7 +353,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def downsample_2d( hidden_states: torch.Tensor, - kernel: Optional[torch.Tensor] = None, + kernel: torch.Tensor | None = None, factor: int = 2, gain: float = 1, ) -> torch.Tensor: diff --git a/src/diffusers/models/embeddings.py b/src/diffusers/models/embeddings.py index 37fc412adcc3..bcd192c1f166 100644 --- a/src/diffusers/models/embeddings.py +++ b/src/diffusers/models/embeddings.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. import math -from typing import List, Optional, Tuple, Union import numpy as np import torch @@ -80,11 +79,11 @@ def get_timestep_embedding( def get_3d_sincos_pos_embed( embed_dim: int, - spatial_size: Union[int, Tuple[int, int]], + spatial_size: int | tuple[int, int], temporal_size: int, spatial_interpolation_scale: float = 1.0, temporal_interpolation_scale: float = 1.0, - device: Optional[torch.device] = None, + device: torch.device | None = None, output_type: str = "np", ) -> torch.Tensor: r""" @@ -93,7 +92,7 @@ def get_3d_sincos_pos_embed( Args: embed_dim (`int`): The embedding dimension of inputs. It must be divisible by 16. - spatial_size (`int` or `Tuple[int, int]`): + spatial_size (`int` or `tuple[int, int]`): The spatial dimension of positional embeddings. If an integer is provided, the same size is applied to both spatial dimensions (height and width). temporal_size (`int`): @@ -154,7 +153,7 @@ def get_3d_sincos_pos_embed( def _get_3d_sincos_pos_embed_np( embed_dim: int, - spatial_size: Union[int, Tuple[int, int]], + spatial_size: int | tuple[int, int], temporal_size: int, spatial_interpolation_scale: float = 1.0, temporal_interpolation_scale: float = 1.0, @@ -165,7 +164,7 @@ def _get_3d_sincos_pos_embed_np( Args: embed_dim (`int`): The embedding dimension of inputs. It must be divisible by 16. - spatial_size (`int` or `Tuple[int, int]`): + spatial_size (`int` or `tuple[int, int]`): The spatial dimension of positional embeddings. If an integer is provided, the same size is applied to both spatial dimensions (height and width). temporal_size (`int`): @@ -225,7 +224,7 @@ def get_2d_sincos_pos_embed( extra_tokens=0, interpolation_scale=1.0, base_size=16, - device: Optional[torch.device] = None, + device: torch.device | None = None, output_type: str = "np", ): """ @@ -609,10 +608,10 @@ def forward(self, x, freqs_cis): Patchifies and embeds the input tensor(s). Args: - x (List[torch.Tensor] | torch.Tensor): The input tensor(s) to be patchified and embedded. + x (list[torch.Tensor] | torch.Tensor): The input tensor(s) to be patchified and embedded. Returns: - Tuple[torch.Tensor, torch.Tensor, List[Tuple[int, int]], torch.Tensor]: A tuple containing the patchified + tuple[torch.Tensor, torch.Tensor, list[tuple[int, int]], torch.Tensor]: A tuple containing the patchified and embedded tensor(s), the mask indicating the valid patches, the original image size(s), and the frequency tensor(s). """ @@ -642,7 +641,7 @@ class CogVideoXPatchEmbed(nn.Module): def __init__( self, patch_size: int = 2, - patch_size_t: Optional[int] = None, + patch_size_t: int | None = None, in_channels: int = 16, embed_dim: int = 1920, text_embed_dim: int = 4096, @@ -689,7 +688,7 @@ def __init__( self.register_buffer("pos_embedding", pos_embedding, persistent=persistent) def _get_positional_embeddings( - self, sample_height: int, sample_width: int, sample_frames: int, device: Optional[torch.device] = None + self, sample_height: int, sample_width: int, sample_frames: int, device: torch.device | None = None ) -> torch.Tensor: post_patch_height = sample_height // self.patch_size post_patch_width = sample_width // self.patch_size @@ -836,18 +835,18 @@ def get_3d_rotary_pos_embed( theta: int = 10000, use_real: bool = True, grid_type: str = "linspace", - max_size: Optional[Tuple[int, int]] = None, - device: Optional[torch.device] = None, -) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + max_size: tuple[int, int] | None = None, + device: torch.device | None = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """ RoPE for video tokens with 3D structure. Args: embed_dim: (`int`): The embedding dimension size, corresponding to hidden_size_head. - crops_coords (`Tuple[int]`): + crops_coords (`tuple[int]`): The top-left and bottom-right coordinates of the crop. - grid_size (`Tuple[int]`): + grid_size (`tuple[int]`): The grid size of the spatial positional embedding (height, width). temporal_size (`int`): The size of the temporal dimension. @@ -934,10 +933,10 @@ def get_3d_rotary_pos_embed_allegro( crops_coords, grid_size, temporal_size, - interpolation_scale: Tuple[float, float, float] = (1.0, 1.0, 1.0), + interpolation_scale: tuple[float, float, float] = (1.0, 1.0, 1.0), theta: int = 10000, - device: Optional[torch.device] = None, -) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + device: torch.device | None = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: # TODO(aryan): docs start, stop = crops_coords grid_size_h, grid_size_w = grid_size @@ -973,7 +972,7 @@ def get_3d_rotary_pos_embed_allegro( def get_2d_rotary_pos_embed( - embed_dim, crops_coords, grid_size, use_real=True, device: Optional[torch.device] = None, output_type: str = "np" + embed_dim, crops_coords, grid_size, use_real=True, device: torch.device | None = None, output_type: str = "np" ): """ RoPE for image tokens with 2d structure. @@ -981,9 +980,9 @@ def get_2d_rotary_pos_embed( Args: embed_dim: (`int`): The embedding dimension size - crops_coords (`Tuple[int]`) + crops_coords (`tuple[int]`) The top-left and bottom-right coordinates of the crop. - grid_size (`Tuple[int]`): + grid_size (`tuple[int]`): The grid size of the positional embedding. use_real (`bool`): If True, return real part and imaginary part separately. Otherwise, return complex numbers. @@ -1029,9 +1028,9 @@ def _get_2d_rotary_pos_embed_np(embed_dim, crops_coords, grid_size, use_real=Tru Args: embed_dim: (`int`): The embedding dimension size - crops_coords (`Tuple[int]`) + crops_coords (`tuple[int]`) The top-left and bottom-right coordinates of the crop. - grid_size (`Tuple[int]`): + grid_size (`tuple[int]`): The grid size of the positional embedding. use_real (`bool`): If True, return real part and imaginary part separately. Otherwise, return complex numbers. @@ -1119,7 +1118,7 @@ def get_2d_rotary_pos_embed_lumina(embed_dim, len_h, len_w, linear_factor=1.0, n def get_1d_rotary_pos_embed( dim: int, - pos: Union[np.ndarray, int], + pos: np.ndarray | int, theta: float = 10000.0, use_real=False, linear_factor=1.0, @@ -1186,11 +1185,11 @@ def get_1d_rotary_pos_embed( def apply_rotary_emb( x: torch.Tensor, - freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]], + freqs_cis: torch.Tensor | tuple[torch.Tensor], use_real: bool = True, use_real_unbind_dim: int = -1, sequence_dim: int = 2, -) -> Tuple[torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor]: """ Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are @@ -1200,10 +1199,10 @@ def apply_rotary_emb( Args: x (`torch.Tensor`): Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply - freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],) + freqs_cis (`tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],) Returns: - Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings. + tuple[torch.Tensor, torch.Tensor]: tuple of modified query tensor and key tensor with rotary embeddings. """ if use_real: cos, sin = freqs_cis # [S, D] @@ -1266,7 +1265,7 @@ def __init__( time_embed_dim: int, act_fn: str = "silu", out_dim: int = None, - post_act_fn: Optional[str] = None, + post_act_fn: str | None = None, cond_proj_dim=None, sample_proj_bias=True, ): @@ -1816,7 +1815,7 @@ def forward( timestep: torch.LongTensor, encoder_hidden_states: torch.Tensor, encoder_attention_mask: torch.Tensor, - hidden_dtype: Optional[torch.dtype] = None, + hidden_dtype: torch.dtype | None = None, ): time_proj = self.time_proj(timestep) time_emb = self.timestep_embedder(time_proj.to(dtype=hidden_dtype)) @@ -1961,7 +1960,7 @@ def __init__( self, num_attention_heads: int, embed_dim: int, - output_dim: Optional[int] = None, + output_dim: int | None = None, ) -> None: super().__init__() @@ -2543,7 +2542,7 @@ def __init__( self.time_proj = Timesteps(timestep_in_dim, timestep_flip_sin_to_cos, timestep_freq_shift) self.time_embedding = TimestepEmbedding(timestep_in_dim, hidden_dim, act_fn="silu") - def forward(self, x: torch.Tensor, timestep: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + def forward(self, x: torch.Tensor, timestep: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Forward pass. Args: @@ -2552,7 +2551,7 @@ def forward(self, x: torch.Tensor, timestep: torch.Tensor) -> Tuple[torch.Tensor timestep (`torch.Tensor`): Timestep in denoising process. Returns: - `Tuple`[`torch.Tensor`, `torch.Tensor`]: The pair (latents, timestep_emb). + `tuple`[`torch.Tensor`, `torch.Tensor`]: The pair (latents, timestep_emb). """ timestep_emb = self.time_proj(timestep).to(dtype=x.dtype) timestep_emb = self.time_embedding(timestep_emb) @@ -2572,7 +2571,7 @@ def forward(self, x: torch.Tensor, timestep: torch.Tensor) -> Tuple[torch.Tensor class MultiIPAdapterImageProjection(nn.Module): - def __init__(self, IPAdapterImageProjectionLayers: Union[List[nn.Module], Tuple[nn.Module]]): + def __init__(self, IPAdapterImageProjectionLayers: list[nn.Module] | tuple[nn.Module]): super().__init__() self.image_projection_layers = nn.ModuleList(IPAdapterImageProjectionLayers) @@ -2581,7 +2580,7 @@ def num_ip_adapters(self) -> int: """Number of IP-Adapters loaded.""" return len(self.image_projection_layers) - def forward(self, image_embeds: List[torch.Tensor]): + def forward(self, image_embeds: list[torch.Tensor]): projected_image_embeds = [] # currently, we accept `image_embeds` as diff --git a/src/diffusers/models/lora.py b/src/diffusers/models/lora.py index 85d61d6d7cdf..489a0f0abea9 100644 --- a/src/diffusers/models/lora.py +++ b/src/diffusers/models/lora.py @@ -21,8 +21,6 @@ # ----------------------------------------------------------------# ################################################################### -from typing import Optional, Tuple, Union - import torch import torch.nn.functional as F from torch import nn @@ -198,9 +196,9 @@ def __init__( in_features: int, out_features: int, rank: int = 4, - network_alpha: Optional[float] = None, - device: Optional[Union[torch.device, str]] = None, - dtype: Optional[torch.dtype] = None, + network_alpha: float | None = None, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, ): super().__init__() @@ -260,10 +258,10 @@ def __init__( in_features: int, out_features: int, rank: int = 4, - kernel_size: Union[int, Tuple[int, int]] = (1, 1), - stride: Union[int, Tuple[int, int]] = (1, 1), - padding: Union[int, Tuple[int, int], str] = 0, - network_alpha: Optional[float] = None, + kernel_size: int | tuple[int, int] = (1, 1), + stride: int | tuple[int, int] = (1, 1), + padding: int | tuple[int, int] | str = 0, + network_alpha: float | None = None, ): super().__init__() @@ -301,14 +299,14 @@ class LoRACompatibleConv(nn.Conv2d): A convolutional layer that can be used with LoRA. """ - def __init__(self, *args, lora_layer: Optional[LoRAConv2dLayer] = None, **kwargs): + def __init__(self, *args, lora_layer: LoRAConv2dLayer | None = None, **kwargs): deprecation_message = "Use of `LoRACompatibleConv` is deprecated. Please switch to PEFT backend by installing PEFT: `pip install peft`." deprecate("LoRACompatibleConv", "1.0.0", deprecation_message) super().__init__(*args, **kwargs) self.lora_layer = lora_layer - def set_lora_layer(self, lora_layer: Optional[LoRAConv2dLayer]): + def set_lora_layer(self, lora_layer: LoRAConv2dLayer | None): deprecation_message = "Use of `set_lora_layer()` is deprecated. Please switch to PEFT backend by installing PEFT: `pip install peft`." deprecate("set_lora_layer", "1.0.0", deprecation_message) @@ -388,14 +386,14 @@ class LoRACompatibleLinear(nn.Linear): A Linear layer that can be used with LoRA. """ - def __init__(self, *args, lora_layer: Optional[LoRALinearLayer] = None, **kwargs): + def __init__(self, *args, lora_layer: LoRALinearLayer | None = None, **kwargs): deprecation_message = "Use of `LoRACompatibleLinear` is deprecated. Please switch to PEFT backend by installing PEFT: `pip install peft`." deprecate("LoRACompatibleLinear", "1.0.0", deprecation_message) super().__init__(*args, **kwargs) self.lora_layer = lora_layer - def set_lora_layer(self, lora_layer: Optional[LoRALinearLayer]): + def set_lora_layer(self, lora_layer: LoRALinearLayer | None): deprecation_message = "Use of `set_lora_layer()` is deprecated. Please switch to PEFT backend by installing PEFT: `pip install peft`." deprecate("set_lora_layer", "1.0.0", deprecation_message) self.lora_layer = lora_layer diff --git a/src/diffusers/models/model_loading_utils.py b/src/diffusers/models/model_loading_utils.py index 8b48ba6b4873..04642ad5d401 100644 --- a/src/diffusers/models/model_loading_utils.py +++ b/src/diffusers/models/model_loading_utils.py @@ -22,7 +22,6 @@ from collections import OrderedDict, defaultdict from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path -from typing import Dict, List, Optional, Union from zipfile import is_zipfile import safetensors @@ -47,6 +46,7 @@ is_torch_version, logging, ) +from ..utils.distributed_utils import is_torch_dist_rank_zero logger = logging.get_logger(__name__) @@ -135,7 +135,7 @@ def _fetch_remapped_cls_from_config(config, old_class): return old_class -def _determine_param_device(param_name: str, device_map: Optional[Dict[str, Union[int, str, torch.device]]]): +def _determine_param_device(param_name: str, device_map: dict[str, int | str | torch.device] | None): """ Find the device of param_name from the device_map. """ @@ -153,10 +153,10 @@ def _determine_param_device(param_name: str, device_map: Optional[Dict[str, Unio def load_state_dict( - checkpoint_file: Union[str, os.PathLike], - dduf_entries: Optional[Dict[str, DDUFEntry]] = None, + checkpoint_file: str | os.PathLike, + dduf_entries: dict[str, DDUFEntry] | None = None, disable_mmap: bool = False, - map_location: Union[str, torch.device] = "cpu", + map_location: str | torch.device = "cpu", ): """ Reads a checkpoint file, returning properly formatted errors if they arise. @@ -213,17 +213,17 @@ def load_state_dict( def load_model_dict_into_meta( model, state_dict: OrderedDict, - dtype: Optional[Union[str, torch.dtype]] = None, - model_name_or_path: Optional[str] = None, - hf_quantizer: Optional[DiffusersQuantizer] = None, - keep_in_fp32_modules: Optional[List] = None, - device_map: Optional[Dict[str, Union[int, str, torch.device]]] = None, - unexpected_keys: Optional[List[str]] = None, - offload_folder: Optional[Union[str, os.PathLike]] = None, - offload_index: Optional[Dict] = None, - state_dict_index: Optional[Dict] = None, - state_dict_folder: Optional[Union[str, os.PathLike]] = None, -) -> List[str]: + dtype: str | torch.dtype | None = None, + model_name_or_path: str | None = None, + hf_quantizer: DiffusersQuantizer | None = None, + keep_in_fp32_modules: list | None = None, + device_map: dict[str, int | str | torch.device] | None = None, + unexpected_keys: list[str] | None = None, + offload_folder: str | os.PathLike | None = None, + offload_index: dict | None = None, + state_dict_index: dict | None = None, + state_dict_folder: str | os.PathLike | None = None, +) -> list[str]: """ This is somewhat similar to `_load_state_dict_into_model`, but deals with a model that has some or all of its params on a `meta` device. It replaces the model params with the data from the `state_dict` @@ -354,8 +354,9 @@ def _load_shard_file( state_dict_folder=None, ignore_mismatched_sizes=False, low_cpu_mem_usage=False, + disable_mmap=False, ): - state_dict = load_state_dict(shard_file, dduf_entries=dduf_entries) + state_dict = load_state_dict(shard_file, dduf_entries=dduf_entries, disable_mmap=disable_mmap) mismatched_keys = _find_mismatched_keys( state_dict, model_state_dict, @@ -401,6 +402,7 @@ def _load_shard_files_with_threadpool( state_dict_folder=None, ignore_mismatched_sizes=False, low_cpu_mem_usage=False, + disable_mmap=False, ): # Do not spawn anymore workers than you need num_workers = min(len(shard_files), DEFAULT_HF_PARALLEL_LOADING_WORKERS) @@ -427,10 +429,15 @@ def _load_shard_files_with_threadpool( state_dict_folder=state_dict_folder, ignore_mismatched_sizes=ignore_mismatched_sizes, low_cpu_mem_usage=low_cpu_mem_usage, + disable_mmap=disable_mmap, ) + tqdm_kwargs = {"total": len(shard_files), "desc": "Loading checkpoint shards"} + if not is_torch_dist_rank_zero(): + tqdm_kwargs["disable"] = True + with ThreadPoolExecutor(max_workers=num_workers) as executor: - with logging.tqdm(total=len(shard_files), desc="Loading checkpoint shards") as pbar: + with logging.tqdm(**tqdm_kwargs) as pbar: futures = [executor.submit(load_one, shard_file) for shard_file in shard_files] for future in as_completed(futures): result = future.result() @@ -466,7 +473,7 @@ def _find_mismatched_keys( def _load_state_dict_into_model( model_to_load, state_dict: OrderedDict, assign_to_params_buffers: bool = False -) -> List[str]: +) -> list[str]: # Convert old format to new format if needed from a PyTorch state_dict # copy state_dict so _load_from_state_dict can modify it state_dict = state_dict.copy() @@ -505,7 +512,7 @@ def _fetch_index_file( revision, user_agent, commit_hash, - dduf_entries: Optional[Dict[str, DDUFEntry]] = None, + dduf_entries: dict[str, DDUFEntry] | None = None, ): if is_local: index_file = Path( @@ -555,7 +562,7 @@ def _fetch_index_file_legacy( revision, user_agent, commit_hash, - dduf_entries: Optional[Dict[str, DDUFEntry]] = None, + dduf_entries: dict[str, DDUFEntry] | None = None, ): if is_local: index_file = Path( @@ -714,7 +721,7 @@ def _expand_device_map(device_map, param_names): # Adapted from: https://github.com/huggingface/transformers/blob/0687d481e2c71544501ef9cb3eef795a6e79b1de/src/transformers/modeling_utils.py#L5859 def _caching_allocator_warmup( - model, expanded_device_map: Dict[str, torch.device], dtype: torch.dtype, hf_quantizer: Optional[DiffusersQuantizer] + model, expanded_device_map: dict[str, torch.device], dtype: torch.dtype, hf_quantizer: DiffusersQuantizer | None ) -> None: """ This function warm-ups the caching allocator based on the size of the model tensors that will reside on each diff --git a/src/diffusers/models/modeling_flax_utils.py b/src/diffusers/models/modeling_flax_utils.py index 3f060993190f..9f62bd7199e0 100644 --- a/src/diffusers/models/modeling_flax_utils.py +++ b/src/diffusers/models/modeling_flax_utils.py @@ -15,7 +15,7 @@ import os from pickle import UnpicklingError -from typing import Any, Dict, Union +from typing import Any import jax import jax.numpy as jnp @@ -68,7 +68,7 @@ def _from_config(cls, config, **kwargs): """ return cls(config, **kwargs) - def _cast_floating_to(self, params: Union[Dict, FrozenDict], dtype: jnp.dtype, mask: Any = None) -> Any: + def _cast_floating_to(self, params: dict | FrozenDict, dtype: jnp.dtype, mask: Any = None) -> Any: """ Helper method to cast floating-point values of given parameter `PyTree` to given `dtype`. """ @@ -92,7 +92,7 @@ def conditional_cast(param): return unflatten_dict(flat_params) - def to_bf16(self, params: Union[Dict, FrozenDict], mask: Any = None): + def to_bf16(self, params: dict | FrozenDict, mask: Any = None): r""" Cast the floating-point `params` to `jax.numpy.bfloat16`. This returns a new `params` tree and does not cast the `params` in place. @@ -101,9 +101,9 @@ def to_bf16(self, params: Union[Dict, FrozenDict], mask: Any = None): half-precision training or to save weights in bfloat16 for inference in order to save memory and improve speed. Arguments: - params (`Union[Dict, FrozenDict]`): + params (`dict | FrozenDict`): A `PyTree` of model parameters. - mask (`Union[Dict, FrozenDict]`): + mask (`dict | FrozenDict`): A `PyTree` with same structure as the `params` tree. The leaves should be booleans. It should be `True` for params you want to cast, and `False` for those you want to skip. @@ -131,15 +131,15 @@ def to_bf16(self, params: Union[Dict, FrozenDict], mask: Any = None): ```""" return self._cast_floating_to(params, jnp.bfloat16, mask) - def to_fp32(self, params: Union[Dict, FrozenDict], mask: Any = None): + def to_fp32(self, params: dict | FrozenDict, mask: Any = None): r""" Cast the floating-point `params` to `jax.numpy.float32`. This method can be used to explicitly convert the model parameters to fp32 precision. This returns a new `params` tree and does not cast the `params` in place. Arguments: - params (`Union[Dict, FrozenDict]`): + params (`dict | FrozenDict`): A `PyTree` of model parameters. - mask (`Union[Dict, FrozenDict]`): + mask (`dict | FrozenDict`): A `PyTree` with same structure as the `params` tree. The leaves should be booleans. It should be `True` for params you want to cast, and `False` for those you want to skip. @@ -158,7 +158,7 @@ def to_fp32(self, params: Union[Dict, FrozenDict], mask: Any = None): ```""" return self._cast_floating_to(params, jnp.float32, mask) - def to_fp16(self, params: Union[Dict, FrozenDict], mask: Any = None): + def to_fp16(self, params: dict | FrozenDict, mask: Any = None): r""" Cast the floating-point `params` to `jax.numpy.float16`. This returns a new `params` tree and does not cast the `params` in place. @@ -167,9 +167,9 @@ def to_fp16(self, params: Union[Dict, FrozenDict], mask: Any = None): half-precision training or to save weights in float16 for inference in order to save memory and improve speed. Arguments: - params (`Union[Dict, FrozenDict]`): + params (`dict | FrozenDict`): A `PyTree` of model parameters. - mask (`Union[Dict, FrozenDict]`): + mask (`dict | FrozenDict`): A `PyTree` with same structure as the `params` tree. The leaves should be booleans. It should be `True` for params you want to cast, and `False` for those you want to skip. @@ -197,14 +197,14 @@ def to_fp16(self, params: Union[Dict, FrozenDict], mask: Any = None): ```""" return self._cast_floating_to(params, jnp.float16, mask) - def init_weights(self, rng: jax.Array) -> Dict: + def init_weights(self, rng: jax.Array) -> dict: raise NotImplementedError(f"init_weights method has to be implemented for {self}") @classmethod @validate_hf_hub_args def from_pretrained( cls, - pretrained_model_name_or_path: Union[str, os.PathLike], + pretrained_model_name_or_path: str | os.PathLike, dtype: jnp.dtype = jnp.float32, *model_args, **kwargs, @@ -233,14 +233,14 @@ def from_pretrained( model_args (sequence of positional arguments, *optional*): All remaining positional arguments are passed to the underlying model's `__init__` method. - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. local_files_only(`bool`, *optional*, defaults to `False`): @@ -493,8 +493,8 @@ def from_pretrained( def save_pretrained( self, - save_directory: Union[str, os.PathLike], - params: Union[Dict, FrozenDict], + save_directory: str | os.PathLike, + params: dict | FrozenDict, is_main_process: bool = True, push_to_hub: bool = False, **kwargs, @@ -506,7 +506,7 @@ def save_pretrained( Arguments: save_directory (`str` or `os.PathLike`): Directory to save a model and its configuration file to. Will be created if it doesn't exist. - params (`Union[Dict, FrozenDict]`): + params (`dict | FrozenDict`): A `PyTree` of model parameters. is_main_process (`bool`, *optional*, defaults to `True`): Whether the process calling this is the main process or not. Useful during distributed training and you @@ -516,7 +516,7 @@ def save_pretrained( Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the repository you want to push to with `repo_id` (will default to the name of `save_directory` in your namespace). - kwargs (`Dict[str, Any]`, *optional*): + kwargs (`dict[str, Any]`, *optional*): Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. """ if os.path.isfile(save_directory): diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 41da95d3a2a2..0901840679e3 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -27,7 +27,7 @@ from contextlib import ExitStack, contextmanager from functools import wraps from pathlib import Path -from typing import Any, Callable, ContextManager, Dict, List, Optional, Tuple, Type, Union +from typing import Any, Callable, ContextManager, Type import safetensors import torch @@ -59,11 +59,8 @@ is_torch_version, logging, ) -from ..utils.hub_utils import ( - PushToHubMixin, - load_or_create_model_card, - populate_model_card, -) +from ..utils.distributed_utils import is_torch_dist_rank_zero +from ..utils.hub_utils import PushToHubMixin, load_or_create_model_card, populate_model_card from ..utils.torch_utils import empty_device_cache from ._modeling_parallel import ContextParallelConfig, ContextParallelModelPlan, ParallelConfig from .model_loading_utils import ( @@ -84,7 +81,7 @@ class ContextManagers: in the `fastcore` library. """ - def __init__(self, context_managers: List[ContextManager]): + def __init__(self, context_managers: list[ContextManager]): self.context_managers = context_managers self.stack = ExitStack() @@ -146,7 +143,7 @@ def get_parameter_device(parameter: torch.nn.Module) -> torch.device: except StopIteration: # For torch.nn.DataParallel compatibility in PyTorch 1.5 - def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]: + def find_tensor_attributes(module: torch.nn.Module) -> list[tuple[str, Tensor]]: tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)] return tuples @@ -194,7 +191,7 @@ def get_parameter_dtype(parameter: torch.nn.Module) -> torch.dtype: return last_dtype # For nn.DataParallel compatibility in PyTorch > 1.5 - def find_tensor_attributes(module: nn.Module) -> List[Tuple[str, Tensor]]: + def find_tensor_attributes(module: nn.Module) -> list[tuple[str, Tensor]]: tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)] return tuples @@ -283,7 +280,7 @@ def is_gradient_checkpointing(self) -> bool: """ return any(hasattr(m, "gradient_checkpointing") and m.gradient_checkpointing for m in self.modules()) - def enable_gradient_checkpointing(self, gradient_checkpointing_func: Optional[Callable] = None) -> None: + def enable_gradient_checkpointing(self, gradient_checkpointing_func: Callable | None = None) -> None: """ Activates gradient checkpointing for the current model (may be referred to as *activation checkpointing* or *checkpoint activations* in other frameworks). @@ -352,7 +349,7 @@ def disable_npu_flash_attention(self) -> None: self.set_use_npu_flash_attention(False) def set_use_xla_flash_attention( - self, use_xla_flash_attention: bool, partition_spec: Optional[Callable] = None, **kwargs + self, use_xla_flash_attention: bool, partition_spec: Callable | None = None, **kwargs ) -> None: # Recursively walk through all the children. # Any children which exposes the set_use_xla_flash_attention method @@ -368,7 +365,7 @@ def fn_recursive_set_flash_attention(module: torch.nn.Module): if isinstance(module, torch.nn.Module): fn_recursive_set_flash_attention(module) - def enable_xla_flash_attention(self, partition_spec: Optional[Callable] = None, **kwargs): + def enable_xla_flash_attention(self, partition_spec: Callable | None = None, **kwargs): r""" Enable the flash attention pallals kernel for torch_xla. """ @@ -380,9 +377,7 @@ def disable_xla_flash_attention(self): """ self.set_use_xla_flash_attention(False) - def set_use_memory_efficient_attention_xformers( - self, valid: bool, attention_op: Optional[Callable] = None - ) -> None: + def set_use_memory_efficient_attention_xformers(self, valid: bool, attention_op: Callable | None = None) -> None: # Recursively walk through all the children. # Any children which exposes the set_use_memory_efficient_attention_xformers method # gets the message @@ -397,7 +392,7 @@ def fn_recursive_set_mem_eff(module: torch.nn.Module): if isinstance(module, torch.nn.Module): fn_recursive_set_mem_eff(module) - def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None) -> None: + def enable_xformers_memory_efficient_attention(self, attention_op: Callable | None = None) -> None: r""" Enable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/). @@ -438,9 +433,9 @@ def disable_xformers_memory_efficient_attention(self) -> None: def enable_layerwise_casting( self, storage_dtype: torch.dtype = torch.float8_e4m3fn, - compute_dtype: Optional[torch.dtype] = None, - skip_modules_pattern: Optional[Tuple[str, ...]] = None, - skip_modules_classes: Optional[Tuple[Type[torch.nn.Module], ...]] = None, + compute_dtype: torch.dtype | None = None, + skip_modules_pattern: tuple[str, ...] | None = None, + skip_modules_classes: tuple[Type[torch.nn.Module], ...] | None = None, non_blocking: bool = False, ) -> None: r""" @@ -476,11 +471,11 @@ def enable_layerwise_casting( The dtype to which the model should be cast for storage. compute_dtype (`torch.dtype`): The dtype to which the model weights should be cast during the forward pass. - skip_modules_pattern (`Tuple[str, ...]`, *optional*): + skip_modules_pattern (`tuple[str, ...]`, *optional*): A list of patterns to match the names of the modules to skip during the layerwise casting process. If set to `None`, default skip patterns are used to ignore certain internal layers of modules and PEFT layers. - skip_modules_classes (`Tuple[Type[torch.nn.Module], ...]`, *optional*): + skip_modules_classes (`tuple[Type[torch.nn.Module], ...]`, *optional*): A list of module classes to skip during the layerwise casting process. non_blocking (`bool`, *optional*, defaults to `False`): If `True`, the weight casting operations are non-blocking. @@ -525,14 +520,14 @@ def enable_group_offload( onload_device: torch.device, offload_device: torch.device = torch.device("cpu"), offload_type: str = "block_level", - num_blocks_per_group: Optional[int] = None, + num_blocks_per_group: int | None = None, non_blocking: bool = False, use_stream: bool = False, record_stream: bool = False, low_cpu_mem_usage=False, - offload_to_disk_path: Optional[str] = None, - block_modules: Optional[str] = None, - exclude_kwargs: Optional[str] = None, + offload_to_disk_path: str | None = None, + block_modules: str | None = None, + exclude_kwargs: str | None = None, ) -> None: r""" Activates group offloading for the current model. @@ -602,6 +597,7 @@ def set_attention_backend(self, backend: str) -> None: from .attention import AttentionModuleMixin from .attention_dispatch import ( AttentionBackendName, + _AttentionBackendRegistry, _check_attention_backend_requirements, _maybe_download_kernel_for_backend, ) @@ -610,6 +606,16 @@ def set_attention_backend(self, backend: str) -> None: from .attention_processor import Attention, MochiAttention logger.warning("Attention backends are an experimental feature and the API may be subject to change.") + attention_classes = (Attention, MochiAttention, AttentionModuleMixin) + + parallel_config_set = False + for module in self.modules(): + if not isinstance(module, attention_classes): + continue + processor = module.processor + if getattr(processor, "_parallel_config", None) is not None: + parallel_config_set = True + break backend = backend.lower() available_backends = {x.value for x in AttentionBackendName.__members__.values()} @@ -617,10 +623,17 @@ def set_attention_backend(self, backend: str) -> None: raise ValueError(f"`{backend=}` must be one of the following: " + ", ".join(available_backends)) backend = AttentionBackendName(backend) + if parallel_config_set and not _AttentionBackendRegistry._is_context_parallel_available(backend): + compatible_backends = sorted(_AttentionBackendRegistry._supports_context_parallel) + raise ValueError( + f"Context parallelism is enabled but current attention backend '{backend.value}' " + f"does not support context parallelism. " + f"Please set a compatible attention backend: {compatible_backends} using `model.set_attention_backend()`." + ) + _check_attention_backend_requirements(backend) _maybe_download_kernel_for_backend(backend) - attention_classes = (Attention, MochiAttention, AttentionModuleMixin) for module in self.modules(): if not isinstance(module, attention_classes): continue @@ -629,6 +642,9 @@ def set_attention_backend(self, backend: str) -> None: continue processor._attention_backend = backend + # Important to set the active backend so that it propagates gracefully throughout. + _AttentionBackendRegistry.set_active_backend(backend) + def reset_attention_backend(self) -> None: """ Resets the attention backend for the model. Following calls to `forward` will use the environment default, if @@ -650,12 +666,12 @@ def reset_attention_backend(self) -> None: def save_pretrained( self, - save_directory: Union[str, os.PathLike], + save_directory: str | os.PathLike, is_main_process: bool = True, - save_function: Optional[Callable] = None, + save_function: Callable | None = None, safe_serialization: bool = True, - variant: Optional[str] = None, - max_shard_size: Union[int, str] = "10GB", + variant: str | None = None, + max_shard_size: int | str = "10GB", push_to_hub: bool = False, **kwargs, ): @@ -689,7 +705,7 @@ def save_pretrained( Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the repository you want to push to with `repo_id` (will default to the name of `save_directory` in your namespace). - kwargs (`Dict[str, Any]`, *optional*): + kwargs (`dict[str, Any]`, *optional*): Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. """ if os.path.isfile(save_directory): @@ -817,7 +833,7 @@ def dequantize(self): @classmethod @validate_hf_hub_args - def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs) -> Self: + def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None, **kwargs) -> Self: r""" Instantiate a pretrained PyTorch model from a pretrained model configuration. @@ -833,7 +849,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved with [`~ModelMixin.save_pretrained`]. - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. torch_dtype (`torch.dtype`, *optional*): @@ -841,7 +857,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. output_loading_info (`bool`, *optional*, defaults to `False`): @@ -863,7 +879,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P Mirror source to resolve accessibility issues if you're downloading a model in China. We do not guarantee the timeliness or safety of the source, and you should refer to the mirror site for more information. - device_map (`Union[int, str, torch.device]` or `Dict[str, Union[int, str, torch.device]]`, *optional*): + device_map (`int | str | torch.device` or `dict[str, int | str | torch.device]`, *optional*): A map that specifies where each submodule should go. It doesn't need to be defined for each parameter/buffer name; once a given module name is inside, every submodule of it will be sent to the same device. Defaults to `None`, meaning that the model will be loaded on CPU. @@ -965,9 +981,9 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P variant = kwargs.pop("variant", None) use_safetensors = kwargs.pop("use_safetensors", None) quantization_config = kwargs.pop("quantization_config", None) - dduf_entries: Optional[Dict[str, DDUFEntry]] = kwargs.pop("dduf_entries", None) + dduf_entries: dict[str, DDUFEntry] | None = kwargs.pop("dduf_entries", None) disable_mmap = kwargs.pop("disable_mmap", False) - parallel_config: Optional[Union[ParallelConfig, ContextParallelConfig]] = kwargs.pop("parallel_config", None) + parallel_config: ParallelConfig | ContextParallelConfig | None = kwargs.pop("parallel_config", None) is_parallel_loading_enabled = HF_ENABLE_PARALLEL_LOADING if is_parallel_loading_enabled and not low_cpu_mem_usage: @@ -1309,6 +1325,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P keep_in_fp32_modules=keep_in_fp32_modules, dduf_entries=dduf_entries, is_parallel_loading_enabled=is_parallel_loading_enabled, + disable_mmap=disable_mmap, ) loading_info = { "missing_keys": missing_keys, @@ -1363,12 +1380,12 @@ def cuda(self, *args, **kwargs): # Checks if the model has been loaded in 4-bit or 8-bit with BNB if getattr(self, "quantization_method", None) == QuantizationMethod.BITS_AND_BYTES: - if getattr(self, "is_loaded_in_8bit", False): + if getattr(self, "is_loaded_in_8bit", False) and is_bitsandbytes_version("<", "0.48.0"): raise ValueError( - "Calling `cuda()` is not supported for `8-bit` quantized models. " - " Please use the model as it is, since the model has already been set to the correct devices." + "Calling `cuda()` is not supported for `8-bit` quantized models with the installed version of bitsandbytes. " + f"The current device is `{self.device}`. If you intended to move the model, please install bitsandbytes >= 0.48.0." ) - elif is_bitsandbytes_version("<", "0.43.2"): + elif getattr(self, "is_loaded_in_4bit", False) and is_bitsandbytes_version("<", "0.43.2"): raise ValueError( "Calling `cuda()` is not supported for `4-bit` quantized models with the installed version of bitsandbytes. " f"The current device is `{self.device}`. If you intended to move the model, please install bitsandbytes >= 0.43.2." @@ -1415,17 +1432,16 @@ def to(self, *args, **kwargs): ) if getattr(self, "quantization_method", None) == QuantizationMethod.BITS_AND_BYTES: - if getattr(self, "is_loaded_in_8bit", False): + if getattr(self, "is_loaded_in_8bit", False) and is_bitsandbytes_version("<", "0.48.0"): raise ValueError( - "`.to` is not supported for `8-bit` bitsandbytes models. Please use the model as it is, since the" - " model has already been set to the correct devices and casted to the correct `dtype`." + "Calling `to()` is not supported for `8-bit` quantized models with the installed version of bitsandbytes. " + f"The current device is `{self.device}`. If you intended to move the model, please install bitsandbytes >= 0.48.0." ) - elif is_bitsandbytes_version("<", "0.43.2"): + elif getattr(self, "is_loaded_in_4bit", False) and is_bitsandbytes_version("<", "0.43.2"): raise ValueError( "Calling `to()` is not supported for `4-bit` quantized models with the installed version of bitsandbytes. " f"The current device is `{self.device}`. If you intended to move the model, please install bitsandbytes >= 0.43.2." ) - if _is_group_offload_enabled(self) and device_arg_or_kwarg_present: logger.warning( f"The module '{self.__class__.__name__}' is group offloaded and moving it using `.to()` is not supported." @@ -1492,8 +1508,8 @@ def compile_repeated_blocks(self, *args, **kwargs): def enable_parallelism( self, *, - config: Union[ParallelConfig, ContextParallelConfig], - cp_plan: Optional[Dict[str, ContextParallelModelPlan]] = None, + config: ParallelConfig | ContextParallelConfig, + cp_plan: dict[str, ContextParallelModelPlan] | None = None, ): logger.warning( "`enable_parallelism` is an experimental feature. The API may change in the future and breaking changes may be introduced at any time without warning." @@ -1541,7 +1557,7 @@ def enable_parallelism( f"Context parallelism is enabled but the attention processor '{processor.__class__.__name__}' " f"is using backend '{attention_backend.value}' which does not support context parallelism. " f"Please set a compatible attention backend: {compatible_backends} using `model.set_attention_backend()` before " - f"calling `enable_parallelism()`." + f"calling `model.enable_parallelism()`." ) # All modules use the same attention processor and backend. We don't need to @@ -1581,20 +1597,21 @@ def _load_pretrained_model( cls, model, state_dict: OrderedDict, - resolved_model_file: List[str], - pretrained_model_name_or_path: Union[str, os.PathLike], - loaded_keys: List[str], + resolved_model_file: list[str], + pretrained_model_name_or_path: str | os.PathLike, + loaded_keys: list[str], ignore_mismatched_sizes: bool = False, assign_to_params_buffers: bool = False, - hf_quantizer: Optional[DiffusersQuantizer] = None, + hf_quantizer: DiffusersQuantizer | None = None, low_cpu_mem_usage: bool = True, - dtype: Optional[Union[str, torch.dtype]] = None, - keep_in_fp32_modules: Optional[List[str]] = None, - device_map: Union[str, int, torch.device, Dict[str, Union[int, str, torch.device]]] = None, - offload_state_dict: Optional[bool] = None, - offload_folder: Optional[Union[str, os.PathLike]] = None, - dduf_entries: Optional[Dict[str, DDUFEntry]] = None, - is_parallel_loading_enabled: Optional[bool] = False, + dtype: str | torch.dtype | None = None, + keep_in_fp32_modules: list[str] | None = None, + device_map: str | int | torch.device | dict[str, str | int | torch.device] = None, + offload_state_dict: bool | None = None, + offload_folder: str | os.PathLike | None = None, + dduf_entries: dict[str, DDUFEntry] | None = None, + is_parallel_loading_enabled: bool | None = False, + disable_mmap: bool = False, ): model_state_dict = model.state_dict() expected_keys = list(model_state_dict.keys()) @@ -1663,6 +1680,7 @@ def _load_pretrained_model( state_dict_folder=state_dict_folder, ignore_mismatched_sizes=ignore_mismatched_sizes, low_cpu_mem_usage=low_cpu_mem_usage, + disable_mmap=disable_mmap, ) if is_parallel_loading_enabled: @@ -1672,7 +1690,10 @@ def _load_pretrained_model( else: shard_files = resolved_model_file if len(resolved_model_file) > 1: - shard_files = logging.tqdm(resolved_model_file, desc="Loading checkpoint shards") + shard_tqdm_kwargs = {"desc": "Loading checkpoint shards"} + if not is_torch_dist_rank_zero(): + shard_tqdm_kwargs["disable"] = True + shard_files = logging.tqdm(resolved_model_file, **shard_tqdm_kwargs) for shard_file in shard_files: offload_index, state_dict_index, _mismatched_keys, _error_msgs = load_fn(shard_file) @@ -1753,7 +1774,7 @@ def _get_no_split_modules(self, device_map: str): The device map value. Options are ["auto", "balanced", "balanced_low_0", "sequential"] Returns: - `List[str]`: List of modules that should not be split + `list[str]`: list of modules that should not be split """ _no_split_modules = set() modules_to_check = [self] @@ -1974,7 +1995,7 @@ class LegacyModelMixin(ModelMixin): @classmethod @validate_hf_hub_args - def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs): + def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None, **kwargs): # To prevent dependency import problem. from .model_loading_utils import _fetch_remapped_cls_from_config diff --git a/src/diffusers/models/normalization.py b/src/diffusers/models/normalization.py index ae2a6298f5f7..84ffb67bfd6a 100644 --- a/src/diffusers/models/normalization.py +++ b/src/diffusers/models/normalization.py @@ -14,7 +14,6 @@ # limitations under the License. import numbers -from typing import Dict, Optional, Tuple import torch import torch.nn as nn @@ -41,8 +40,8 @@ class AdaLayerNorm(nn.Module): def __init__( self, embedding_dim: int, - num_embeddings: Optional[int] = None, - output_dim: Optional[int] = None, + num_embeddings: int | None = None, + output_dim: int | None = None, norm_elementwise_affine: bool = False, norm_eps: float = 1e-5, chunk_dim: int = 0, @@ -62,7 +61,7 @@ def __init__( self.norm = nn.LayerNorm(output_dim // 2, norm_eps, norm_elementwise_affine) def forward( - self, x: torch.Tensor, timestep: Optional[torch.Tensor] = None, temb: Optional[torch.Tensor] = None + self, x: torch.Tensor, timestep: torch.Tensor | None = None, temb: torch.Tensor | None = None ) -> torch.Tensor: if self.emb is not None: temb = self.emb(timestep) @@ -116,8 +115,8 @@ def __init__(self, embedding_dim: int, norm_type: str = "layer_norm", bias: bool def forward( self, hidden_states: torch.Tensor, - emb: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, ...]: + emb: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, ...]: emb = self.linear(self.silu(emb)) shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp, shift_msa2, scale_msa2, gate_msa2 = emb.chunk( 9, dim=1 @@ -137,7 +136,7 @@ class AdaLayerNormZero(nn.Module): num_embeddings (`int`): The size of the embeddings dictionary. """ - def __init__(self, embedding_dim: int, num_embeddings: Optional[int] = None, norm_type="layer_norm", bias=True): + def __init__(self, embedding_dim: int, num_embeddings: int | None = None, norm_type="layer_norm", bias=True): super().__init__() if num_embeddings is not None: self.emb = CombinedTimestepLabelEmbeddings(num_embeddings, embedding_dim) @@ -158,11 +157,11 @@ def __init__(self, embedding_dim: int, num_embeddings: Optional[int] = None, nor def forward( self, x: torch.Tensor, - timestep: Optional[torch.Tensor] = None, - class_labels: Optional[torch.LongTensor] = None, - hidden_dtype: Optional[torch.dtype] = None, - emb: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + timestep: torch.Tensor | None = None, + class_labels: torch.LongTensor | None = None, + hidden_dtype: torch.dtype | None = None, + emb: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: if self.emb is not None: emb = self.emb(timestep, class_labels, hidden_dtype=hidden_dtype) emb = self.linear(self.silu(emb)) @@ -195,8 +194,8 @@ def __init__(self, embedding_dim: int, norm_type="layer_norm", bias=True): def forward( self, x: torch.Tensor, - emb: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + emb: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: emb = self.linear(self.silu(emb)) shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=1) x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] @@ -224,8 +223,8 @@ def __init__(self, embedding_dim: int, norm_eps: float, norm_elementwise_affine: def forward( self, x: torch.Tensor, - emb: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + emb: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: emb = self.linear(self.silu(emb)) scale_msa, gate_msa, scale_mlp, gate_mlp = emb.chunk(4, dim=1) x = self.norm(x) * (1 + scale_msa[:, None]) @@ -257,10 +256,10 @@ def __init__(self, embedding_dim: int, use_additional_conditions: bool = False): def forward( self, timestep: torch.Tensor, - added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, - batch_size: Optional[int] = None, - hidden_dtype: Optional[torch.dtype] = None, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + added_cond_kwargs: dict[str, torch.Tensor] | None = None, + batch_size: int | None = None, + hidden_dtype: torch.dtype | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: # No modulation happening here. added_cond_kwargs = added_cond_kwargs or {"resolution": None, "aspect_ratio": None} embedded_timestep = self.emb(timestep, **added_cond_kwargs, batch_size=batch_size, hidden_dtype=hidden_dtype) @@ -280,7 +279,7 @@ class AdaGroupNorm(nn.Module): """ def __init__( - self, embedding_dim: int, out_dim: int, num_groups: int, act_fn: Optional[str] = None, eps: float = 1e-5 + self, embedding_dim: int, out_dim: int, num_groups: int, act_fn: str | None = None, eps: float = 1e-5 ): super().__init__() self.num_groups = num_groups @@ -366,7 +365,7 @@ def __init__( eps=1e-5, bias=True, norm_type="layer_norm", - out_dim: Optional[int] = None, + out_dim: int | None = None, ): super().__init__() @@ -422,8 +421,8 @@ def forward( self, x: torch.Tensor, context: torch.Tensor, - emb: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + emb: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: emb = self.linear(self.silu(emb)) ( shift_msa, @@ -463,7 +462,7 @@ def __init__( def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: shift, scale, gate, enc_shift, enc_scale, enc_gate = self.linear(self.silu(temb)).chunk(6, dim=1) hidden_states = self.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :] encoder_hidden_states = self.norm(encoder_hidden_states) * (1 + enc_scale)[:, None, :] + enc_shift[:, None, :] @@ -632,7 +631,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: def get_normalization( norm_type: str = "batch_norm", - num_features: Optional[int] = None, + num_features: int | None = None, eps: float = 1e-5, elementwise_affine: bool = True, bias: bool = True, diff --git a/src/diffusers/models/resnet.py b/src/diffusers/models/resnet.py index c0b4ad40055a..df793b534ebb 100644 --- a/src/diffusers/models/resnet.py +++ b/src/diffusers/models/resnet.py @@ -14,7 +14,6 @@ # limitations under the License. from functools import partial -from typing import Optional, Tuple, Union import torch import torch.nn as nn @@ -75,21 +74,21 @@ def __init__( self, *, in_channels: int, - out_channels: Optional[int] = None, + out_channels: int | None = None, conv_shortcut: bool = False, dropout: float = 0.0, temb_channels: int = 512, groups: int = 32, - groups_out: Optional[int] = None, + groups_out: int | None = None, eps: float = 1e-6, non_linearity: str = "swish", time_embedding_norm: str = "ada_group", # ada_group, spatial output_scale_factor: float = 1.0, - use_in_shortcut: Optional[bool] = None, + use_in_shortcut: bool | None = None, up: bool = False, down: bool = False, conv_shortcut_bias: bool = True, - conv_2d_out_channels: Optional[int] = None, + conv_2d_out_channels: int | None = None, ): super().__init__() self.in_channels = in_channels @@ -221,24 +220,24 @@ def __init__( self, *, in_channels: int, - out_channels: Optional[int] = None, + out_channels: int | None = None, conv_shortcut: bool = False, dropout: float = 0.0, temb_channels: int = 512, groups: int = 32, - groups_out: Optional[int] = None, + groups_out: int | None = None, pre_norm: bool = True, eps: float = 1e-6, non_linearity: str = "swish", skip_time_act: bool = False, time_embedding_norm: str = "default", # default, scale_shift, - kernel: Optional[torch.Tensor] = None, + kernel: torch.Tensor | None = None, output_scale_factor: float = 1.0, - use_in_shortcut: Optional[bool] = None, + use_in_shortcut: bool | None = None, up: bool = False, down: bool = False, conv_shortcut_bias: bool = True, - conv_2d_out_channels: Optional[int] = None, + conv_2d_out_channels: int | None = None, ): super().__init__() if time_embedding_norm == "ada_group": @@ -366,7 +365,12 @@ def forward(self, input_tensor: torch.Tensor, temb: torch.Tensor, *args, **kwarg hidden_states = self.conv2(hidden_states) if self.conv_shortcut is not None: - input_tensor = self.conv_shortcut(input_tensor.contiguous()) + # Only use contiguous() during training to avoid DDP gradient stride mismatch warning. + # In inference mode (eval or no_grad), skip contiguous() for better performance, especially on CPU. + # Issue: https://github.com/huggingface/diffusers/issues/12975 + if self.training: + input_tensor = input_tensor.contiguous() + input_tensor = self.conv_shortcut(input_tensor) output_tensor = (input_tensor + hidden_states) / self.output_scale_factor @@ -401,7 +405,7 @@ def __init__( self, inp_channels: int, out_channels: int, - kernel_size: Union[int, Tuple[int, int]], + kernel_size: int | tuple[int, int], n_groups: int = 8, activation: str = "mish", ): @@ -438,7 +442,7 @@ def __init__( inp_channels: int, out_channels: int, embed_dim: int, - kernel_size: Union[int, Tuple[int, int]] = 5, + kernel_size: int | tuple[int, int] = 5, activation: str = "mish", ): super().__init__() @@ -482,7 +486,7 @@ class TemporalConvLayer(nn.Module): def __init__( self, in_dim: int, - out_dim: Optional[int] = None, + out_dim: int | None = None, dropout: float = 0.0, norm_num_groups: int = 32, ): @@ -554,7 +558,7 @@ class TemporalResnetBlock(nn.Module): def __init__( self, in_channels: int, - out_channels: Optional[int] = None, + out_channels: int | None = None, temb_channels: int = 512, eps: float = 1e-6, ): @@ -653,10 +657,10 @@ class SpatioTemporalResBlock(nn.Module): def __init__( self, in_channels: int, - out_channels: Optional[int] = None, + out_channels: int | None = None, temb_channels: int = 512, eps: float = 1e-6, - temporal_eps: Optional[float] = None, + temporal_eps: float | None = None, merge_factor: float = 0.5, merge_strategy="learned_with_images", switch_spatial_to_temporal_mix: bool = False, @@ -686,8 +690,8 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - image_only_indicator: Optional[torch.Tensor] = None, + temb: torch.Tensor | None = None, + image_only_indicator: torch.Tensor | None = None, ): num_frames = image_only_indicator.shape[-1] hidden_states = self.spatial_res_block(hidden_states, temb) @@ -785,7 +789,7 @@ def forward( self, x_spatial: torch.Tensor, x_temporal: torch.Tensor, - image_only_indicator: Optional[torch.Tensor] = None, + image_only_indicator: torch.Tensor | None = None, ) -> torch.Tensor: alpha = self.get_alpha(image_only_indicator, x_spatial.ndim) alpha = alpha.to(x_spatial.dtype) diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index 40b5d4a0dfc9..45157ee91808 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -27,6 +27,8 @@ from .transformer_easyanimate import EasyAnimateTransformer3DModel from .transformer_flux import FluxTransformer2DModel from .transformer_flux2 import Flux2Transformer2DModel + from .transformer_glm_image import GlmImageTransformer2DModel + from .transformer_helios import HeliosTransformer3DModel from .transformer_hidream_image import HiDreamImageTransformer2DModel from .transformer_hunyuan_video import HunyuanVideoTransformer3DModel from .transformer_hunyuan_video15 import HunyuanVideo15Transformer3DModel @@ -35,6 +37,7 @@ from .transformer_kandinsky import Kandinsky5Transformer3DModel from .transformer_longcat_image import LongCatImageTransformer2DModel from .transformer_ltx import LTXVideoTransformer3DModel + from .transformer_ltx2 import LTX2VideoTransformer3DModel from .transformer_lumina2 import Lumina2Transformer2DModel from .transformer_mochi import MochiTransformer3DModel from .transformer_omnigen import OmniGenTransformer2DModel diff --git a/src/diffusers/models/transformers/auraflow_transformer_2d.py b/src/diffusers/models/transformers/auraflow_transformer_2d.py index e3732662e408..3fa4df738784 100644 --- a/src/diffusers/models/transformers/auraflow_transformer_2d.py +++ b/src/diffusers/models/transformers/auraflow_transformer_2d.py @@ -13,7 +13,7 @@ # limitations under the License. -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -21,7 +21,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, logging from ...utils.torch_utils import maybe_allow_in_graph from ..attention import AttentionMixin from ..attention_processor import ( @@ -172,7 +172,7 @@ def forward( self, hidden_states: torch.FloatTensor, temb: torch.FloatTensor, - attention_kwargs: Optional[Dict[str, Any]] = None, + attention_kwargs: dict[str, Any] | None = None, ) -> torch.Tensor: residual = hidden_states attention_kwargs = attention_kwargs or {} @@ -241,8 +241,8 @@ def forward( hidden_states: torch.FloatTensor, encoder_hidden_states: torch.FloatTensor, temb: torch.FloatTensor, - attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + attention_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: residual = hidden_states residual_context = encoder_hidden_states attention_kwargs = attention_kwargs or {} @@ -397,29 +397,15 @@ def unfuse_qkv_projections(self): if self.original_attn_processors is not None: self.set_attn_processor(self.original_attn_processors) + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.FloatTensor, encoder_hidden_states: torch.FloatTensor = None, timestep: torch.LongTensor = None, - attention_kwargs: Optional[Dict[str, Any]] = None, + attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) - + ) -> tuple[torch.Tensor] | Transformer2DModelOutput: height, width = hidden_states.shape[-2:] # Apply patch embedding, timestep embedding, and project the caption embeddings. @@ -486,10 +472,6 @@ def forward( shape=(hidden_states.shape[0], out_channels, height * patch_size, width * patch_size) ) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) diff --git a/src/diffusers/models/transformers/cogvideox_transformer_3d.py b/src/diffusers/models/transformers/cogvideox_transformer_3d.py index 14b38cd46c52..4b8beeeb6fe3 100644 --- a/src/diffusers/models/transformers/cogvideox_transformer_3d.py +++ b/src/diffusers/models/transformers/cogvideox_transformer_3d.py @@ -13,14 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any import torch from torch import nn from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, logging from ...utils.torch_utils import maybe_allow_in_graph from ..attention import Attention, AttentionMixin, FeedForward from ..attention_processor import CogVideoXAttnProcessor2_0, FusedCogVideoXAttnProcessor2_0 @@ -83,7 +83,7 @@ def __init__( norm_elementwise_affine: bool = True, norm_eps: float = 1e-5, final_dropout: bool = True, - ff_inner_dim: Optional[int] = None, + ff_inner_dim: int | None = None, ff_bias: bool = True, attention_out_bias: bool = True, ): @@ -120,9 +120,9 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: text_seq_length = encoder_hidden_states.size(1) attention_kwargs = attention_kwargs or {} @@ -223,11 +223,11 @@ def __init__( num_attention_heads: int = 30, attention_head_dim: int = 64, in_channels: int = 16, - out_channels: Optional[int] = 16, + out_channels: int | None = 16, flip_sin_to_cos: bool = True, freq_shift: int = 0, time_embed_dim: int = 512, - ofs_embed_dim: Optional[int] = None, + ofs_embed_dim: int | None = None, text_embed_dim: int = 4096, num_layers: int = 30, dropout: float = 0.0, @@ -236,7 +236,7 @@ def __init__( sample_height: int = 60, sample_frames: int = 49, patch_size: int = 2, - patch_size_t: Optional[int] = None, + patch_size_t: int | None = None, temporal_compression_ratio: int = 4, max_text_seq_length: int = 226, activation_fn: str = "gelu-approximate", @@ -363,32 +363,18 @@ def unfuse_qkv_projections(self): if self.original_attn_processors is not None: self.set_attn_processor(self.original_attn_processors) + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, - timestep: Union[int, float, torch.LongTensor], - timestep_cond: Optional[torch.Tensor] = None, - ofs: Optional[Union[int, float, torch.LongTensor]] = None, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - attention_kwargs: Optional[Dict[str, Any]] = None, + timestep: int | float | torch.LongTensor, + timestep_cond: torch.Tensor | None = None, + ofs: int | float | torch.LongTensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) - + ) -> tuple[torch.Tensor] | Transformer2DModelOutput: batch_size, num_frames, channels, height, width = hidden_states.shape # 1. Time embedding @@ -454,10 +440,6 @@ def forward( ) output = output.permute(0, 1, 5, 4, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(1, 2) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) return Transformer2DModelOutput(sample=output) diff --git a/src/diffusers/models/transformers/consisid_transformer_3d.py b/src/diffusers/models/transformers/consisid_transformer_3d.py index be20b0a3eacf..64a58e394366 100644 --- a/src/diffusers/models/transformers/consisid_transformer_3d.py +++ b/src/diffusers/models/transformers/consisid_transformer_3d.py @@ -13,14 +13,14 @@ # limitations under the License. import math -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch from torch import nn from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, logging from ...utils.torch_utils import maybe_allow_in_graph from ..attention import Attention, AttentionMixin, FeedForward from ..attention_processor import CogVideoXAttnProcessor2_0 @@ -34,7 +34,7 @@ class PerceiverAttention(nn.Module): - def __init__(self, dim: int, dim_head: int = 64, heads: int = 8, kv_dim: Optional[int] = None): + def __init__(self, dim: int, dim_head: int = 64, heads: int = 8, kv_dim: int | None = None): super().__init__() self.scale = dim_head**-0.5 @@ -152,7 +152,7 @@ def __init__( nn.Linear(vit_dim, vit_dim * num_id_token), ) - def forward(self, id_embeds: torch.Tensor, vit_hidden_states: List[torch.Tensor]) -> torch.Tensor: + def forward(self, id_embeds: torch.Tensor, vit_hidden_states: list[torch.Tensor]) -> torch.Tensor: # Repeat latent queries for the batch size latents = self.latents.repeat(id_embeds.size(0), 1, 1) @@ -277,7 +277,7 @@ def __init__( norm_elementwise_affine: bool = True, norm_eps: float = 1e-5, final_dropout: bool = True, - ff_inner_dim: Optional[int] = None, + ff_inner_dim: int | None = None, ff_bias: bool = True, attention_out_bias: bool = True, ): @@ -314,8 +314,8 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: text_seq_length = encoder_hidden_states.size(1) # norm & modulate @@ -465,7 +465,7 @@ def __init__( num_attention_heads: int = 30, attention_head_dim: int = 64, in_channels: int = 16, - out_channels: Optional[int] = 16, + out_channels: int | None = 16, flip_sin_to_cos: bool = True, freq_shift: int = 0, time_embed_dim: int = 512, @@ -620,33 +620,19 @@ def _init_face_inputs(self): ] ) + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, - timestep: Union[int, float, torch.LongTensor], - timestep_cond: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - attention_kwargs: Optional[Dict[str, Any]] = None, - id_cond: Optional[torch.Tensor] = None, - id_vit_hidden: Optional[torch.Tensor] = None, + timestep: int | float | torch.LongTensor, + timestep_cond: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_kwargs: dict[str, Any] | None = None, + id_cond: torch.Tensor | None = None, + id_vit_hidden: torch.Tensor | None = None, return_dict: bool = True, - ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) - + ) -> tuple[torch.Tensor] | Transformer2DModelOutput: # fuse clip and insightface valid_face_emb = None if self.is_train_face: @@ -720,10 +706,6 @@ def forward( output = hidden_states.reshape(batch_size, num_frames, height // p, width // p, -1, p, p) output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) return Transformer2DModelOutput(sample=output) diff --git a/src/diffusers/models/transformers/dit_transformer_2d.py b/src/diffusers/models/transformers/dit_transformer_2d.py index 68f6f769436e..3d10c278cdbb 100644 --- a/src/diffusers/models/transformers/dit_transformer_2d.py +++ b/src/diffusers/models/transformers/dit_transformer_2d.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, Optional +from typing import Any import torch import torch.nn.functional as F @@ -74,7 +74,7 @@ def __init__( num_attention_heads: int = 16, attention_head_dim: int = 72, in_channels: int = 4, - out_channels: Optional[int] = None, + out_channels: int | None = None, num_layers: int = 28, dropout: float = 0.0, norm_num_groups: int = 32, @@ -82,7 +82,7 @@ def __init__( sample_size: int = 32, patch_size: int = 2, activation_fn: str = "gelu-approximate", - num_embeds_ada_norm: Optional[int] = 1000, + num_embeds_ada_norm: int | None = 1000, upcast_attention: bool = False, norm_type: str = "ada_norm_zero", norm_elementwise_affine: bool = False, @@ -148,9 +148,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - timestep: Optional[torch.LongTensor] = None, - class_labels: Optional[torch.LongTensor] = None, - cross_attention_kwargs: Dict[str, Any] = None, + timestep: torch.LongTensor | None = None, + class_labels: torch.LongTensor | None = None, + cross_attention_kwargs: dict[str, Any] = None, return_dict: bool = True, ): """ @@ -164,7 +164,7 @@ def forward( class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*): Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in `AdaLayerZeroNorm`. - cross_attention_kwargs ( `Dict[str, Any]`, *optional*): + cross_attention_kwargs ( `dict[str, Any]`, *optional*): A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). diff --git a/src/diffusers/models/transformers/dual_transformer_2d.py b/src/diffusers/models/transformers/dual_transformer_2d.py index 24eed2168229..c25c6e9c4227 100644 --- a/src/diffusers/models/transformers/dual_transformer_2d.py +++ b/src/diffusers/models/transformers/dual_transformer_2d.py @@ -11,8 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional - from torch import nn from ..modeling_outputs import Transformer2DModelOutput @@ -50,16 +48,16 @@ def __init__( self, num_attention_heads: int = 16, attention_head_dim: int = 88, - in_channels: Optional[int] = None, + in_channels: int | None = None, num_layers: int = 1, dropout: float = 0.0, norm_num_groups: int = 32, - cross_attention_dim: Optional[int] = None, + cross_attention_dim: int | None = None, attention_bias: bool = False, - sample_size: Optional[int] = None, - num_vector_embeds: Optional[int] = None, + sample_size: int | None = None, + num_vector_embeds: int | None = None, activation_fn: str = "geglu", - num_embeds_ada_norm: Optional[int] = None, + num_embeds_ada_norm: int | None = None, ): super().__init__() self.transformers = nn.ModuleList( diff --git a/src/diffusers/models/transformers/hunyuan_transformer_2d.py b/src/diffusers/models/transformers/hunyuan_transformer_2d.py index cecb675b32b7..a25aa99fb8b9 100644 --- a/src/diffusers/models/transformers/hunyuan_transformer_2d.py +++ b/src/diffusers/models/transformers/hunyuan_transformer_2d.py @@ -11,8 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional - import torch from torch import nn @@ -98,7 +96,7 @@ def __init__( norm_elementwise_affine: bool = True, norm_eps: float = 1e-6, final_dropout: bool = False, - ff_inner_dim: Optional[int] = None, + ff_inner_dim: int | None = None, ff_bias: bool = True, skip: bool = False, qk_norm: bool = True, @@ -158,7 +156,7 @@ def __init__( self._chunk_dim = 0 # Copied from diffusers.models.attention.BasicTransformerBlock.set_chunk_feed_forward - def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0): + def set_chunk_feed_forward(self, chunk_size: int | None, dim: int = 0): # Sets chunk feed-forward self._chunk_size = chunk_size self._chunk_dim = dim @@ -166,8 +164,8 @@ def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0): def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - temb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + temb: torch.Tensor | None = None, image_rotary_emb=None, skip=None, ) -> torch.Tensor: @@ -252,8 +250,8 @@ def __init__( self, num_attention_heads: int = 16, attention_head_dim: int = 88, - in_channels: Optional[int] = None, - patch_size: Optional[int] = None, + in_channels: int | None = None, + patch_size: int | None = None, activation_fn: str = "gelu-approximate", sample_size=32, hidden_size=1152, @@ -469,7 +467,7 @@ def forward( return Transformer2DModelOutput(sample=output) # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking - def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None: + def enable_forward_chunking(self, chunk_size: int | None = None, dim: int = 0) -> None: """ Sets the attention processor to use [feed forward chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers). diff --git a/src/diffusers/models/transformers/latte_transformer_3d.py b/src/diffusers/models/transformers/latte_transformer_3d.py index 990c90512e39..32e97aff8fb7 100644 --- a/src/diffusers/models/transformers/latte_transformer_3d.py +++ b/src/diffusers/models/transformers/latte_transformer_3d.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional - import torch from torch import nn @@ -73,16 +71,16 @@ def __init__( self, num_attention_heads: int = 16, attention_head_dim: int = 88, - in_channels: Optional[int] = None, - out_channels: Optional[int] = None, + in_channels: int | None = None, + out_channels: int | None = None, num_layers: int = 1, dropout: float = 0.0, - cross_attention_dim: Optional[int] = None, + cross_attention_dim: int | None = None, attention_bias: bool = False, sample_size: int = 64, - patch_size: Optional[int] = None, + patch_size: int | None = None, activation_fn: str = "geglu", - num_embeds_ada_norm: Optional[int] = None, + num_embeds_ada_norm: int | None = None, norm_type: str = "layer_norm", norm_elementwise_affine: bool = True, norm_eps: float = 1e-5, @@ -168,9 +166,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - timestep: Optional[torch.LongTensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, + timestep: torch.LongTensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, enable_temporal_attentions: bool = True, return_dict: bool = True, ): diff --git a/src/diffusers/models/transformers/lumina_nextdit2d.py b/src/diffusers/models/transformers/lumina_nextdit2d.py index bed5e69c2d36..46a6753b4cb1 100644 --- a/src/diffusers/models/transformers/lumina_nextdit2d.py +++ b/src/diffusers/models/transformers/lumina_nextdit2d.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -123,7 +123,7 @@ def forward( encoder_hidden_states: torch.Tensor, encoder_mask: torch.Tensor, temb: torch.Tensor, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, ) -> torch.Tensor: """ Perform a forward pass through the LuminaNextDiTBlock. @@ -135,7 +135,7 @@ def forward( encoder_hidden_states: (`torch.Tensor`): The hidden_states of text prompt are processed by Gemma encoder. encoder_mask (`torch.Tensor`): The hidden_states of text prompt attention mask. temb (`torch.Tensor`): Timestep embedding with text prompt embedding. - cross_attention_kwargs (`Dict[str, Any]`): kwargs for cross attention. + cross_attention_kwargs (`dict[str, Any]`): kwargs for cross attention. """ residual = hidden_states @@ -227,19 +227,19 @@ class LuminaNextDiT2DModel(ModelMixin, ConfigMixin): def __init__( self, sample_size: int = 128, - patch_size: Optional[int] = 2, - in_channels: Optional[int] = 4, - hidden_size: Optional[int] = 2304, - num_layers: Optional[int] = 32, - num_attention_heads: Optional[int] = 32, - num_kv_heads: Optional[int] = None, - multiple_of: Optional[int] = 256, - ffn_dim_multiplier: Optional[float] = None, - norm_eps: Optional[float] = 1e-5, - learn_sigma: Optional[bool] = True, - qk_norm: Optional[bool] = True, - cross_attention_dim: Optional[int] = 2048, - scaling_factor: Optional[float] = 1.0, + patch_size: int | None = 2, + in_channels: int | None = 4, + hidden_size: int | None = 2304, + num_layers: int | None = 32, + num_attention_heads: int | None = 32, + num_kv_heads: int | None = None, + multiple_of: int | None = 256, + ffn_dim_multiplier: float | None = None, + norm_eps: float | None = 1e-5, + learn_sigma: bool | None = True, + qk_norm: bool | None = True, + cross_attention_dim: int | None = 2048, + scaling_factor: float | None = 1.0, ) -> None: super().__init__() self.sample_size = sample_size @@ -295,9 +295,9 @@ def forward( encoder_hidden_states: torch.Tensor, encoder_mask: torch.Tensor, image_rotary_emb: torch.Tensor, - cross_attention_kwargs: Dict[str, Any] = None, + cross_attention_kwargs: dict[str, Any] = None, return_dict=True, - ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: + ) -> tuple[torch.Tensor] | Transformer2DModelOutput: """ Forward pass of LuminaNextDiT. diff --git a/src/diffusers/models/transformers/pixart_transformer_2d.py b/src/diffusers/models/transformers/pixart_transformer_2d.py index 072670ee0c30..2476668ba307 100644 --- a/src/diffusers/models/transformers/pixart_transformer_2d.py +++ b/src/diffusers/models/transformers/pixart_transformer_2d.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, Optional +from typing import Any import torch from torch import nn @@ -87,24 +87,24 @@ def __init__( num_attention_heads: int = 16, attention_head_dim: int = 72, in_channels: int = 4, - out_channels: Optional[int] = 8, + out_channels: int | None = 8, num_layers: int = 28, dropout: float = 0.0, norm_num_groups: int = 32, - cross_attention_dim: Optional[int] = 1152, + cross_attention_dim: int | None = 1152, attention_bias: bool = True, sample_size: int = 128, patch_size: int = 2, activation_fn: str = "gelu-approximate", - num_embeds_ada_norm: Optional[int] = 1000, + num_embeds_ada_norm: int | None = 1000, upcast_attention: bool = False, norm_type: str = "ada_norm_single", norm_elementwise_affine: bool = False, norm_eps: float = 1e-6, - interpolation_scale: Optional[int] = None, - use_additional_conditions: Optional[bool] = None, - caption_channels: Optional[int] = None, - attention_type: Optional[str] = "default", + interpolation_scale: int | None = None, + use_additional_conditions: bool | None = None, + caption_channels: int | None = None, + attention_type: str | None = "default", ): super().__init__() @@ -227,12 +227,12 @@ def unfuse_qkv_projections(self): def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - timestep: Optional[torch.LongTensor] = None, - added_cond_kwargs: Dict[str, torch.Tensor] = None, - cross_attention_kwargs: Dict[str, Any] = None, - attention_mask: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + timestep: torch.LongTensor | None = None, + added_cond_kwargs: dict[str, torch.Tensor] = None, + cross_attention_kwargs: dict[str, Any] = None, + attention_mask: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, return_dict: bool = True, ): """ @@ -246,8 +246,8 @@ def forward( self-attention. timestep (`torch.LongTensor`, *optional*): Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`. - added_cond_kwargs: (`Dict[str, Any]`, *optional*): Additional conditions to be used as inputs. - cross_attention_kwargs ( `Dict[str, Any]`, *optional*): + added_cond_kwargs: (`dict[str, Any]`, *optional*): Additional conditions to be used as inputs. + cross_attention_kwargs ( `dict[str, Any]`, *optional*): A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). diff --git a/src/diffusers/models/transformers/prior_transformer.py b/src/diffusers/models/transformers/prior_transformer.py index 757bb436040f..ace2b529c4f2 100644 --- a/src/diffusers/models/transformers/prior_transformer.py +++ b/src/diffusers/models/transformers/prior_transformer.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import Optional, Union import torch import torch.nn.functional as F @@ -80,13 +79,13 @@ def __init__( additional_embeddings=4, dropout: float = 0.0, time_embed_act_fn: str = "silu", - norm_in_type: Optional[str] = None, # layer - embedding_proj_norm_type: Optional[str] = None, # layer - encoder_hid_proj_type: Optional[str] = "linear", # linear - added_emb_type: Optional[str] = "prd", # prd - time_embed_dim: Optional[int] = None, - embedding_proj_dim: Optional[int] = None, - clip_embed_dim: Optional[int] = None, + norm_in_type: str | None = None, # layer + embedding_proj_norm_type: str | None = None, # layer + encoder_hid_proj_type: str | None = "linear", # linear + added_emb_type: str | None = "prd", # prd + time_embed_dim: int | None = None, + embedding_proj_dim: int | None = None, + clip_embed_dim: int | None = None, ): super().__init__() self.num_attention_heads = num_attention_heads @@ -184,10 +183,10 @@ def set_default_attn_processor(self): def forward( self, hidden_states, - timestep: Union[torch.Tensor, float, int], + timestep: torch.Tensor | float | int, proj_embedding: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.BoolTensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.BoolTensor | None = None, return_dict: bool = True, ): """ diff --git a/src/diffusers/models/transformers/sana_transformer.py b/src/diffusers/models/transformers/sana_transformer.py index d45badb1b121..ff078dc695d7 100644 --- a/src/diffusers/models/transformers/sana_transformer.py +++ b/src/diffusers/models/transformers/sana_transformer.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any import torch import torch.nn.functional as F @@ -20,7 +20,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, logging from ..attention import AttentionMixin from ..attention_processor import ( Attention, @@ -41,7 +41,7 @@ def __init__( in_channels: int, out_channels: int, expand_ratio: float = 4, - norm_type: Optional[str] = None, + norm_type: str | None = None, residual_connection: bool = True, ) -> None: super().__init__() @@ -132,8 +132,8 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: batch_size, sequence_length, _ = ( hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape @@ -196,15 +196,15 @@ def __init__( num_attention_heads: int = 70, attention_head_dim: int = 32, dropout: float = 0.0, - num_cross_attention_heads: Optional[int] = 20, - cross_attention_head_dim: Optional[int] = 112, - cross_attention_dim: Optional[int] = 2240, + num_cross_attention_heads: int | None = 20, + cross_attention_head_dim: int | None = 112, + cross_attention_dim: int | None = 2240, attention_bias: bool = True, norm_elementwise_affine: bool = False, norm_eps: float = 1e-6, attention_out_bias: bool = True, mlp_ratio: float = 2.5, - qk_norm: Optional[str] = None, + qk_norm: str | None = None, ) -> None: super().__init__() @@ -246,10 +246,10 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - timestep: Optional[torch.LongTensor] = None, + attention_mask: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + timestep: torch.LongTensor | None = None, height: int = None, width: int = None, ) -> torch.Tensor: @@ -340,13 +340,13 @@ class SanaTransformer2DModel(ModelMixin, AttentionMixin, ConfigMixin, PeftAdapte def __init__( self, in_channels: int = 32, - out_channels: Optional[int] = 32, + out_channels: int | None = 32, num_attention_heads: int = 70, attention_head_dim: int = 32, num_layers: int = 20, - num_cross_attention_heads: Optional[int] = 20, - cross_attention_head_dim: Optional[int] = 112, - cross_attention_dim: Optional[int] = 2240, + num_cross_attention_heads: int | None = 20, + cross_attention_head_dim: int | None = 112, + cross_attention_dim: int | None = 2240, caption_channels: int = 2304, mlp_ratio: float = 2.5, dropout: float = 0.0, @@ -355,10 +355,10 @@ def __init__( patch_size: int = 1, norm_elementwise_affine: bool = False, norm_eps: float = 1e-6, - interpolation_scale: Optional[int] = None, + interpolation_scale: int | None = None, guidance_embeds: bool = False, guidance_embeds_scale: float = 0.1, - qk_norm: Optional[str] = None, + qk_norm: str | None = None, timestep_scale: float = 1.0, ) -> None: super().__init__() @@ -414,33 +414,19 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, timestep: torch.Tensor, - guidance: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_block_samples: Optional[Tuple[torch.Tensor]] = None, + guidance: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, + controlnet_block_samples: tuple[torch.Tensor] | None = None, return_dict: bool = True, - ) -> Union[Tuple[torch.Tensor, ...], Transformer2DModelOutput]: - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) - + ) -> tuple[torch.Tensor, ...] | Transformer2DModelOutput: # ensure attention_mask is a bias, and give it a singleton query_tokens dimension. # we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward. # we can tell by counting dims; if ndim == 2: it's a mask rather than a bias. @@ -527,10 +513,6 @@ def forward( hidden_states = hidden_states.permute(0, 5, 1, 3, 2, 4) output = hidden_states.reshape(batch_size, -1, post_patch_height * p, post_patch_width * p) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) diff --git a/src/diffusers/models/transformers/stable_audio_transformer.py b/src/diffusers/models/transformers/stable_audio_transformer.py index 2c3b6b5df91d..f4974926ec72 100644 --- a/src/diffusers/models/transformers/stable_audio_transformer.py +++ b/src/diffusers/models/transformers/stable_audio_transformer.py @@ -13,8 +13,6 @@ # limitations under the License. -from typing import Optional, Union - import numpy as np import torch import torch.nn as nn @@ -87,10 +85,10 @@ def __init__( num_key_value_attention_heads: int, attention_head_dim: int, dropout=0.0, - cross_attention_dim: Optional[int] = None, + cross_attention_dim: int | None = None, upcast_attention: bool = False, norm_eps: float = 1e-5, - ff_inner_dim: Optional[int] = None, + ff_inner_dim: int | None = None, ): super().__init__() # Define 3 blocks. Each block has its own normalization layer. @@ -138,7 +136,7 @@ def __init__( self._chunk_size = None self._chunk_dim = 0 - def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0): + def set_chunk_feed_forward(self, chunk_size: int | None, dim: int = 0): # Sets chunk feed-forward self._chunk_size = chunk_size self._chunk_dim = dim @@ -146,10 +144,10 @@ def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0): def forward( self, hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - rotary_embedding: Optional[torch.FloatTensor] = None, + attention_mask: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + rotary_embedding: torch.FloatTensor | None = None, ) -> torch.Tensor: # Notice that normalization is always applied before the real computation in the following blocks. # 0. Self-Attention @@ -289,9 +287,9 @@ def forward( global_hidden_states: torch.FloatTensor = None, rotary_embedding: torch.FloatTensor = None, return_dict: bool = True, - attention_mask: Optional[torch.LongTensor] = None, - encoder_attention_mask: Optional[torch.LongTensor] = None, - ) -> Union[torch.FloatTensor, Transformer2DModelOutput]: + attention_mask: torch.LongTensor | None = None, + encoder_attention_mask: torch.LongTensor | None = None, + ) -> torch.FloatTensor | Transformer2DModelOutput: """ The [`StableAudioDiTModel`] forward method. diff --git a/src/diffusers/models/transformers/t5_film_transformer.py b/src/diffusers/models/transformers/t5_film_transformer.py index 7a9608735e32..1ae2b1e3fedb 100644 --- a/src/diffusers/models/transformers/t5_film_transformer.py +++ b/src/diffusers/models/transformers/t5_film_transformer.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. import math -from typing import Optional, Tuple import torch from torch import nn @@ -196,12 +195,12 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - conditioning_emb: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, + conditioning_emb: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, encoder_decoder_position_bias=None, - ) -> Tuple[torch.Tensor]: + ) -> tuple[torch.Tensor]: hidden_states = self.layer[0]( hidden_states, conditioning_emb=conditioning_emb, @@ -250,8 +249,8 @@ def __init__(self, d_model: int, d_kv: int, num_heads: int, dropout_rate: float) def forward( self, hidden_states: torch.Tensor, - conditioning_emb: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + conditioning_emb: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: # pre_self_attention_layer_norm normed_hidden_states = self.layer_norm(hidden_states) @@ -293,8 +292,8 @@ def __init__(self, d_model: int, d_kv: int, num_heads: int, dropout_rate: float, def forward( self, hidden_states: torch.Tensor, - key_value_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + key_value_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: normed_hidden_states = self.layer_norm(hidden_states) attention_output = self.attention( @@ -328,7 +327,7 @@ def __init__(self, d_model: int, d_ff: int, dropout_rate: float, layer_norm_epsi self.layer_norm = T5LayerNorm(d_model, eps=layer_norm_epsilon) self.dropout = nn.Dropout(dropout_rate) - def forward(self, hidden_states: torch.Tensor, conditioning_emb: Optional[torch.Tensor] = None) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, conditioning_emb: torch.Tensor | None = None) -> torch.Tensor: forwarded_states = self.layer_norm(hidden_states) if conditioning_emb is not None: forwarded_states = self.film(forwarded_states, conditioning_emb) diff --git a/src/diffusers/models/transformers/transformer_2d.py b/src/diffusers/models/transformers/transformer_2d.py index 67fe9a33109b..12f89201d752 100644 --- a/src/diffusers/models/transformers/transformer_2d.py +++ b/src/diffusers/models/transformers/transformer_2d.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, Optional +from typing import Any import torch import torch.nn.functional as F @@ -73,18 +73,18 @@ def __init__( self, num_attention_heads: int = 16, attention_head_dim: int = 88, - in_channels: Optional[int] = None, - out_channels: Optional[int] = None, + in_channels: int | None = None, + out_channels: int | None = None, num_layers: int = 1, dropout: float = 0.0, norm_num_groups: int = 32, - cross_attention_dim: Optional[int] = None, + cross_attention_dim: int | None = None, attention_bias: bool = False, - sample_size: Optional[int] = None, - num_vector_embeds: Optional[int] = None, - patch_size: Optional[int] = None, + sample_size: int | None = None, + num_vector_embeds: int | None = None, + patch_size: int | None = None, activation_fn: str = "geglu", - num_embeds_ada_norm: Optional[int] = None, + num_embeds_ada_norm: int | None = None, use_linear_projection: bool = False, only_cross_attention: bool = False, double_self_attention: bool = False, @@ -95,7 +95,7 @@ def __init__( attention_type: str = "default", caption_channels: int = None, interpolation_scale: float = None, - use_additional_conditions: Optional[bool] = None, + use_additional_conditions: bool | None = None, ): super().__init__() @@ -324,13 +324,13 @@ def _init_patched_inputs(self, norm_type): def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - timestep: Optional[torch.LongTensor] = None, - added_cond_kwargs: Dict[str, torch.Tensor] = None, - class_labels: Optional[torch.LongTensor] = None, - cross_attention_kwargs: Dict[str, Any] = None, - attention_mask: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + timestep: torch.LongTensor | None = None, + added_cond_kwargs: dict[str, torch.Tensor] = None, + class_labels: torch.LongTensor | None = None, + cross_attention_kwargs: dict[str, Any] = None, + attention_mask: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, return_dict: bool = True, ): """ @@ -347,7 +347,7 @@ def forward( class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*): Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in `AdaLayerZeroNorm`. - cross_attention_kwargs ( `Dict[str, Any]`, *optional*): + cross_attention_kwargs ( `dict[str, Any]`, *optional*): A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). diff --git a/src/diffusers/models/transformers/transformer_allegro.py b/src/diffusers/models/transformers/transformer_allegro.py index 5fa59a71d977..934e2787674f 100644 --- a/src/diffusers/models/transformers/transformer_allegro.py +++ b/src/diffusers/models/transformers/transformer_allegro.py @@ -13,8 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Tuple - import torch import torch.nn as nn import torch.nn.functional as F @@ -69,7 +67,7 @@ def __init__( num_attention_heads: int, attention_head_dim: int, dropout=0.0, - cross_attention_dim: Optional[int] = None, + cross_attention_dim: int | None = None, activation_fn: str = "geglu", attention_bias: bool = False, norm_elementwise_affine: bool = True, @@ -117,10 +115,10 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - temb: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + temb: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, image_rotary_emb=None, ) -> torch.Tensor: # 0. Self-Attention @@ -309,9 +307,9 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, timestep: torch.LongTensor, - attention_mask: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + attention_mask: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, return_dict: bool = True, ): batch_size, num_channels, num_frames, height, width = hidden_states.shape diff --git a/src/diffusers/models/transformers/transformer_bria.py b/src/diffusers/models/transformers/transformer_bria.py index d54679306e64..99b7bbfd64cf 100644 --- a/src/diffusers/models/transformers/transformer_bria.py +++ b/src/diffusers/models/transformers/transformer_bria.py @@ -1,5 +1,5 @@ import inspect -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import numpy as np import torch @@ -8,7 +8,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, logging from ...utils.torch_utils import maybe_allow_in_graph from ..attention import AttentionModuleMixin, FeedForward from ..attention_dispatch import dispatch_attention_fn @@ -54,7 +54,7 @@ def _get_qkv_projections(attn: "BriaAttention", hidden_states, encoder_hidden_st def get_1d_rotary_pos_embed( dim: int, - pos: Union[np.ndarray, int], + pos: np.ndarray | int, theta: float = 10000.0, use_real=False, linear_factor=1.0, @@ -131,8 +131,8 @@ def __call__( attn: "BriaAttention", hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections( attn, hidden_states, encoder_hidden_states @@ -198,12 +198,12 @@ def __init__( dim_head: int = 64, dropout: float = 0.0, bias: bool = False, - added_kv_proj_dim: Optional[int] = None, - added_proj_bias: Optional[bool] = True, + added_kv_proj_dim: int | None = None, + added_proj_bias: bool | None = True, out_bias: bool = True, eps: float = 1e-5, out_dim: int = None, - context_pre_only: Optional[bool] = None, + context_pre_only: bool | None = None, pre_only: bool = False, elementwise_affine: bool = True, processor=None, @@ -248,9 +248,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) @@ -266,7 +266,7 @@ def forward( class BriaEmbedND(torch.nn.Module): # modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11 - def __init__(self, theta: int, axes_dim: List[int]): + def __init__(self, theta: int, axes_dim: list[int]): super().__init__() self.theta = theta self.axes_dim = axes_dim @@ -334,7 +334,7 @@ def forward(self, timestep, dtype): class BriaPosEmbed(torch.nn.Module): # modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11 - def __init__(self, theta: int, axes_dim: List[int]): + def __init__(self, theta: int, axes_dim: list[int]): super().__init__() self.theta = theta self.axes_dim = axes_dim @@ -395,9 +395,9 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( @@ -476,9 +476,9 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: text_seq_len = encoder_hidden_states.shape[1] hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) @@ -537,7 +537,7 @@ def __init__( joint_attention_dim: int = 4096, pooled_projection_dim: int = None, guidance_embeds: bool = False, - axes_dims_rope: List[int] = [16, 56, 56], + axes_dims_rope: list[int] = [16, 56, 56], rope_theta=10000, time_theta=10000, ): @@ -581,6 +581,7 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, @@ -590,11 +591,11 @@ def forward( img_ids: torch.Tensor = None, txt_ids: torch.Tensor = None, guidance: torch.Tensor = None, - attention_kwargs: Optional[Dict[str, Any]] = None, + attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, controlnet_block_samples=None, controlnet_single_block_samples=None, - ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: + ) -> tuple[torch.Tensor] | Transformer2DModelOutput: """ The [`BriaTransformer2DModel`] forward method. @@ -621,20 +622,6 @@ def forward( If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a `tuple` where the first element is the sample tensor. """ - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) hidden_states = self.x_embedder(hidden_states) timestep = timestep.to(hidden_states.dtype) @@ -715,10 +702,6 @@ def forward( hidden_states = self.norm_out(hidden_states, temb) output = self.proj_out(hidden_states) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) diff --git a/src/diffusers/models/transformers/transformer_bria_fibo.py b/src/diffusers/models/transformers/transformer_bria_fibo.py index 09f79619320d..7ddbccfa47c5 100644 --- a/src/diffusers/models/transformers/transformer_bria_fibo.py +++ b/src/diffusers/models/transformers/transformer_bria_fibo.py @@ -8,7 +8,7 @@ # # See the license for further details. import inspect -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -22,10 +22,8 @@ from ...models.modeling_utils import ModelMixin from ...models.transformers.transformer_bria import BriaAttnProcessor from ...utils import ( - USE_PEFT_BACKEND, + apply_lora_scale, logging, - scale_lora_layers, - unscale_lora_layers, ) from ...utils.torch_utils import maybe_allow_in_graph from ..attention import AttentionModuleMixin, FeedForward @@ -80,8 +78,8 @@ def __call__( attn: "BriaFiboAttention", hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections( attn, hidden_states, encoder_hidden_states @@ -125,9 +123,9 @@ def __call__( encoder_hidden_states, hidden_states = hidden_states.split_with_sizes( [encoder_hidden_states.shape[1], hidden_states.shape[1] - encoder_hidden_states.shape[1]], dim=1 ) - hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[0](hidden_states.contiguous()) hidden_states = attn.to_out[1](hidden_states) - encoder_hidden_states = attn.to_add_out(encoder_hidden_states) + encoder_hidden_states = attn.to_add_out(encoder_hidden_states.contiguous()) return hidden_states, encoder_hidden_states else: @@ -146,12 +144,12 @@ def __init__( dim_head: int = 64, dropout: float = 0.0, bias: bool = False, - added_kv_proj_dim: Optional[int] = None, - added_proj_bias: Optional[bool] = True, + added_kv_proj_dim: int | None = None, + added_proj_bias: bool | None = True, out_bias: bool = True, eps: float = 1e-5, out_dim: int = None, - context_pre_only: Optional[bool] = None, + context_pre_only: bool | None = None, pre_only: bool = False, elementwise_affine: bool = True, processor=None, @@ -196,9 +194,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) @@ -214,7 +212,7 @@ def forward( class BriaFiboEmbedND(torch.nn.Module): # modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11 - def __init__(self, theta: int, axes_dim: List[int]): + def __init__(self, theta: int, axes_dim: list[int]): super().__init__() self.theta = theta self.axes_dim = axes_dim @@ -272,8 +270,8 @@ def forward( self, hidden_states: torch.Tensor, temb: torch.Tensor, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, ) -> torch.Tensor: residual = hidden_states norm_hidden_states, gate = self.norm(hidden_states, emb=temb) @@ -339,9 +337,9 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( @@ -458,7 +456,7 @@ def __init__( joint_attention_dim: int = 4096, pooled_projection_dim: int = None, guidance_embeds: bool = False, - axes_dims_rope: List[int] = [16, 56, 56], + axes_dims_rope: list[int] = [16, 56, 56], rope_theta=10000, time_theta=10000, text_encoder_dim: int = 2048, @@ -510,6 +508,7 @@ def __init__( ] self.caption_projection = nn.ModuleList(caption_projection) + @apply_lora_scale("joint_attention_kwargs") def forward( self, hidden_states: torch.Tensor, @@ -520,9 +519,9 @@ def forward( img_ids: torch.Tensor = None, txt_ids: torch.Tensor = None, guidance: torch.Tensor = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, + joint_attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - ) -> Union[torch.FloatTensor, Transformer2DModelOutput]: + ) -> torch.FloatTensor | Transformer2DModelOutput: """ Args: @@ -545,20 +544,7 @@ def forward( If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a `tuple` where the first element is the sample tensor. """ - if joint_attention_kwargs is not None: - joint_attention_kwargs = joint_attention_kwargs.copy() - lora_scale = joint_attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective." - ) hidden_states = self.x_embedder(hidden_states) timestep = timestep.to(hidden_states.dtype) @@ -645,10 +631,6 @@ def forward( hidden_states = self.norm_out(hidden_states, temb) output = self.proj_out(hidden_states) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) diff --git a/src/diffusers/models/transformers/transformer_chroma.py b/src/diffusers/models/transformers/transformer_chroma.py index 2ef3643dafbd..c229d1fefad9 100644 --- a/src/diffusers/models/transformers/transformer_chroma.py +++ b/src/diffusers/models/transformers/transformer_chroma.py @@ -13,7 +13,7 @@ # limitations under the License. -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any import numpy as np import torch @@ -21,7 +21,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FluxTransformer2DLoadersMixin, FromOriginalModelMixin, PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, deprecate, logging from ...utils.import_utils import is_torch_npu_available from ...utils.torch_utils import maybe_allow_in_graph from ..attention import AttentionMixin, FeedForward @@ -36,6 +36,154 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name +class Nerf(nn.Module): + def __init__( + self, + in_channels: int, + num_layers: int, + nerf_hidden_size: int, + transformer_hidden_size: int, + max_freqs: int, + mlp_ratio: int, + eps=1e-6, + ): + super().__init__() + self.nerf_embedder = NerfEmbedder( + in_channels=in_channels, + hidden_size=nerf_hidden_size, + max_freqs=max_freqs, + ) + self.blocks = nn.ModuleList( + [ + NerfGLUBlock( + transformer_hidden_size=transformer_hidden_size, + nerf_hidden_size=nerf_hidden_size, + mlp_ratio=mlp_ratio, + eps=eps, + ) + for _ in range(num_layers) + ] + ) + self.final_layer = NerfFinalLayer( + hidden_size=nerf_hidden_size, + out_channels=in_channels, + eps=eps, + ) + self.transformer_hidden_size = transformer_hidden_size + + def __call__( + self, + pixels, + latents, + patch_size, + num_patches, + ): + batch_size, channels, height, width = pixels.shape + + pixels = nn.functional.unfold(pixels, kernel_size=patch_size, stride=patch_size) + pixels = pixels.transpose(1, 2) + + hidden = latents.reshape(batch_size * num_patches, self.transformer_hidden_size) + pixels = pixels.reshape(batch_size * num_patches, channels, patch_size**2).transpose(1, 2) + + # Get pixel embeddings + latents_dct = self.nerf_embedder(pixels) + + # Pass through blocks + for block in self.blocks: + latents_dct = block(latents_dct, hidden) + + latents_dct = self.final_layer.norm(latents_dct) + + latents_dct = latents_dct.transpose(1, 2).reshape(batch_size, num_patches, -1).transpose(1, 2) + latents_dct = nn.functional.fold( + latents_dct, + output_size=(height, width), + kernel_size=patch_size, + stride=patch_size, + ) + return self.final_layer.conv(latents_dct) + + +class NerfEmbedder(nn.Module): + def __init__( + self, + in_channels: int, + hidden_size: int, + max_freqs: int, + ): + super().__init__() + self.max_freqs = max_freqs + self.hidden_size = hidden_size + self.embedder = nn.Sequential(nn.Linear(in_channels + max_freqs**2, hidden_size)) + + def fetch_pos(self, patch_size, device, dtype) -> torch.Tensor: + pos_x = torch.linspace(0, 1, patch_size, device=device, dtype=dtype) + pos_y = torch.linspace(0, 1, patch_size, device=device, dtype=dtype) + pos_y, pos_x = torch.meshgrid(pos_y, pos_x, indexing="ij") + pos_x = pos_x.reshape(-1, 1, 1) + pos_y = pos_y.reshape(-1, 1, 1) + freqs = torch.linspace(0, self.max_freqs - 1, self.max_freqs, device=device, dtype=dtype) + freqs_x = freqs[None, :, None] + freqs_y = freqs[None, None, :] + coeffs = (1 + freqs_x * freqs_y) ** -1 + dct_x = torch.cos(pos_x * freqs_x * torch.pi) + dct_y = torch.cos(pos_y * freqs_y * torch.pi) + dct = (dct_x * dct_y * coeffs).view(1, -1, self.max_freqs**2) + return dct + + def __call__(self, inputs: torch.Tensor) -> torch.Tensor: + batch, pixels, channels = inputs.shape + patch_size = int(pixels**0.5) + input_dtype = inputs.dtype + with torch.autocast(str(inputs.device), enabled=False): + inputs = inputs.float() + dct = self.fetch_pos(patch_size, inputs.device, torch.float32) + dct = dct.repeat(batch, 1, 1) + inputs = torch.cat((inputs, dct), dim=-1) + inputs = self.embedder.float()(inputs) + return inputs.to(dtype=input_dtype) + + +class NerfGLUBlock(nn.Module): + def __init__(self, transformer_hidden_size: int, nerf_hidden_size: int, mlp_ratio, eps): + super().__init__() + total_params = 3 * nerf_hidden_size**2 * mlp_ratio + self.param_generator = nn.Linear(transformer_hidden_size, total_params) + self.norm = RMSNorm(nerf_hidden_size, eps=eps) + self.mlp_ratio = mlp_ratio + + def forward(self, x: torch.Tensor, s: torch.Tensor) -> torch.Tensor: + batch_size, num_x, hidden_size_x = x.shape + mlp_params = self.param_generator(s) + fc1_gate_params, fc1_value_params, fc2_params = mlp_params.chunk(3, dim=-1) + fc1_gate = fc1_gate_params.view(batch_size, hidden_size_x, hidden_size_x * self.mlp_ratio) + fc1_value = fc1_value_params.view(batch_size, hidden_size_x, hidden_size_x * self.mlp_ratio) + fc2 = fc2_params.view(batch_size, hidden_size_x * self.mlp_ratio, hidden_size_x) + fc1_gate = torch.nn.functional.normalize(fc1_gate, dim=-2) + fc1_value = torch.nn.functional.normalize(fc1_value, dim=-2) + fc2 = torch.nn.functional.normalize(fc2, dim=-2) + res_x = x + x = self.norm(x) + x = torch.bmm(torch.nn.functional.silu(torch.bmm(x, fc1_gate)) * torch.bmm(x, fc1_value), fc2) + return x + res_x + + +class NerfFinalLayer(nn.Module): + def __init__(self, hidden_size: int, out_channels: int, eps): + super().__init__() + self.norm = RMSNorm(hidden_size, eps=eps) + self.conv = nn.Conv2d( + in_channels=hidden_size, + out_channels=out_channels, + kernel_size=3, + padding=1, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(self.norm(x.movedim(1, -1)).movedim(-1, 1)) + + class ChromaAdaLayerNormZeroPruned(nn.Module): r""" Norm layer adaptive layer norm zero (adaLN-Zero). @@ -45,7 +193,7 @@ class ChromaAdaLayerNormZeroPruned(nn.Module): num_embeddings (`int`): The size of the embeddings dictionary. """ - def __init__(self, embedding_dim: int, num_embeddings: Optional[int] = None, norm_type="layer_norm", bias=True): + def __init__(self, embedding_dim: int, num_embeddings: int | None = None, norm_type="layer_norm", bias=True): super().__init__() if num_embeddings is not None: self.emb = CombinedTimestepLabelEmbeddings(num_embeddings, embedding_dim) @@ -64,11 +212,11 @@ def __init__(self, embedding_dim: int, num_embeddings: Optional[int] = None, nor def forward( self, x: torch.Tensor, - timestep: Optional[torch.Tensor] = None, - class_labels: Optional[torch.LongTensor] = None, - hidden_dtype: Optional[torch.dtype] = None, - emb: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + timestep: torch.Tensor | None = None, + class_labels: torch.LongTensor | None = None, + hidden_dtype: torch.dtype | None = None, + emb: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: if self.emb is not None: emb = self.emb(timestep, class_labels, hidden_dtype=hidden_dtype) shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.flatten(1, 2).chunk(6, dim=1) @@ -98,8 +246,8 @@ def __init__(self, embedding_dim: int, norm_type="layer_norm", bias=True): def forward( self, x: torch.Tensor, - emb: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + emb: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: shift_msa, scale_msa, gate_msa = emb.flatten(1, 2).chunk(3, dim=1) x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] return x, gate_msa @@ -243,9 +391,9 @@ def forward( self, hidden_states: torch.Tensor, temb: torch.Tensor, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - attention_mask: Optional[torch.Tensor] = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, ) -> torch.Tensor: residual = hidden_states norm_hidden_states, gate = self.norm(hidden_states, emb=temb) @@ -309,10 +457,10 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - attention_mask: Optional[torch.Tensor] = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: temb_img, temb_txt = temb[:, :6], temb[:, 6:] norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb_img) @@ -399,7 +547,7 @@ class ChromaTransformer2DModel( joint_attention_dim (`int`, defaults to `4096`): The number of dimensions to use for the joint attention (embedding/channel dimension of `encoder_hidden_states`). - axes_dims_rope (`Tuple[int]`, defaults to `(16, 56, 56)`): + axes_dims_rope (`tuple[int]`, defaults to `(16, 56, 56)`): The dimensions to use for the rotary positional embeddings. """ @@ -413,13 +561,13 @@ def __init__( self, patch_size: int = 1, in_channels: int = 64, - out_channels: Optional[int] = None, + out_channels: int | None = None, num_layers: int = 19, num_single_layers: int = 38, attention_head_dim: int = 128, num_attention_heads: int = 24, joint_attention_dim: int = 4096, - axes_dims_rope: Tuple[int, ...] = (16, 56, 56), + axes_dims_rope: tuple[int, ...] = (16, 56, 56), approximator_num_channels: int = 64, approximator_hidden_dim: int = 5120, approximator_layers: int = 5, @@ -473,6 +621,282 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("joint_attention_kwargs") + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor = None, + timestep: torch.LongTensor = None, + img_ids: torch.Tensor = None, + txt_ids: torch.Tensor = None, + attention_mask: torch.Tensor = None, + joint_attention_kwargs: dict[str, Any] | None = None, + controlnet_block_samples=None, + controlnet_single_block_samples=None, + return_dict: bool = True, + controlnet_blocks_repeat: bool = False, + ) -> torch.Tensor | Transformer2DModelOutput: + """ + The [`FluxTransformer2DModel`] forward method. + + Args: + hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`): + Input `hidden_states`. + encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`): + Conditional embeddings (embeddings computed from the input conditions such as prompts) to use. + timestep ( `torch.LongTensor`): + Used to indicate denoising step. + block_controlnet_hidden_states: (`list` of `torch.Tensor`): + A list of tensors that if specified are added to the residuals of transformer blocks. + joint_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain + tuple. + + Returns: + If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a + `tuple` where the first element is the sample tensor. + """ + + hidden_states = self.x_embedder(hidden_states) + + timestep = timestep.to(hidden_states.dtype) * 1000 + + input_vec = self.time_text_embed(timestep) + pooled_temb = self.distilled_guidance_layer(input_vec) + + encoder_hidden_states = self.context_embedder(encoder_hidden_states) + + if txt_ids.ndim == 3: + logger.warning( + "Passing `txt_ids` 3d torch.Tensor is deprecated." + "Please remove the batch dimension and pass it as a 2d torch Tensor" + ) + txt_ids = txt_ids[0] + if img_ids.ndim == 3: + logger.warning( + "Passing `img_ids` 3d torch.Tensor is deprecated." + "Please remove the batch dimension and pass it as a 2d torch Tensor" + ) + img_ids = img_ids[0] + + ids = torch.cat((txt_ids, img_ids), dim=0) + image_rotary_emb = self.pos_embed(ids) + + if joint_attention_kwargs is not None and "ip_adapter_image_embeds" in joint_attention_kwargs: + ip_adapter_image_embeds = joint_attention_kwargs.pop("ip_adapter_image_embeds") + ip_hidden_states = self.encoder_hid_proj(ip_adapter_image_embeds) + joint_attention_kwargs.update({"ip_hidden_states": ip_hidden_states}) + + for index_block, block in enumerate(self.transformer_blocks): + img_offset = 3 * len(self.single_transformer_blocks) + txt_offset = img_offset + 6 * len(self.transformer_blocks) + img_modulation = img_offset + 6 * index_block + text_modulation = txt_offset + 6 * index_block + temb = torch.cat( + ( + pooled_temb[:, img_modulation : img_modulation + 6], + pooled_temb[:, text_modulation : text_modulation + 6], + ), + dim=1, + ) + if torch.is_grad_enabled() and self.gradient_checkpointing: + encoder_hidden_states, hidden_states = self._gradient_checkpointing_func( + block, hidden_states, encoder_hidden_states, temb, image_rotary_emb, attention_mask + ) + + else: + encoder_hidden_states, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + attention_mask=attention_mask, + joint_attention_kwargs=joint_attention_kwargs, + ) + + # controlnet residual + if controlnet_block_samples is not None: + interval_control = len(self.transformer_blocks) / len(controlnet_block_samples) + interval_control = int(np.ceil(interval_control)) + # For Xlabs ControlNet. + if controlnet_blocks_repeat: + hidden_states = ( + hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)] + ) + else: + hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control] + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + for index_block, block in enumerate(self.single_transformer_blocks): + start_idx = 3 * index_block + temb = pooled_temb[:, start_idx : start_idx + 3] + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func( + block, + hidden_states, + temb, + image_rotary_emb, + ) + + else: + hidden_states = block( + hidden_states=hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + attention_mask=attention_mask, + joint_attention_kwargs=joint_attention_kwargs, + ) + + # controlnet residual + if controlnet_single_block_samples is not None: + interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples) + interval_control = int(np.ceil(interval_control)) + hidden_states[:, encoder_hidden_states.shape[1] :, ...] = ( + hidden_states[:, encoder_hidden_states.shape[1] :, ...] + + controlnet_single_block_samples[index_block // interval_control] + ) + + hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...] + + temb = pooled_temb[:, -2:] + hidden_states = self.norm_out(hidden_states, temb) + output = self.proj_out(hidden_states) + + if not return_dict: + return (output,) + + return Transformer2DModelOutput(sample=output) + + +class ChromaRadianceTransformer2DModel( + ModelMixin, + ConfigMixin, + PeftAdapterMixin, + FromOriginalModelMixin, + FluxTransformer2DLoadersMixin, + CacheMixin, + AttentionMixin, +): + """ + The Transformer model introduced in Flux, modified for Chroma Radiance. + + Reference: https://huggingface.co/lodestones/Chroma1-Radiance + + Args: + patch_size (`int`, defaults to `1`): + Patch size to turn the input data into small patches. + in_channels (`int`, defaults to `3`): + The number of channels in the input. + out_channels (`int`, *optional*, defaults to `None`): + The number of channels in the output. If not specified, it defaults to `in_channels`. + num_layers (`int`, defaults to `19`): + The number of layers of dual stream DiT blocks to use. + num_single_layers (`int`, defaults to `38`): + The number of layers of single stream DiT blocks to use. + attention_head_dim (`int`, defaults to `128`): + The number of dimensions to use for each attention head. + num_attention_heads (`int`, defaults to `24`): + The number of attention heads to use. + joint_attention_dim (`int`, defaults to `4096`): + The number of dimensions to use for the joint attention (embedding/channel dimension of + `encoder_hidden_states`). + axes_dims_rope (`Tuple[int]`, defaults to `(16, 56, 56)`): + The dimensions to use for the rotary positional embeddings. + x0 (`bool`, defaults to `True`): + Whether or not to use x0 prediction + """ + + _supports_gradient_checkpointing = True + _no_split_modules = ["ChromaTransformerBlock", "ChromaSingleTransformerBlock"] + _repeated_blocks = ["ChromaTransformerBlock", "ChromaSingleTransformerBlock"] + _skip_layerwise_casting_patterns = ["pos_embed", "norm"] + + @register_to_config + def __init__( + self, + patch_size: int = 16, + in_channels: int = 3, + out_channels: Optional[int] = None, + num_layers: int = 19, + num_single_layers: int = 38, + attention_head_dim: int = 128, + num_attention_heads: int = 24, + joint_attention_dim: int = 4096, + axes_dims_rope: Tuple[int, ...] = (16, 56, 56), + approximator_num_channels: int = 64, + approximator_hidden_dim: int = 5120, + approximator_layers: int = 5, + nerf_layers: int = 4, + nerf_hidden_dim: int = 64, + nerf_max_freqs: int = 8, + nerf_mlp_ratio: int = 4, + x0: bool = True, + ): + super().__init__() + self.out_channels = out_channels or in_channels + self.inner_dim = num_attention_heads * attention_head_dim + + self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_dims_rope) + + self.time_text_embed = ChromaCombinedTimestepTextProjEmbeddings( + num_channels=approximator_num_channels // 4, + out_dim=3 * num_single_layers + 2 * 6 * num_layers + 2, + ) + self.distilled_guidance_layer = ChromaApproximator( + in_dim=approximator_num_channels, + out_dim=self.inner_dim, + hidden_dim=approximator_hidden_dim, + n_layers=approximator_layers, + ) + + self.nerf = Nerf( + in_channels, + nerf_layers, + nerf_hidden_dim, + self.inner_dim, + nerf_max_freqs, + nerf_mlp_ratio, + ) + + self.x_embedder_patch = nn.Conv2d( + in_channels, + self.inner_dim, + kernel_size=patch_size, + stride=patch_size, + bias=True, + ) + + self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim) + + self.transformer_blocks = nn.ModuleList( + [ + ChromaTransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + ) + for _ in range(num_layers) + ] + ) + + self.single_transformer_blocks = nn.ModuleList( + [ + ChromaSingleTransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + ) + for _ in range(num_single_layers) + ] + ) + + self.gradient_checkpointing = False + self.x0 = x0 + def forward( self, hidden_states: torch.Tensor, @@ -511,6 +935,7 @@ def forward( If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a `tuple` where the first element is the sample tensor. """ + pixels = hidden_states.to(self.device) if joint_attention_kwargs is not None: joint_attention_kwargs = joint_attention_kwargs.copy() lora_scale = joint_attention_kwargs.pop("scale", 1.0) @@ -526,7 +951,12 @@ def forward( "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective." ) - hidden_states = self.x_embedder(hidden_states) + if self.config.patch_size == 32: + _, _, height, width = hidden_states.shape + hidden_states = nn.functional.interpolate(hidden_states, size=(height // 2, width // 2), mode="nearest") + hidden_states = self.x_embedder_patch(hidden_states) + num_patches = hidden_states.shape[2] * hidden_states.shape[3] + hidden_states = hidden_states.flatten(2).transpose(1, 2) timestep = timestep.to(hidden_states.dtype) * 1000 @@ -627,9 +1057,12 @@ def forward( hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...] - temb = pooled_temb[:, -2:] - hidden_states = self.norm_out(hidden_states, temb) - output = self.proj_out(hidden_states) + output = self.nerf(pixels, hidden_states, self.config.patch_size, num_patches) + + # using x0 prediction + + if self.x0: + output = (pixels - output) / (timestep / 1000).view(-1, 1, 1, 1) if USE_PEFT_BACKEND: # remove `lora_scale` from each PEFT layer diff --git a/src/diffusers/models/transformers/transformer_chronoedit.py b/src/diffusers/models/transformers/transformer_chronoedit.py index 79828b6464f4..25eb6f87a93a 100644 --- a/src/diffusers/models/transformers/transformer_chronoedit.py +++ b/src/diffusers/models/transformers/transformer_chronoedit.py @@ -13,7 +13,7 @@ # limitations under the License. import math -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -21,7 +21,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, deprecate, logging from ...utils.torch_utils import maybe_allow_in_graph from .._modeling_parallel import ContextParallelInput, ContextParallelOutput from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward @@ -43,7 +43,7 @@ def _get_qkv_projections(attn: "WanAttention", hidden_states: torch.Tensor, enco encoder_hidden_states = hidden_states if attn.fused_projections: - if attn.cross_attention_dim_head is None: + if not attn.is_cross_attention: # In self-attention layers, we can fuse the entire QKV projection into a single linear query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1) else: @@ -82,9 +82,9 @@ def __call__( self, attn: "WanAttention", hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: encoder_hidden_states_img = None if attn.add_k_proj is not None: @@ -188,8 +188,8 @@ def __init__( dim_head: int = 64, eps: float = 1e-5, dropout: float = 0.0, - added_kv_proj_dim: Optional[int] = None, - cross_attention_dim_head: Optional[int] = None, + added_kv_proj_dim: int | None = None, + cross_attention_dim_head: int | None = None, processor=None, is_cross_attention=None, ): @@ -219,7 +219,10 @@ def __init__( self.add_v_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=True) self.norm_added_k = torch.nn.RMSNorm(dim_head * heads, eps=eps) - self.is_cross_attention = cross_attention_dim_head is not None + if is_cross_attention is not None: + self.is_cross_attention = is_cross_attention + else: + self.is_cross_attention = cross_attention_dim_head is not None self.set_processor(processor) @@ -227,7 +230,7 @@ def fuse_projections(self): if getattr(self, "fused_projections", False): return - if self.cross_attention_dim_head is None: + if not self.is_cross_attention: concatenated_weights = torch.cat([self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data]) concatenated_bias = torch.cat([self.to_q.bias.data, self.to_k.bias.data, self.to_v.bias.data]) out_features, in_features = concatenated_weights.shape @@ -275,9 +278,9 @@ def unfuse_projections(self): def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, **kwargs, ) -> torch.Tensor: return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, rotary_emb, **kwargs) @@ -316,8 +319,8 @@ def __init__( time_freq_dim: int, time_proj_dim: int, text_embed_dim: int, - image_embed_dim: Optional[int] = None, - pos_embed_seq_len: Optional[int] = None, + image_embed_dim: int | None = None, + pos_embed_seq_len: int | None = None, ): super().__init__() @@ -335,8 +338,8 @@ def forward( self, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, - encoder_hidden_states_image: Optional[torch.Tensor] = None, - timestep_seq_len: Optional[int] = None, + encoder_hidden_states_image: torch.Tensor | None = None, + timestep_seq_len: int | None = None, ): timestep = self.timesteps_proj(timestep) if timestep_seq_len is not None: @@ -359,7 +362,7 @@ class ChronoEditRotaryPosEmbed(nn.Module): def __init__( self, attention_head_dim: int, - patch_size: Tuple[int, int, int], + patch_size: tuple[int, int, int], max_seq_len: int, theta: float = 10000.0, temporal_skip_len: int = 8, @@ -438,7 +441,7 @@ def __init__( qk_norm: str = "rms_norm_across_heads", cross_attn_norm: bool = False, eps: float = 1e-6, - added_kv_proj_dim: Optional[int] = None, + added_kv_proj_dim: int | None = None, ): super().__init__() @@ -524,7 +527,7 @@ class ChronoEditTransformer3DModel( A Transformer model for video-like data used in the ChronoEdit model. Args: - patch_size (`Tuple[int]`, defaults to `(1, 2, 2)`): + patch_size (`tuple[int]`, defaults to `(1, 2, 2)`): 3D patch dimensions for video embedding (t_patch, h_patch, w_patch). num_attention_heads (`int`, defaults to `40`): Fixed length for text embeddings. @@ -542,7 +545,7 @@ class ChronoEditTransformer3DModel( Intermediate dimension in feed-forward network. num_layers (`int`, defaults to `40`): The number of layers of transformer blocks to use. - window_size (`Tuple[int]`, defaults to `(-1, -1)`): + window_size (`tuple[int]`, defaults to `(-1, -1)`): Window size for local attention (-1 indicates global attention). cross_attn_norm (`bool`, defaults to `True`): Enable cross-attention normalization. @@ -581,7 +584,7 @@ class ChronoEditTransformer3DModel( @register_to_config def __init__( self, - patch_size: Tuple[int] = (1, 2, 2), + patch_size: tuple[int] = (1, 2, 2), num_attention_heads: int = 40, attention_head_dim: int = 128, in_channels: int = 16, @@ -591,12 +594,12 @@ def __init__( ffn_dim: int = 13824, num_layers: int = 40, cross_attn_norm: bool = True, - qk_norm: Optional[str] = "rms_norm_across_heads", + qk_norm: str | None = "rms_norm_across_heads", eps: float = 1e-6, - image_dim: Optional[int] = None, - added_kv_proj_dim: Optional[int] = None, + image_dim: int | None = None, + added_kv_proj_dim: int | None = None, rope_max_seq_len: int = 1024, - pos_embed_seq_len: Optional[int] = None, + pos_embed_seq_len: int | None = None, rope_temporal_skip_len: int = 8, ) -> None: super().__init__() @@ -638,30 +641,16 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, timestep: torch.LongTensor, encoder_hidden_states: torch.Tensor, - encoder_hidden_states_image: Optional[torch.Tensor] = None, + encoder_hidden_states_image: torch.Tensor | None = None, return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) - + attention_kwargs: dict[str, Any] | None = None, + ) -> torch.Tensor | dict[str, torch.Tensor]: batch_size, num_channels, num_frames, height, width = hidden_states.shape p_t, p_h, p_w = self.config.patch_size post_patch_num_frames = num_frames // p_t @@ -729,10 +718,6 @@ def forward( hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) diff --git a/src/diffusers/models/transformers/transformer_cogview3plus.py b/src/diffusers/models/transformers/transformer_cogview3plus.py index e48290fb39d4..ad6a442acbcc 100644 --- a/src/diffusers/models/transformers/transformer_cogview3plus.py +++ b/src/diffusers/models/transformers/transformer_cogview3plus.py @@ -13,8 +13,6 @@ # limitations under the License. -from typing import Tuple, Union - import torch import torch.nn as nn @@ -79,7 +77,7 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, emb: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: text_seq_length = encoder_hidden_states.size(1) # norm & modulate @@ -233,7 +231,7 @@ def forward( target_size: torch.Tensor, crop_coords: torch.Tensor, return_dict: bool = True, - ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: + ) -> tuple[torch.Tensor] | Transformer2DModelOutput: """ The [`CogView3PlusTransformer2DModel`] forward method. diff --git a/src/diffusers/models/transformers/transformer_cogview4.py b/src/diffusers/models/transformers/transformer_cogview4.py index 64e9a538a7c2..308e0e6cccaf 100644 --- a/src/diffusers/models/transformers/transformer_cogview4.py +++ b/src/diffusers/models/transformers/transformer_cogview4.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -20,7 +20,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, logging from ...utils.torch_utils import maybe_allow_in_graph from ..attention import FeedForward from ..attention_processor import Attention @@ -73,7 +73,7 @@ def __init__(self, embedding_dim: int, dim: int) -> None: def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: dtype = hidden_states.dtype norm_hidden_states = self.norm(hidden_states).to(dtype=dtype) norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to(dtype=dtype) @@ -129,9 +129,9 @@ def __call__( attn: Attention, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + attention_mask: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: dtype = encoder_hidden_states.dtype batch_size, text_seq_length, embed_dim = encoder_hidden_states.shape @@ -211,14 +211,12 @@ def __call__( attn: Attention, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, - latent_attn_mask: Optional[torch.Tensor] = None, - text_attn_mask: Optional[torch.Tensor] = None, - batch_flag: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[ - Union[Tuple[torch.Tensor, torch.Tensor], List[Tuple[torch.Tensor, torch.Tensor]]] - ] = None, + latent_attn_mask: torch.Tensor | None = None, + text_attn_mask: torch.Tensor | None = None, + batch_flag: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]] | None = None, **kwargs, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: """ Args: attn (`Attention`): @@ -238,10 +236,10 @@ def __call__( Values from 0 to n-1 indicating which samples belong to the same batch. Samples with the same batch_flag are packed together. Example: [0, 1, 1, 2, 2] means sample 0 forms batch0, samples 1-2 form batch1, and samples 3-4 form batch2. If None, no packing is used. - image_rotary_emb (`Tuple[torch.Tensor, torch.Tensor]` or `list[Tuple[torch.Tensor, torch.Tensor]]`, *optional*): + image_rotary_emb (`tuple[torch.Tensor, torch.Tensor]` or `list[tuple[torch.Tensor, torch.Tensor]]`, *optional*): The rotary embedding for the image part of the input. Returns: - `Tuple[torch.Tensor, torch.Tensor]`: The processed hidden states for both image and text streams. + `tuple[torch.Tensor, torch.Tensor]`: The processed hidden states for both image and text streams. """ # Get dimensions and device info @@ -488,13 +486,11 @@ def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[ - Union[Tuple[torch.Tensor, torch.Tensor], List[Tuple[torch.Tensor, torch.Tensor]]] - ] = None, - attention_mask: Optional[Dict[str, torch.Tensor]] = None, - attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + temb: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]] | None = None, + attention_mask: dict[str, torch.Tensor] | None = None, + attention_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: # 1. Timestep conditioning ( norm_hidden_states, @@ -537,7 +533,7 @@ def forward( class CogView4RotaryPosEmbed(nn.Module): - def __init__(self, dim: int, patch_size: int, rope_axes_dim: Tuple[int, int], theta: float = 10000.0) -> None: + def __init__(self, dim: int, patch_size: int, rope_axes_dim: tuple[int, int], theta: float = 10000.0) -> None: super().__init__() self.dim = dim @@ -545,7 +541,7 @@ def __init__(self, dim: int, patch_size: int, rope_axes_dim: Tuple[int, int], th self.rope_axes_dim = rope_axes_dim self.theta = theta - def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: batch_size, num_channels, height, width = hidden_states.shape height, width = height // self.patch_size, width // self.patch_size @@ -666,7 +662,7 @@ def __init__( condition_dim: int = 256, pos_embed_max_size: int = 128, sample_size: int = 128, - rope_axes_dim: Tuple[int, int] = (256, 256), + rope_axes_dim: tuple[int, int] = (256, 256), ): super().__init__() @@ -703,6 +699,7 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, @@ -711,28 +708,11 @@ def forward( original_size: torch.Tensor, target_size: torch.Tensor, crop_coords: torch.Tensor, - attention_kwargs: Optional[Dict[str, Any]] = None, + attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[ - Union[Tuple[torch.Tensor, torch.Tensor], List[Tuple[torch.Tensor, torch.Tensor]]] - ] = None, - ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) - + attention_mask: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]] | None = None, + ) -> tuple[torch.Tensor] | Transformer2DModelOutput: batch_size, num_channels, height, width = hidden_states.shape # 1. RoPE @@ -779,10 +759,6 @@ def forward( hidden_states = hidden_states.reshape(batch_size, post_patch_height, post_patch_width, -1, p, p) output = hidden_states.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) return Transformer2DModelOutput(sample=output) diff --git a/src/diffusers/models/transformers/transformer_cosmos.py b/src/diffusers/models/transformers/transformer_cosmos.py index 2b0c2667072b..46746a19a678 100644 --- a/src/diffusers/models/transformers/transformer_cosmos.py +++ b/src/diffusers/models/transformers/transformer_cosmos.py @@ -12,17 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Tuple - import numpy as np import torch import torch.nn as nn -import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin from ...utils import is_torchvision_available from ..attention import FeedForward +from ..attention_dispatch import dispatch_attention_fn from ..attention_processor import Attention from ..embeddings import Timesteps from ..modeling_outputs import Transformer2DModelOutput @@ -36,7 +34,7 @@ class CosmosPatchEmbed(nn.Module): def __init__( - self, in_channels: int, out_channels: int, patch_size: Tuple[int, int, int], bias: bool = True + self, in_channels: int, out_channels: int, patch_size: tuple[int, int, int], bias: bool = True ) -> None: super().__init__() self.patch_size = patch_size @@ -94,7 +92,7 @@ def __init__(self, in_features: int, hidden_features: int) -> None: self.linear_2 = nn.Linear(hidden_features, 2 * in_features, bias=False) def forward( - self, hidden_states: torch.Tensor, embedded_timestep: torch.Tensor, temb: Optional[torch.Tensor] = None + self, hidden_states: torch.Tensor, embedded_timestep: torch.Tensor, temb: torch.Tensor | None = None ) -> torch.Tensor: embedded_timestep = self.activation(embedded_timestep) embedded_timestep = self.linear_1(embedded_timestep) @@ -114,7 +112,7 @@ def forward( class CosmosAdaLayerNormZero(nn.Module): - def __init__(self, in_features: int, hidden_features: Optional[int] = None) -> None: + def __init__(self, in_features: int, hidden_features: int | None = None) -> None: super().__init__() self.norm = nn.LayerNorm(in_features, elementwise_affine=False, eps=1e-6) @@ -131,7 +129,7 @@ def forward( self, hidden_states: torch.Tensor, embedded_timestep: torch.Tensor, - temb: Optional[torch.Tensor] = None, + temb: torch.Tensor | None = None, ) -> torch.Tensor: embedded_timestep = self.activation(embedded_timestep) embedded_timestep = self.linear_1(embedded_timestep) @@ -152,16 +150,16 @@ def forward( class CosmosAttnProcessor2_0: def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): + if not hasattr(torch.nn.functional, "scaled_dot_product_attention"): raise ImportError("CosmosAttnProcessor2_0 requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0.") def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: # 1. QKV projections if encoder_hidden_states is None: @@ -191,7 +189,6 @@ def __call__( query_idx = torch.tensor(query.size(3), device=query.device) key_idx = torch.tensor(key.size(3), device=key.device) value_idx = torch.tensor(value.size(3), device=value.device) - else: query_idx = query.size(3) key_idx = key.size(3) @@ -200,18 +197,148 @@ def __call__( value = value.repeat_interleave(query_idx // value_idx, dim=3) # 5. Attention - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False + hidden_states = dispatch_attention_fn( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, ) - hidden_states = hidden_states.transpose(1, 2).flatten(2, 3).type_as(query) - - # 6. Output projection + hidden_states = hidden_states.flatten(2, 3).type_as(query) hidden_states = attn.to_out[0](hidden_states) hidden_states = attn.to_out[1](hidden_states) return hidden_states +class CosmosAttnProcessor2_5: + def __init__(self): + if not hasattr(torch.nn.functional, "scaled_dot_product_attention"): + raise ImportError("CosmosAttnProcessor2_5 requires PyTorch 2.0. Please upgrade PyTorch to 2.0 or newer.") + + def __call__( + self, + attn: Attention, + hidden_states: torch.Tensor, + encoder_hidden_states: tuple[torch.Tensor, torch.Tensor], + attention_mask: tuple[torch.Tensor, torch.Tensor], + image_rotary_emb=None, + ) -> torch.Tensor: + if not isinstance(encoder_hidden_states, tuple): + raise ValueError("Expected encoder_hidden_states as (text_context, img_context) tuple.") + + text_context, img_context = encoder_hidden_states if encoder_hidden_states else (None, None) + text_mask, img_mask = attention_mask if attention_mask else (None, None) + + if text_context is None: + text_context = hidden_states + + query = attn.to_q(hidden_states) + key = attn.to_k(text_context) + value = attn.to_v(text_context) + + query = query.unflatten(2, (attn.heads, -1)).transpose(1, 2) + key = key.unflatten(2, (attn.heads, -1)).transpose(1, 2) + value = value.unflatten(2, (attn.heads, -1)).transpose(1, 2) + + query = attn.norm_q(query) + key = attn.norm_k(key) + + if image_rotary_emb is not None: + from ..embeddings import apply_rotary_emb + + query = apply_rotary_emb(query, image_rotary_emb, use_real=True, use_real_unbind_dim=-2) + key = apply_rotary_emb(key, image_rotary_emb, use_real=True, use_real_unbind_dim=-2) + + if torch.onnx.is_in_onnx_export(): + query_idx = torch.tensor(query.size(3), device=query.device) + key_idx = torch.tensor(key.size(3), device=key.device) + value_idx = torch.tensor(value.size(3), device=value.device) + else: + query_idx = query.size(3) + key_idx = key.size(3) + value_idx = value.size(3) + key = key.repeat_interleave(query_idx // key_idx, dim=3) + value = value.repeat_interleave(query_idx // value_idx, dim=3) + + attn_out = dispatch_attention_fn( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + attn_mask=text_mask, + dropout_p=0.0, + is_causal=False, + ) + attn_out = attn_out.flatten(2, 3).type_as(query) + + if img_context is not None: + q_img = attn.q_img(hidden_states) + k_img = attn.k_img(img_context) + v_img = attn.v_img(img_context) + + batch_size = hidden_states.shape[0] + dim_head = attn.out_dim // attn.heads + + q_img = q_img.view(batch_size, -1, attn.heads, dim_head).transpose(1, 2) + k_img = k_img.view(batch_size, -1, attn.heads, dim_head).transpose(1, 2) + v_img = v_img.view(batch_size, -1, attn.heads, dim_head).transpose(1, 2) + + q_img = attn.q_img_norm(q_img) + k_img = attn.k_img_norm(k_img) + + q_img_idx = q_img.size(3) + k_img_idx = k_img.size(3) + v_img_idx = v_img.size(3) + k_img = k_img.repeat_interleave(q_img_idx // k_img_idx, dim=3) + v_img = v_img.repeat_interleave(q_img_idx // v_img_idx, dim=3) + + img_out = dispatch_attention_fn( + q_img.transpose(1, 2), + k_img.transpose(1, 2), + v_img.transpose(1, 2), + attn_mask=img_mask, + dropout_p=0.0, + is_causal=False, + ) + img_out = img_out.flatten(2, 3).type_as(q_img) + hidden_states = attn_out + img_out + else: + hidden_states = attn_out + + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + return hidden_states + + +class CosmosAttention(Attention): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # add parameters for image q/k/v + inner_dim = self.heads * self.to_q.out_features // self.heads + self.q_img = nn.Linear(self.query_dim, inner_dim, bias=False) + self.k_img = nn.Linear(self.query_dim, inner_dim, bias=False) + self.v_img = nn.Linear(self.query_dim, inner_dim, bias=False) + self.q_img_norm = RMSNorm(self.to_q.out_features // self.heads, eps=1e-6, elementwise_affine=True) + self.k_img_norm = RMSNorm(self.to_k.out_features // self.heads, eps=1e-6, elementwise_affine=True) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None = None, + **cross_attention_kwargs, + ) -> torch.Tensor: + return super().forward( + hidden_states=hidden_states, + # NOTE: type-hint in base class can be ignored + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + + class CosmosTransformerBlock(nn.Module): def __init__( self, @@ -222,12 +349,16 @@ def __init__( adaln_lora_dim: int = 256, qk_norm: str = "rms_norm", out_bias: bool = False, + img_context: bool = False, + before_proj: bool = False, + after_proj: bool = False, ) -> None: super().__init__() hidden_size = num_attention_heads * attention_head_dim self.norm1 = CosmosAdaLayerNormZero(in_features=hidden_size, hidden_features=adaln_lora_dim) + self.img_context = img_context self.attn1 = Attention( query_dim=hidden_size, cross_attention_dim=None, @@ -240,30 +371,56 @@ def __init__( ) self.norm2 = CosmosAdaLayerNormZero(in_features=hidden_size, hidden_features=adaln_lora_dim) - self.attn2 = Attention( - query_dim=hidden_size, - cross_attention_dim=cross_attention_dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - qk_norm=qk_norm, - elementwise_affine=True, - out_bias=out_bias, - processor=CosmosAttnProcessor2_0(), - ) + if img_context: + self.attn2 = CosmosAttention( + query_dim=hidden_size, + cross_attention_dim=cross_attention_dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + qk_norm=qk_norm, + elementwise_affine=True, + out_bias=out_bias, + processor=CosmosAttnProcessor2_5(), + ) + else: + self.attn2 = Attention( + query_dim=hidden_size, + cross_attention_dim=cross_attention_dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + qk_norm=qk_norm, + elementwise_affine=True, + out_bias=out_bias, + processor=CosmosAttnProcessor2_0(), + ) self.norm3 = CosmosAdaLayerNormZero(in_features=hidden_size, hidden_features=adaln_lora_dim) self.ff = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu", bias=out_bias) + # NOTE: zero conv for CosmosControlNet + self.before_proj = None + self.after_proj = None + if before_proj: + self.before_proj = nn.Linear(hidden_size, hidden_size) + if after_proj: + self.after_proj = nn.Linear(hidden_size, hidden_size) + def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None | tuple[torch.Tensor | None, torch.Tensor | None], embedded_timestep: torch.Tensor, - temb: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - extra_pos_emb: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: + temb: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, + extra_pos_emb: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + controlnet_residual: torch.Tensor | None = None, + latents: torch.Tensor | None = None, + block_idx: int | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if self.before_proj is not None: + hidden_states = self.before_proj(hidden_states) + latents + if extra_pos_emb is not None: hidden_states = hidden_states + extra_pos_emb @@ -284,6 +441,16 @@ def forward( ff_output = self.ff(norm_hidden_states) hidden_states = hidden_states + gate * ff_output + if controlnet_residual is not None: + assert self.after_proj is None + # NOTE: this is assumed to be scaled by the controlnet + hidden_states += controlnet_residual + + if self.after_proj is not None: + assert controlnet_residual is None + hs_proj = self.after_proj(hidden_states) + return hidden_states, hs_proj + return hidden_states @@ -291,10 +458,10 @@ class CosmosRotaryPosEmbed(nn.Module): def __init__( self, hidden_size: int, - max_size: Tuple[int, int, int] = (128, 240, 240), - patch_size: Tuple[int, int, int] = (1, 2, 2), + max_size: tuple[int, int, int] = (128, 240, 240), + patch_size: tuple[int, int, int] = (1, 2, 2), base_fps: int = 24, - rope_scale: Tuple[float, float, float] = (2.0, 1.0, 1.0), + rope_scale: tuple[float, float, float] = (2.0, 1.0, 1.0), ) -> None: super().__init__() @@ -310,7 +477,7 @@ def __init__( self.w_ntk_factor = rope_scale[2] ** (self.dim_w / (self.dim_w - 2)) self.t_ntk_factor = rope_scale[0] ** (self.dim_t / (self.dim_t - 2)) - def forward(self, hidden_states: torch.Tensor, fps: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]: + def forward(self, hidden_states: torch.Tensor, fps: int | None = None) -> tuple[torch.Tensor, torch.Tensor]: batch_size, num_channels, num_frames, height, width = hidden_states.shape pe_size = [num_frames // self.patch_size[0], height // self.patch_size[1], width // self.patch_size[2]] device = hidden_states.device @@ -355,8 +522,8 @@ class CosmosLearnablePositionalEmbed(nn.Module): def __init__( self, hidden_size: int, - max_size: Tuple[int, int, int], - patch_size: Tuple[int, int, int], + max_size: tuple[int, int, int], + patch_size: tuple[int, int, int], eps: float = 1e-6, ) -> None: super().__init__() @@ -405,17 +572,28 @@ class CosmosTransformer3DModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): Input dimension of text embeddings from the text encoder. adaln_lora_dim (`int`, defaults to `256`): The hidden dimension of the Adaptive LayerNorm LoRA layer. - max_size (`Tuple[int, int, int]`, defaults to `(128, 240, 240)`): + max_size (`tuple[int, int, int]`, defaults to `(128, 240, 240)`): The maximum size of the input latent tensors in the temporal, height, and width dimensions. - patch_size (`Tuple[int, int, int]`, defaults to `(1, 2, 2)`): + patch_size (`tuple[int, int, int]`, defaults to `(1, 2, 2)`): The patch size to use for patchifying the input latent tensors in the temporal, height, and width dimensions. - rope_scale (`Tuple[float, float, float]`, defaults to `(2.0, 1.0, 1.0)`): + rope_scale (`tuple[float, float, float]`, defaults to `(2.0, 1.0, 1.0)`): The scaling factor to use for RoPE in the temporal, height, and width dimensions. concat_padding_mask (`bool`, defaults to `True`): Whether to concatenate the padding mask to the input latent tensors. extra_pos_embed_type (`str`, *optional*, defaults to `learnable`): The type of extra positional embeddings to use. Can be one of `None` or `learnable`. + controlnet_block_every_n (`int`, *optional*): + Interval between transformer blocks that should receive control residuals (for example, `7` to inject after + every seventh block). Required for Cosmos Transfer2.5. + img_context_dim_in (`int`, *optional*): + The dimension of the input image context feature vector, i.e. it is the D in [B, N, D]. + img_context_num_tokens (`int`): + The number of tokens in the image context feature vector, i.e. it is the N in [B, N, D]. If + `img_context_dim_in` is not provided, then this parameter is ignored. + img_context_dim_out (`int`): + The output dimension of the image context projection layer. If `img_context_dim_in` is not provided, then + this parameter is ignored. """ _supports_gradient_checkpointing = True @@ -434,14 +612,18 @@ def __init__( mlp_ratio: float = 4.0, text_embed_dim: int = 1024, adaln_lora_dim: int = 256, - max_size: Tuple[int, int, int] = (128, 240, 240), - patch_size: Tuple[int, int, int] = (1, 2, 2), - rope_scale: Tuple[float, float, float] = (2.0, 1.0, 1.0), + max_size: tuple[int, int, int] = (128, 240, 240), + patch_size: tuple[int, int, int] = (1, 2, 2), + rope_scale: tuple[float, float, float] = (2.0, 1.0, 1.0), concat_padding_mask: bool = True, - extra_pos_embed_type: Optional[str] = "learnable", + extra_pos_embed_type: str | None = "learnable", use_crossattn_projection: bool = False, crossattn_proj_in_channels: int = 1024, encoder_hidden_states_channels: int = 1024, + controlnet_block_every_n: int | None = None, + img_context_dim_in: int | None = None, + img_context_num_tokens: int = 256, + img_context_dim_out: int = 2048, ) -> None: super().__init__() hidden_size = num_attention_heads * attention_head_dim @@ -477,6 +659,7 @@ def __init__( adaln_lora_dim=adaln_lora_dim, qk_norm="rms_norm", out_bias=False, + img_context=self.config.img_context_dim_in is not None and self.config.img_context_dim_in > 0, ) for _ in range(num_layers) ] @@ -496,17 +679,24 @@ def __init__( self.gradient_checkpointing = False + if self.config.img_context_dim_in: + self.img_context_proj = nn.Sequential( + nn.Linear(self.config.img_context_dim_in, self.config.img_context_dim_out, bias=True), + nn.GELU(), + ) + def forward( self, hidden_states: torch.Tensor, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - fps: Optional[int] = None, - condition_mask: Optional[torch.Tensor] = None, - padding_mask: Optional[torch.Tensor] = None, + block_controlnet_hidden_states: list[torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + fps: int | None = None, + condition_mask: torch.Tensor | None = None, + padding_mask: torch.Tensor | None = None, return_dict: bool = True, - ) -> torch.Tensor: + ) -> tuple[torch.Tensor] | Transformer2DModelOutput: batch_size, num_channels, num_frames, height, width = hidden_states.shape # 1. Concatenate padding mask if needed & prepare attention mask @@ -514,11 +704,11 @@ def forward( hidden_states = torch.cat([hidden_states, condition_mask], dim=1) if self.config.concat_padding_mask: - padding_mask = transforms.functional.resize( + padding_mask_resized = transforms.functional.resize( padding_mask, list(hidden_states.shape[-2:]), interpolation=transforms.InterpolationMode.NEAREST ) hidden_states = torch.cat( - [hidden_states, padding_mask.unsqueeze(2).repeat(batch_size, 1, num_frames, 1, 1)], dim=1 + [hidden_states, padding_mask_resized.unsqueeze(2).repeat(batch_size, 1, num_frames, 1, 1)], dim=1 ) if attention_mask is not None: @@ -554,36 +744,59 @@ def forward( for x in (temb, embedded_timestep) ) # [BT, C] -> [B, T, 1, 1, C] -> [B, T, H, W, C] -> [B, THW, C] else: - assert False + raise ValueError(f"Expected timestep to have shape [B, 1, T, 1, 1] or [T], but got {timestep.shape}") + # 5. Process encoder hidden states + text_context, img_context = ( + encoder_hidden_states if isinstance(encoder_hidden_states, tuple) else (encoder_hidden_states, None) + ) if self.config.use_crossattn_projection: - encoder_hidden_states = self.crossattn_proj(encoder_hidden_states) + text_context = self.crossattn_proj(text_context) + + if img_context is not None and self.config.img_context_dim_in: + img_context = self.img_context_proj(img_context) - # 5. Transformer blocks - for block in self.transformer_blocks: + processed_encoder_hidden_states = ( + (text_context, img_context) if isinstance(encoder_hidden_states, tuple) else text_context + ) + + # 6. Build controlnet block index map + controlnet_block_index_map = {} + if block_controlnet_hidden_states is not None: + n_blocks = len(self.transformer_blocks) + controlnet_block_index_map = { + block_idx: block_controlnet_hidden_states[idx] + for idx, block_idx in list(enumerate(range(0, n_blocks, self.config.controlnet_block_every_n))) + } + + # 7. Transformer blocks + for block_idx, block in enumerate(self.transformer_blocks): + controlnet_residual = controlnet_block_index_map.get(block_idx) if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = self._gradient_checkpointing_func( block, hidden_states, - encoder_hidden_states, + processed_encoder_hidden_states, embedded_timestep, temb, image_rotary_emb, extra_pos_emb, attention_mask, + controlnet_residual, ) else: hidden_states = block( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - embedded_timestep=embedded_timestep, - temb=temb, - image_rotary_emb=image_rotary_emb, - extra_pos_emb=extra_pos_emb, - attention_mask=attention_mask, + hidden_states, + processed_encoder_hidden_states, + embedded_timestep, + temb, + image_rotary_emb, + extra_pos_emb, + attention_mask, + controlnet_residual, ) - # 6. Output norm & projection & unpatchify + # 8. Output norm & projection & unpatchify hidden_states = self.norm_out(hidden_states, embedded_timestep, temb) hidden_states = self.proj_out(hidden_states) hidden_states = hidden_states.unflatten(2, (p_h, p_w, p_t, -1)) diff --git a/src/diffusers/models/transformers/transformer_easyanimate.py b/src/diffusers/models/transformers/transformer_easyanimate.py index 545fa29730db..a665d420c230 100755 --- a/src/diffusers/models/transformers/transformer_easyanimate.py +++ b/src/diffusers/models/transformers/transformer_easyanimate.py @@ -13,8 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Optional, Tuple, Union - import torch import torch.nn.functional as F from torch import nn @@ -58,7 +56,7 @@ def __init__( def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: shift, scale, gate, enc_shift, enc_scale, enc_gate = self.linear(self.silu(temb)).chunk(6, dim=1) hidden_states = self.norm(hidden_states) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) encoder_hidden_states = self.norm(encoder_hidden_states) * (1 + enc_scale.unsqueeze(1)) + enc_shift.unsqueeze( @@ -68,7 +66,7 @@ def forward( class EasyAnimateRotaryPosEmbed(nn.Module): - def __init__(self, patch_size: int, rope_dim: List[int]) -> None: + def __init__(self, patch_size: int, rope_dim: list[int]) -> None: super().__init__() self.patch_size = patch_size @@ -128,8 +126,8 @@ def __call__( attn: Attention, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: if attn.add_q_proj is None and encoder_hidden_states is not None: hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) @@ -220,7 +218,7 @@ def __init__( norm_elementwise_affine: bool = True, norm_eps: float = 1e-6, final_dropout: bool = True, - ff_inner_dim: Optional[int] = None, + ff_inner_dim: int | None = None, ff_bias: bool = True, qk_norm: bool = True, after_norm: bool = False, @@ -280,8 +278,8 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: # 1. Attention norm_hidden_states, norm_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1( hidden_states, encoder_hidden_states, temb @@ -375,9 +373,9 @@ def __init__( self, num_attention_heads: int = 48, attention_head_dim: int = 64, - in_channels: Optional[int] = None, - out_channels: Optional[int] = None, - patch_size: Optional[int] = None, + in_channels: int | None = None, + out_channels: int | None = None, + patch_size: int | None = None, sample_width: int = 90, sample_height: int = 60, activation_fn: str = "gelu-approximate", @@ -464,13 +462,13 @@ def forward( self, hidden_states: torch.Tensor, timestep: torch.Tensor, - timestep_cond: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_hidden_states_t5: Optional[torch.Tensor] = None, - inpaint_latents: Optional[torch.Tensor] = None, - control_latents: Optional[torch.Tensor] = None, + timestep_cond: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_hidden_states_t5: torch.Tensor | None = None, + inpaint_latents: torch.Tensor | None = None, + control_latents: torch.Tensor | None = None, return_dict: bool = True, - ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: + ) -> tuple[torch.Tensor] | Transformer2DModelOutput: batch_size, channels, video_length, height, width = hidden_states.size() p = self.config.patch_size post_patch_height = height // p diff --git a/src/diffusers/models/transformers/transformer_flux.py b/src/diffusers/models/transformers/transformer_flux.py index 16c526f437f2..78a77ebcfea9 100644 --- a/src/diffusers/models/transformers/transformer_flux.py +++ b/src/diffusers/models/transformers/transformer_flux.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import numpy as np import torch @@ -22,7 +22,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FluxTransformer2DLoadersMixin, FromOriginalModelMixin, PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, is_torch_npu_available, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, logging from ...utils.torch_utils import maybe_allow_in_graph from .._modeling_parallel import ContextParallelInput, ContextParallelOutput from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward @@ -85,8 +85,8 @@ def __call__( attn: "FluxAttention", hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections( attn, hidden_states, encoder_hidden_states @@ -130,9 +130,9 @@ def __call__( encoder_hidden_states, hidden_states = hidden_states.split_with_sizes( [encoder_hidden_states.shape[1], hidden_states.shape[1] - encoder_hidden_states.shape[1]], dim=1 ) - hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[0](hidden_states.contiguous()) hidden_states = attn.to_out[1](hidden_states) - encoder_hidden_states = attn.to_add_out(encoder_hidden_states) + encoder_hidden_states = attn.to_add_out(encoder_hidden_states.contiguous()) return hidden_states, encoder_hidden_states else: @@ -185,10 +185,10 @@ def __call__( attn: "FluxAttention", hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - ip_hidden_states: Optional[List[torch.Tensor]] = None, - ip_adapter_masks: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, + ip_hidden_states: list[torch.Tensor] | None = None, + ip_adapter_masks: torch.Tensor | None = None, ) -> torch.Tensor: batch_size = hidden_states.shape[0] @@ -286,12 +286,12 @@ def __init__( dim_head: int = 64, dropout: float = 0.0, bias: bool = False, - added_kv_proj_dim: Optional[int] = None, - added_proj_bias: Optional[bool] = True, + added_kv_proj_dim: int | None = None, + added_proj_bias: bool | None = True, out_bias: bool = True, eps: float = 1e-5, out_dim: int = None, - context_pre_only: Optional[bool] = None, + context_pre_only: bool | None = None, pre_only: bool = False, elementwise_affine: bool = True, processor=None, @@ -336,9 +336,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) @@ -379,9 +379,9 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: text_seq_len = encoder_hidden_states.shape[1] hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) @@ -439,9 +439,9 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( @@ -493,7 +493,7 @@ def forward( class FluxPosEmbed(nn.Module): # modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11 - def __init__(self, theta: int, axes_dim: List[int]): + def __init__(self, theta: int, axes_dim: list[int]): super().__init__() self.theta = theta self.axes_dim = axes_dim @@ -558,7 +558,7 @@ class FluxTransformer2DModel( The number of dimensions to use for the pooled projection. guidance_embeds (`bool`, defaults to `False`): Whether to use guidance embeddings for guidance-distilled variant of the model. - axes_dims_rope (`Tuple[int]`, defaults to `(16, 56, 56)`): + axes_dims_rope (`tuple[int]`, defaults to `(16, 56, 56)`): The dimensions to use for the rotary positional embeddings. """ @@ -581,7 +581,7 @@ def __init__( self, patch_size: int = 1, in_channels: int = 64, - out_channels: Optional[int] = None, + out_channels: int | None = None, num_layers: int = 19, num_single_layers: int = 38, attention_head_dim: int = 128, @@ -589,7 +589,7 @@ def __init__( joint_attention_dim: int = 4096, pooled_projection_dim: int = 768, guidance_embeds: bool = False, - axes_dims_rope: Tuple[int, int, int] = (16, 56, 56), + axes_dims_rope: tuple[int, int, int] = (16, 56, 56), ): super().__init__() self.out_channels = out_channels or in_channels @@ -634,6 +634,7 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("joint_attention_kwargs") def forward( self, hidden_states: torch.Tensor, @@ -643,12 +644,12 @@ def forward( img_ids: torch.Tensor = None, txt_ids: torch.Tensor = None, guidance: torch.Tensor = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, + joint_attention_kwargs: dict[str, Any] | None = None, controlnet_block_samples=None, controlnet_single_block_samples=None, return_dict: bool = True, controlnet_blocks_repeat: bool = False, - ) -> Union[torch.Tensor, Transformer2DModelOutput]: + ) -> torch.Tensor | Transformer2DModelOutput: """ The [`FluxTransformer2DModel`] forward method. @@ -675,20 +676,6 @@ def forward( If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a `tuple` where the first element is the sample tensor. """ - if joint_attention_kwargs is not None: - joint_attention_kwargs = joint_attention_kwargs.copy() - lora_scale = joint_attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective." - ) hidden_states = self.x_embedder(hidden_states) @@ -717,11 +704,7 @@ def forward( img_ids = img_ids[0] ids = torch.cat((txt_ids, img_ids), dim=0) - if is_torch_npu_available(): - freqs_cos, freqs_sin = self.pos_embed(ids.cpu()) - image_rotary_emb = (freqs_cos.npu(), freqs_sin.npu()) - else: - image_rotary_emb = self.pos_embed(ids) + image_rotary_emb = self.pos_embed(ids) if joint_attention_kwargs is not None and "ip_adapter_image_embeds" in joint_attention_kwargs: ip_adapter_image_embeds = joint_attention_kwargs.pop("ip_adapter_image_embeds") @@ -789,10 +772,6 @@ def forward( hidden_states = self.norm_out(hidden_states, temb) output = self.proj_out(hidden_states) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) diff --git a/src/diffusers/models/transformers/transformer_flux2.py b/src/diffusers/models/transformers/transformer_flux2.py index c10bf3ed4f7b..f77498c74fc1 100644 --- a/src/diffusers/models/transformers/transformer_flux2.py +++ b/src/diffusers/models/transformers/transformer_flux2.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -21,7 +21,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FluxTransformer2DLoadersMixin, FromOriginalModelMixin, PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, is_torch_npu_available, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, logging from .._modeling_parallel import ContextParallelInput, ContextParallelOutput from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import dispatch_attention_fn @@ -90,9 +90,9 @@ class Flux2FeedForward(nn.Module): def __init__( self, dim: int, - dim_out: Optional[int] = None, + dim_out: int | None = None, mult: float = 3.0, - inner_dim: Optional[int] = None, + inner_dim: int | None = None, bias: bool = False, ): super().__init__() @@ -125,8 +125,8 @@ def __call__( attn: "Flux2Attention", hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections( attn, hidden_states, encoder_hidden_states @@ -192,8 +192,8 @@ def __init__( dim_head: int = 64, dropout: float = 0.0, bias: bool = False, - added_kv_proj_dim: Optional[int] = None, - added_proj_bias: Optional[bool] = True, + added_kv_proj_dim: int | None = None, + added_proj_bias: bool | None = True, out_bias: bool = True, eps: float = 1e-5, out_dim: int = None, @@ -241,9 +241,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) @@ -268,8 +268,8 @@ def __call__( self, attn: "Flux2ParallelSelfAttention", hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: # Parallel in (QKV + MLP in) projection hidden_states = attn.to_qkv_mlp_proj(hidden_states) @@ -376,8 +376,8 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) @@ -423,20 +423,20 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor], - temb_mod_params: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_hidden_states: torch.Tensor | None, + temb_mod: torch.Tensor, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, split_hidden_states: bool = False, - text_seq_len: Optional[int] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + text_seq_len: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: # If encoder_hidden_states is None, hidden_states is assumed to have encoder_hidden_states already # concatenated if encoder_hidden_states is not None: text_seq_len = encoder_hidden_states.shape[1] hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) - mod_shift, mod_scale, mod_gate = temb_mod_params + mod_shift, mod_scale, mod_gate = Flux2Modulation.split(temb_mod, 1)[0] norm_hidden_states = self.norm(hidden_states) norm_hidden_states = (1 + mod_scale) * norm_hidden_states + mod_shift @@ -498,16 +498,18 @@ def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, - temb_mod_params_img: Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...], - temb_mod_params_txt: Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...], - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + temb_mod_img: torch.Tensor, + temb_mod_txt: torch.Tensor, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: joint_attention_kwargs = joint_attention_kwargs or {} # Modulation parameters shape: [1, 1, self.dim] - (shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = temb_mod_params_img - (c_shift_msa, c_scale_msa, c_gate_msa), (c_shift_mlp, c_scale_mlp, c_gate_mlp) = temb_mod_params_txt + (shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = Flux2Modulation.split(temb_mod_img, 2) + (c_shift_msa, c_scale_msa, c_gate_msa), (c_shift_mlp, c_scale_mlp, c_gate_mlp) = Flux2Modulation.split( + temb_mod_txt, 2 + ) # Img stream norm_hidden_states = self.norm1(hidden_states) @@ -554,7 +556,7 @@ def forward( class Flux2PosEmbed(nn.Module): # modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11 - def __init__(self, theta: int, axes_dim: List[int]): + def __init__(self, theta: int, axes_dim: list[int]): super().__init__() self.theta = theta self.axes_dim = axes_dim @@ -585,7 +587,13 @@ def forward(self, ids: torch.Tensor) -> torch.Tensor: class Flux2TimestepGuidanceEmbeddings(nn.Module): - def __init__(self, in_channels: int = 256, embedding_dim: int = 6144, bias: bool = False): + def __init__( + self, + in_channels: int = 256, + embedding_dim: int = 6144, + bias: bool = False, + guidance_embeds: bool = True, + ): super().__init__() self.time_proj = Timesteps(num_channels=in_channels, flip_sin_to_cos=True, downscale_freq_shift=0) @@ -593,20 +601,24 @@ def __init__(self, in_channels: int = 256, embedding_dim: int = 6144, bias: bool in_channels=in_channels, time_embed_dim=embedding_dim, sample_proj_bias=bias ) - self.guidance_embedder = TimestepEmbedding( - in_channels=in_channels, time_embed_dim=embedding_dim, sample_proj_bias=bias - ) + if guidance_embeds: + self.guidance_embedder = TimestepEmbedding( + in_channels=in_channels, time_embed_dim=embedding_dim, sample_proj_bias=bias + ) + else: + self.guidance_embedder = None def forward(self, timestep: torch.Tensor, guidance: torch.Tensor) -> torch.Tensor: timesteps_proj = self.time_proj(timestep) timesteps_emb = self.timestep_embedder(timesteps_proj.to(timestep.dtype)) # (N, D) - guidance_proj = self.time_proj(guidance) - guidance_emb = self.guidance_embedder(guidance_proj.to(guidance.dtype)) # (N, D) - - time_guidance_emb = timesteps_emb + guidance_emb - - return time_guidance_emb + if guidance is not None and self.guidance_embedder is not None: + guidance_proj = self.time_proj(guidance) + guidance_emb = self.guidance_embedder(guidance_proj.to(guidance.dtype)) # (N, D) + time_guidance_emb = timesteps_emb + guidance_emb + return time_guidance_emb + else: + return timesteps_emb class Flux2Modulation(nn.Module): @@ -617,15 +629,19 @@ def __init__(self, dim: int, mod_param_sets: int = 2, bias: bool = False): self.linear = nn.Linear(dim, dim * 3 * self.mod_param_sets, bias=bias) self.act_fn = nn.SiLU() - def forward(self, temb: torch.Tensor) -> Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]: + def forward(self, temb: torch.Tensor) -> torch.Tensor: mod = self.act_fn(temb) mod = self.linear(mod) + return mod + @staticmethod + # split inside the transformer blocks, to avoid passing tuples into checkpoints https://github.com/huggingface/diffusers/issues/12776 + def split(mod: torch.Tensor, mod_param_sets: int) -> tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]: if mod.ndim == 2: mod = mod.unsqueeze(1) - mod_params = torch.chunk(mod, 3 * self.mod_param_sets, dim=-1) + mod_params = torch.chunk(mod, 3 * mod_param_sets, dim=-1) # Return tuple of 3-tuples of modulation params shift/scale/gate - return tuple(mod_params[3 * i : 3 * (i + 1)] for i in range(self.mod_param_sets)) + return tuple(mod_params[3 * i : 3 * (i + 1)] for i in range(mod_param_sets)) class Flux2Transformer2DModel( @@ -664,7 +680,7 @@ class Flux2Transformer2DModel( The number of dimensions to use for the pooled projection. guidance_embeds (`bool`, defaults to `True`): Whether to use guidance embeddings for guidance-distilled variant of the model. - axes_dims_rope (`Tuple[int]`, defaults to `(32, 32, 32, 32)`): + axes_dims_rope (`tuple[int]`, defaults to `(32, 32, 32, 32)`): The dimensions to use for the rotary positional embeddings. """ @@ -687,7 +703,7 @@ def __init__( self, patch_size: int = 1, in_channels: int = 128, - out_channels: Optional[int] = None, + out_channels: int | None = None, num_layers: int = 8, num_single_layers: int = 48, attention_head_dim: int = 128, @@ -695,9 +711,10 @@ def __init__( joint_attention_dim: int = 15360, timestep_guidance_channels: int = 256, mlp_ratio: float = 3.0, - axes_dims_rope: Tuple[int, ...] = (32, 32, 32, 32), + axes_dims_rope: tuple[int, ...] = (32, 32, 32, 32), rope_theta: int = 2000, eps: float = 1e-6, + guidance_embeds: bool = True, ): super().__init__() self.out_channels = out_channels or in_channels @@ -708,7 +725,10 @@ def __init__( # 2. Combined timestep + guidance embedding self.time_guidance_embed = Flux2TimestepGuidanceEmbeddings( - in_channels=timestep_guidance_channels, embedding_dim=self.inner_dim, bias=False + in_channels=timestep_guidance_channels, + embedding_dim=self.inner_dim, + bias=False, + guidance_embeds=guidance_embeds, ) # 3. Modulation (double stream and single stream blocks share modulation parameters, resp.) @@ -760,6 +780,7 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("joint_attention_kwargs") def forward( self, hidden_states: torch.Tensor, @@ -768,9 +789,9 @@ def forward( img_ids: torch.Tensor = None, txt_ids: torch.Tensor = None, guidance: torch.Tensor = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, + joint_attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - ) -> Union[torch.Tensor, Transformer2DModelOutput]: + ) -> torch.Tensor | Transformer2DModelOutput: """ The [`FluxTransformer2DModel`] forward method. @@ -796,32 +817,20 @@ def forward( `tuple` where the first element is the sample tensor. """ # 0. Handle input arguments - if joint_attention_kwargs is not None: - joint_attention_kwargs = joint_attention_kwargs.copy() - lora_scale = joint_attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective." - ) num_txt_tokens = encoder_hidden_states.shape[1] # 1. Calculate timestep embedding and modulation parameters timestep = timestep.to(hidden_states.dtype) * 1000 - guidance = guidance.to(hidden_states.dtype) * 1000 + + if guidance is not None: + guidance = guidance.to(hidden_states.dtype) * 1000 temb = self.time_guidance_embed(timestep, guidance) double_stream_mod_img = self.double_stream_modulation_img(temb) double_stream_mod_txt = self.double_stream_modulation_txt(temb) - single_stream_mod = self.single_stream_modulation(temb)[0] + single_stream_mod = self.single_stream_modulation(temb) # 2. Input projection for image (hidden_states) and conditioning text (encoder_hidden_states) hidden_states = self.x_embedder(hidden_states) @@ -835,14 +844,8 @@ def forward( if txt_ids.ndim == 3: txt_ids = txt_ids[0] - if is_torch_npu_available(): - freqs_cos_image, freqs_sin_image = self.pos_embed(img_ids.cpu()) - image_rotary_emb = (freqs_cos_image.npu(), freqs_sin_image.npu()) - freqs_cos_text, freqs_sin_text = self.pos_embed(txt_ids.cpu()) - text_rotary_emb = (freqs_cos_text.npu(), freqs_sin_text.npu()) - else: - image_rotary_emb = self.pos_embed(img_ids) - text_rotary_emb = self.pos_embed(txt_ids) + image_rotary_emb = self.pos_embed(img_ids) + text_rotary_emb = self.pos_embed(txt_ids) concat_rotary_emb = ( torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0), torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0), @@ -864,8 +867,8 @@ def forward( encoder_hidden_states, hidden_states = block( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, - temb_mod_params_img=double_stream_mod_img, - temb_mod_params_txt=double_stream_mod_txt, + temb_mod_img=double_stream_mod_img, + temb_mod_txt=double_stream_mod_txt, image_rotary_emb=concat_rotary_emb, joint_attention_kwargs=joint_attention_kwargs, ) @@ -887,7 +890,7 @@ def forward( hidden_states = block( hidden_states=hidden_states, encoder_hidden_states=None, - temb_mod_params=single_stream_mod, + temb_mod=single_stream_mod, image_rotary_emb=concat_rotary_emb, joint_attention_kwargs=joint_attention_kwargs, ) @@ -898,10 +901,6 @@ def forward( hidden_states = self.norm_out(hidden_states, temb) output = self.proj_out(hidden_states) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) diff --git a/src/diffusers/models/transformers/transformer_glm_image.py b/src/diffusers/models/transformers/transformer_glm_image.py new file mode 100644 index 000000000000..8413b24fef45 --- /dev/null +++ b/src/diffusers/models/transformers/transformer_glm_image.py @@ -0,0 +1,668 @@ +# Copyright 2025 The CogView team, Tsinghua University & ZhipuAI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import PeftAdapterMixin +from ...utils import logging +from ...utils.torch_utils import maybe_allow_in_graph +from ..attention import FeedForward +from ..attention_dispatch import dispatch_attention_fn +from ..attention_processor import Attention +from ..cache_utils import CacheMixin +from ..embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin +from ..normalization import LayerNorm, RMSNorm + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class GlmImageCombinedTimestepSizeEmbeddings(nn.Module): + def __init__(self, embedding_dim: int, condition_dim: int, pooled_projection_dim: int, timesteps_dim: int = 256): + super().__init__() + + self.time_proj = Timesteps(num_channels=timesteps_dim, flip_sin_to_cos=True, downscale_freq_shift=0) + self.condition_proj = Timesteps(num_channels=condition_dim, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=timesteps_dim, time_embed_dim=embedding_dim) + self.condition_embedder = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim, act_fn="silu") + + def forward( + self, + timestep: torch.Tensor, + target_size: torch.Tensor, + crop_coords: torch.Tensor, + hidden_dtype: torch.dtype, + ) -> torch.Tensor: + timesteps_proj = self.time_proj(timestep) + + crop_coords_proj = self.condition_proj(crop_coords.flatten()).view(crop_coords.size(0), -1) + target_size_proj = self.condition_proj(target_size.flatten()).view(target_size.size(0), -1) + + # (B, 2 * condition_dim) + condition_proj = torch.cat([crop_coords_proj, target_size_proj], dim=1) + + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (B, embedding_dim) + condition_emb = self.condition_embedder(condition_proj.to(dtype=hidden_dtype)) # (B, embedding_dim) + + conditioning = timesteps_emb + condition_emb + conditioning = F.silu(conditioning) + + return conditioning + + +class GlmImageImageProjector(nn.Module): + def __init__( + self, + in_channels: int = 16, + hidden_size: int = 2560, + patch_size: int = 2, + ): + super().__init__() + self.patch_size = patch_size + + self.proj = nn.Linear(in_channels * patch_size**2, hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, channel, height, width = hidden_states.shape + post_patch_height = height // self.patch_size + post_patch_width = width // self.patch_size + + hidden_states = hidden_states.reshape( + batch_size, channel, post_patch_height, self.patch_size, post_patch_width, self.patch_size + ) + hidden_states = hidden_states.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2) + hidden_states = self.proj(hidden_states) + + return hidden_states + + +class GlmImageAdaLayerNormZero(nn.Module): + def __init__(self, embedding_dim: int, dim: int) -> None: + super().__init__() + + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) + self.norm_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) + self.linear = nn.Linear(embedding_dim, 12 * dim, bias=True) + + def forward( + self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + dtype = hidden_states.dtype + norm_hidden_states = self.norm(hidden_states).to(dtype=dtype) + norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to(dtype=dtype) + + emb = self.linear(temb) + ( + shift_msa, + c_shift_msa, + scale_msa, + c_scale_msa, + gate_msa, + c_gate_msa, + shift_mlp, + c_shift_mlp, + scale_mlp, + c_scale_mlp, + gate_mlp, + c_gate_mlp, + ) = emb.chunk(12, dim=1) + + hidden_states = norm_hidden_states * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1) + encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_msa.unsqueeze(1)) + c_shift_msa.unsqueeze(1) + + return ( + hidden_states, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + encoder_hidden_states, + c_gate_msa, + c_shift_mlp, + c_scale_mlp, + c_gate_mlp, + ) + + +class GlmImageLayerKVCache: + """KV cache for GlmImage model. + Supports per-sample caching for batch processing where each sample may have different condition images. + """ + + def __init__(self): + self.k_caches: list[torch.Tensor | None] = [] + self.v_caches: list[torch.Tensor | None] = [] + self.mode: str | None = None # "write", "read", "skip" + self.current_sample_idx: int = 0 # Current sample index for writing + + def store(self, k: torch.Tensor, v: torch.Tensor): + """Store KV cache for the current sample.""" + # k, v shape: (1, seq_len, num_heads, head_dim) + if len(self.k_caches) <= self.current_sample_idx: + # First time storing for this sample + self.k_caches.append(k) + self.v_caches.append(v) + else: + # Append to existing cache for this sample (multiple condition images) + self.k_caches[self.current_sample_idx] = torch.cat([self.k_caches[self.current_sample_idx], k], dim=1) + self.v_caches[self.current_sample_idx] = torch.cat([self.v_caches[self.current_sample_idx], v], dim=1) + + def get(self, k: torch.Tensor, v: torch.Tensor): + """Get combined KV cache for all samples in the batch. + + Args: + k: Current key tensor, shape (batch_size, seq_len, num_heads, head_dim) + v: Current value tensor, shape (batch_size, seq_len, num_heads, head_dim) + Returns: + Combined key and value tensors with cached values prepended. + """ + batch_size = k.shape[0] + num_cached_samples = len(self.k_caches) + if num_cached_samples == 0: + return k, v + if num_cached_samples == 1: + # Single cache, expand for all batch samples (shared condition images) + k_cache_expanded = self.k_caches[0].expand(batch_size, -1, -1, -1) + v_cache_expanded = self.v_caches[0].expand(batch_size, -1, -1, -1) + elif num_cached_samples == batch_size: + # Per-sample cache, concatenate along batch dimension + k_cache_expanded = torch.cat(self.k_caches, dim=0) + v_cache_expanded = torch.cat(self.v_caches, dim=0) + else: + # Mismatch: try to handle by repeating the caches + # This handles cases like num_images_per_prompt > 1 + repeat_factor = batch_size // num_cached_samples + if batch_size % num_cached_samples == 0: + k_cache_list = [] + v_cache_list = [] + for i in range(num_cached_samples): + k_cache_list.append(self.k_caches[i].expand(repeat_factor, -1, -1, -1)) + v_cache_list.append(self.v_caches[i].expand(repeat_factor, -1, -1, -1)) + k_cache_expanded = torch.cat(k_cache_list, dim=0) + v_cache_expanded = torch.cat(v_cache_list, dim=0) + else: + raise ValueError( + f"Cannot match {num_cached_samples} cached samples to batch size {batch_size}. " + f"Batch size must be a multiple of the number of cached samples." + ) + + k_combined = torch.cat([k_cache_expanded, k], dim=1) + v_combined = torch.cat([v_cache_expanded, v], dim=1) + return k_combined, v_combined + + def clear(self): + self.k_caches = [] + self.v_caches = [] + self.mode = None + self.current_sample_idx = 0 + + def next_sample(self): + """Move to the next sample for writing.""" + self.current_sample_idx += 1 + + +class GlmImageKVCache: + """Container for all layers' KV caches. + Supports per-sample caching for batch processing where each sample may have different condition images. + """ + + def __init__(self, num_layers: int): + self.num_layers = num_layers + self.caches = [GlmImageLayerKVCache() for _ in range(num_layers)] + + def __getitem__(self, layer_idx: int) -> GlmImageLayerKVCache: + return self.caches[layer_idx] + + def set_mode(self, mode: str): + if mode is not None and mode not in ["write", "read", "skip"]: + raise ValueError(f"Invalid mode: {mode}, must be one of 'write', 'read', 'skip'") + for cache in self.caches: + cache.mode = mode + + def next_sample(self): + """Move to the next sample for writing. Call this after processing + all condition images for one batch sample.""" + for cache in self.caches: + cache.next_sample() + + def clear(self): + for cache in self.caches: + cache.clear() + + +class GlmImageAttnProcessor: + """ + Processor for implementing scaled dot-product attention for the GlmImage model. It applies a rotary embedding on + query and key vectors, but does not include spatial normalization. + + The processor supports passing an attention mask for text tokens. The attention mask should have shape (batch_size, + text_seq_length) where 1 indicates a non-padded token and 0 indicates a padded token. + """ + + _attention_backend = None + _parallel_config = None + + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError("GlmImageAttnProcessor requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0.") + + def __call__( + self, + attn: Attention, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + kv_cache: GlmImageLayerKVCache | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + dtype = encoder_hidden_states.dtype + + batch_size, text_seq_length, embed_dim = encoder_hidden_states.shape + batch_size, image_seq_length, embed_dim = hidden_states.shape + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + # 1. QKV projections + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + + query = query.unflatten(2, (attn.heads, -1)) + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + + # 2. QK normalization + if attn.norm_q is not None: + query = attn.norm_q(query).to(dtype=dtype) + if attn.norm_k is not None: + key = attn.norm_k(key).to(dtype=dtype) + + # 3. Rotational positional embeddings applied to latent stream + if image_rotary_emb is not None: + from ..embeddings import apply_rotary_emb + + query[:, text_seq_length:, :, :] = apply_rotary_emb( + query[:, text_seq_length:, :, :], image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2 + ) + key[:, text_seq_length:, :, :] = apply_rotary_emb( + key[:, text_seq_length:, :, :], image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2 + ) + + if kv_cache is not None: + if kv_cache.mode == "write": + kv_cache.store(key, value) + elif kv_cache.mode == "read": + key, value = kv_cache.get(key, value) + elif kv_cache.mode == "skip": + pass + + # 4. Attention + if attention_mask is not None: + text_attn_mask = attention_mask + assert text_attn_mask.dim() == 2, "the shape of text_attn_mask should be (batch_size, text_seq_length)" + text_attn_mask = text_attn_mask.float().to(query.device) + mix_attn_mask = torch.ones((batch_size, text_seq_length + image_seq_length), device=query.device) + mix_attn_mask[:, :text_seq_length] = text_attn_mask + mix_attn_mask = mix_attn_mask.unsqueeze(2) + attn_mask_matrix = mix_attn_mask @ mix_attn_mask.transpose(1, 2) + attention_mask = (attn_mask_matrix > 0).unsqueeze(1).to(query.dtype) + + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(query.dtype) + + # 5. Output projection + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + + encoder_hidden_states, hidden_states = hidden_states.split( + [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1 + ) + return hidden_states, encoder_hidden_states + + +@maybe_allow_in_graph +class GlmImageTransformerBlock(nn.Module): + def __init__( + self, + dim: int = 2560, + num_attention_heads: int = 64, + attention_head_dim: int = 40, + time_embed_dim: int = 512, + ) -> None: + super().__init__() + + # 1. Attention + self.norm1 = GlmImageAdaLayerNormZero(time_embed_dim, dim) + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + out_dim=dim, + bias=True, + qk_norm="layer_norm", + elementwise_affine=False, + eps=1e-5, + processor=GlmImageAttnProcessor(), + ) + + # 2. Feedforward + self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) + self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) + self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate") + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]] | None = None, + attention_mask: dict[str, torch.Tensor] | None = None, + attention_kwargs: dict[str, Any] | None = None, + kv_cache: GlmImageLayerKVCache | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # 1. Timestep conditioning + ( + norm_hidden_states, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + norm_encoder_hidden_states, + c_gate_msa, + c_shift_mlp, + c_scale_mlp, + c_gate_mlp, + ) = self.norm1(hidden_states, encoder_hidden_states, temb) + + # 2. Attention + attention_kwargs = attention_kwargs or {} + + attn_hidden_states, attn_encoder_hidden_states = self.attn1( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + attention_mask=attention_mask, + kv_cache=kv_cache, + **attention_kwargs, + ) + hidden_states = hidden_states + attn_hidden_states * gate_msa.unsqueeze(1) + encoder_hidden_states = encoder_hidden_states + attn_encoder_hidden_states * c_gate_msa.unsqueeze(1) + + # 3. Feedforward + norm_hidden_states = self.norm2(hidden_states) * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1) + norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) * ( + 1 + c_scale_mlp.unsqueeze(1) + ) + c_shift_mlp.unsqueeze(1) + + ff_output = self.ff(norm_hidden_states) + ff_output_context = self.ff(norm_encoder_hidden_states) + hidden_states = hidden_states + ff_output * gate_mlp.unsqueeze(1) + encoder_hidden_states = encoder_hidden_states + ff_output_context * c_gate_mlp.unsqueeze(1) + + return hidden_states, encoder_hidden_states + + +class GlmImageRotaryPosEmbed(nn.Module): + def __init__(self, dim: int, patch_size: int, theta: float = 10000.0) -> None: + super().__init__() + + self.dim = dim + self.patch_size = patch_size + self.theta = theta + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, num_channels, height, width = hidden_states.shape + height, width = height // self.patch_size, width // self.patch_size + + dim_h, dim_w = self.dim // 2, self.dim // 2 + h_inv_freq = 1.0 / ( + self.theta ** (torch.arange(0, dim_h, 2, dtype=torch.float32)[: (dim_h // 2)].float() / dim_h) + ) + w_inv_freq = 1.0 / ( + self.theta ** (torch.arange(0, dim_w, 2, dtype=torch.float32)[: (dim_w // 2)].float() / dim_w) + ) + h_seq = torch.arange(height) + w_seq = torch.arange(width) + freqs_h = torch.outer(h_seq, h_inv_freq) + freqs_w = torch.outer(w_seq, w_inv_freq) + + # Create position matrices for height and width + # [height, 1, dim//4] and [1, width, dim//4] + freqs_h = freqs_h.unsqueeze(1) + freqs_w = freqs_w.unsqueeze(0) + # Broadcast freqs_h and freqs_w to [height, width, dim//4] + freqs_h = freqs_h.expand(height, width, -1) + freqs_w = freqs_w.expand(height, width, -1) + + # Concatenate along last dimension to get [height, width, dim//2] + freqs = torch.cat([freqs_h, freqs_w], dim=-1) + freqs = torch.cat([freqs, freqs], dim=-1) # [height, width, dim] + freqs = freqs.reshape(height * width, -1) + return (freqs.cos(), freqs.sin()) + + +class GlmImageAdaLayerNormContinuous(nn.Module): + """ + GlmImage-only final AdaLN: LN(x) -> Linear(cond) -> chunk -> affine. Matches Megatron: **no activation** before the + Linear on conditioning embedding. + """ + + def __init__( + self, + embedding_dim: int, + conditioning_embedding_dim: int, + elementwise_affine: bool = True, + eps: float = 1e-5, + bias: bool = True, + norm_type: str = "layer_norm", + ): + super().__init__() + self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias) + if norm_type == "layer_norm": + self.norm = LayerNorm(embedding_dim, eps, elementwise_affine, bias) + elif norm_type == "rms_norm": + self.norm = RMSNorm(embedding_dim, eps, elementwise_affine) + else: + raise ValueError(f"unknown norm_type {norm_type}") + + def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor: + # *** NO SiLU here *** + emb = self.linear(conditioning_embedding.to(x.dtype)) + scale, shift = torch.chunk(emb, 2, dim=1) + x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] + return x + + +class GlmImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, CacheMixin): + r""" + Args: + patch_size (`int`, defaults to `2`): + The size of the patches to use in the patch embedding layer. + in_channels (`int`, defaults to `16`): + The number of channels in the input. + num_layers (`int`, defaults to `30`): + The number of layers of Transformer blocks to use. + attention_head_dim (`int`, defaults to `40`): + The number of channels in each head. + num_attention_heads (`int`, defaults to `64`): + The number of heads to use for multi-head attention. + out_channels (`int`, defaults to `16`): + The number of channels in the output. + text_embed_dim (`int`, defaults to `1472`): + Input dimension of text embeddings from the text encoder. + time_embed_dim (`int`, defaults to `512`): + Output dimension of timestep embeddings. + condition_dim (`int`, defaults to `256`): + The embedding dimension of the input SDXL-style resolution conditions (original_size, target_size, + crop_coords). + pos_embed_max_size (`int`, defaults to `128`): + The maximum resolution of the positional embeddings, from which slices of shape `H x W` are taken and added + to input patched latents, where `H` and `W` are the latent height and width respectively. A value of 128 + means that the maximum supported height and width for image generation is `128 * vae_scale_factor * + patch_size => 128 * 8 * 2 => 2048`. + sample_size (`int`, defaults to `128`): + The base resolution of input latents. If height/width is not provided during generation, this value is used + to determine the resolution as `sample_size * vae_scale_factor => 128 * 8 => 1024` + """ + + _supports_gradient_checkpointing = True + _no_split_modules = [ + "GlmImageTransformerBlock", + "GlmImageImageProjector", + "GlmImageImageProjector", + ] + _skip_layerwise_casting_patterns = ["patch_embed", "norm", "proj_out"] + _skip_keys = ["kv_caches"] + + @register_to_config + def __init__( + self, + patch_size: int = 2, + in_channels: int = 16, + out_channels: int = 16, + num_layers: int = 30, + attention_head_dim: int = 40, + num_attention_heads: int = 64, + text_embed_dim: int = 1472, + time_embed_dim: int = 512, + condition_dim: int = 256, + prior_vq_quantizer_codebook_size: int = 16384, + ): + super().__init__() + + # GlmImage uses 2 additional SDXL-like conditions - target_size, crop_coords + # Each of these are sincos embeddings of shape 2 * condition_dim + pooled_projection_dim = 2 * 2 * condition_dim + inner_dim = num_attention_heads * attention_head_dim + out_channels = out_channels + + # 1. RoPE + self.rope = GlmImageRotaryPosEmbed(attention_head_dim, patch_size, theta=10000.0) + + # 2. Patch & Text-timestep embedding + self.image_projector = GlmImageImageProjector(in_channels, inner_dim, patch_size) + self.glyph_projector = FeedForward(text_embed_dim, inner_dim, inner_dim=inner_dim, activation_fn="gelu") + self.prior_token_embedding = nn.Embedding(prior_vq_quantizer_codebook_size, inner_dim) + self.prior_projector = FeedForward(inner_dim, inner_dim, inner_dim=inner_dim, activation_fn="linear-silu") + + self.time_condition_embed = GlmImageCombinedTimestepSizeEmbeddings( + embedding_dim=time_embed_dim, + condition_dim=condition_dim, + pooled_projection_dim=pooled_projection_dim, + timesteps_dim=time_embed_dim, + ) + + # 3. Transformer blocks + self.transformer_blocks = nn.ModuleList( + [ + GlmImageTransformerBlock(inner_dim, num_attention_heads, attention_head_dim, time_embed_dim) + for _ in range(num_layers) + ] + ) + + # 4. Output projection + self.norm_out = GlmImageAdaLayerNormContinuous(inner_dim, time_embed_dim, elementwise_affine=False) + self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels, bias=True) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + prior_token_id: torch.Tensor, + prior_token_drop: torch.Tensor, + timestep: torch.LongTensor, + target_size: torch.Tensor, + crop_coords: torch.Tensor, + attention_kwargs: dict[str, Any] | None = None, + return_dict: bool = True, + attention_mask: torch.Tensor | None = None, + kv_caches: GlmImageKVCache | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]] | None = None, + ) -> tuple[torch.Tensor] | Transformer2DModelOutput: + batch_size, num_channels, height, width = hidden_states.shape + + # 1. RoPE + if image_rotary_emb is None: + image_rotary_emb = self.rope(hidden_states) + + # 2. Patch & Timestep embeddings + p = self.config.patch_size + post_patch_height = height // p + post_patch_width = width // p + + hidden_states = self.image_projector(hidden_states) + encoder_hidden_states = self.glyph_projector(encoder_hidden_states) + prior_embedding = self.prior_token_embedding(prior_token_id) + prior_embedding[prior_token_drop] *= 0.0 + prior_hidden_states = self.prior_projector(prior_embedding) + + hidden_states = hidden_states + prior_hidden_states + + temb = self.time_condition_embed(timestep, target_size, crop_coords, hidden_states.dtype) + + # 3. Transformer blocks + for idx, block in enumerate(self.transformer_blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states, encoder_hidden_states = self._gradient_checkpointing_func( + block, + hidden_states, + encoder_hidden_states, + temb, + image_rotary_emb, + attention_mask, + attention_kwargs, + kv_caches[idx] if kv_caches is not None else None, + ) + else: + hidden_states, encoder_hidden_states = block( + hidden_states, + encoder_hidden_states, + temb, + image_rotary_emb, + attention_mask, + attention_kwargs, + kv_cache=kv_caches[idx] if kv_caches is not None else None, + ) + + # 4. Output norm & projection + hidden_states = self.norm_out(hidden_states, temb) + hidden_states = self.proj_out(hidden_states) + + # 5. Unpatchify + hidden_states = hidden_states.reshape(batch_size, post_patch_height, post_patch_width, -1, p, p) + + # Rearrange tensor from (B, H_p, W_p, C, p, p) to (B, C, H_p * p, W_p * p) + output = hidden_states.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) + + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) diff --git a/src/diffusers/models/transformers/transformer_helios.py b/src/diffusers/models/transformers/transformer_helios.py new file mode 100644 index 000000000000..6d81f8a13af7 --- /dev/null +++ b/src/diffusers/models/transformers/transformer_helios.py @@ -0,0 +1,821 @@ +# Copyright 2025 The Helios Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from functools import lru_cache +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin +from ...utils import apply_lora_scale, logging +from ...utils.torch_utils import maybe_allow_in_graph +from .._modeling_parallel import ContextParallelInput, ContextParallelOutput +from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward +from ..attention_dispatch import dispatch_attention_fn +from ..cache_utils import CacheMixin +from ..embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin +from ..normalization import FP32LayerNorm + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def pad_for_3d_conv(x, kernel_size): + b, c, t, h, w = x.shape + pt, ph, pw = kernel_size + pad_t = (pt - (t % pt)) % pt + pad_h = (ph - (h % ph)) % ph + pad_w = (pw - (w % pw)) % pw + return torch.nn.functional.pad(x, (0, pad_w, 0, pad_h, 0, pad_t), mode="replicate") + + +def center_down_sample_3d(x, kernel_size): + return torch.nn.functional.avg_pool3d(x, kernel_size, stride=kernel_size) + + +def apply_rotary_emb_transposed( + hidden_states: torch.Tensor, + freqs_cis: torch.Tensor, +): + x_1, x_2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) + cos, sin = freqs_cis.unsqueeze(-2).chunk(2, dim=-1) + out = torch.empty_like(hidden_states) + out[..., 0::2] = x_1 * cos[..., 0::2] - x_2 * sin[..., 1::2] + out[..., 1::2] = x_1 * sin[..., 1::2] + x_2 * cos[..., 0::2] + return out.type_as(hidden_states) + + +def _get_qkv_projections(attn: "HeliosAttention", hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor): + # encoder_hidden_states is only passed for cross-attention + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + + if attn.fused_projections: + if not attn.is_cross_attention: + # In self-attention layers, we can fuse the entire QKV projection into a single linear + query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1) + else: + # In cross-attention layers, we can only fuse the KV projections into a single linear + query = attn.to_q(hidden_states) + key, value = attn.to_kv(encoder_hidden_states).chunk(2, dim=-1) + else: + query = attn.to_q(hidden_states) + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + return query, key, value + + +class HeliosOutputNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = False): + super().__init__() + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + self.norm = FP32LayerNorm(dim, eps, elementwise_affine=False) + + def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor, original_context_length: int): + temb = temb[:, -original_context_length:, :] + shift, scale = (self.scale_shift_table.unsqueeze(0).to(temb.device) + temb.unsqueeze(2)).chunk(2, dim=2) + shift, scale = shift.squeeze(2).to(hidden_states.device), scale.squeeze(2).to(hidden_states.device) + hidden_states = hidden_states[:, -original_context_length:, :] + hidden_states = (self.norm(hidden_states.float()) * (1 + scale) + shift).type_as(hidden_states) + return hidden_states + + +class HeliosAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError( + "HeliosAttnProcessor requires PyTorch 2.0. To use it, please upgrade PyTorch to version 2.0 or higher." + ) + + def __call__( + self, + attn: "HeliosAttention", + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + original_context_length: int = None, + ) -> torch.Tensor: + query, key, value = _get_qkv_projections(attn, hidden_states, encoder_hidden_states) + + query = attn.norm_q(query) + key = attn.norm_k(key) + + query = query.unflatten(2, (attn.heads, -1)) + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + + if rotary_emb is not None: + query = apply_rotary_emb_transposed(query, rotary_emb) + key = apply_rotary_emb_transposed(key, rotary_emb) + + if not attn.is_cross_attention and attn.is_amplify_history: + history_seq_len = hidden_states.shape[1] - original_context_length + + if history_seq_len > 0: + scale_key = 1.0 + torch.sigmoid(attn.history_key_scale) * (attn.max_scale - 1.0) + if attn.history_scale_mode == "per_head": + scale_key = scale_key.view(1, 1, -1, 1) + key = torch.cat([key[:, :history_seq_len] * scale_key, key[:, history_seq_len:]], dim=1) + + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + # Reference: https://github.com/huggingface/diffusers/pull/12909 + parallel_config=(self._parallel_config if encoder_hidden_states is None else None), + ) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.type_as(query) + + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + return hidden_states + + +class HeliosAttention(torch.nn.Module, AttentionModuleMixin): + _default_processor_cls = HeliosAttnProcessor + _available_processors = [HeliosAttnProcessor] + + def __init__( + self, + dim: int, + heads: int = 8, + dim_head: int = 64, + eps: float = 1e-5, + dropout: float = 0.0, + added_kv_proj_dim: int | None = None, + cross_attention_dim_head: int | None = None, + processor=None, + is_cross_attention=None, + is_amplify_history=False, + history_scale_mode="per_head", # [scalar, per_head] + ): + super().__init__() + + self.inner_dim = dim_head * heads + self.heads = heads + self.added_kv_proj_dim = added_kv_proj_dim + self.cross_attention_dim_head = cross_attention_dim_head + self.kv_inner_dim = self.inner_dim if cross_attention_dim_head is None else cross_attention_dim_head * heads + + self.to_q = torch.nn.Linear(dim, self.inner_dim, bias=True) + self.to_k = torch.nn.Linear(dim, self.kv_inner_dim, bias=True) + self.to_v = torch.nn.Linear(dim, self.kv_inner_dim, bias=True) + self.to_out = torch.nn.ModuleList( + [ + torch.nn.Linear(self.inner_dim, dim, bias=True), + torch.nn.Dropout(dropout), + ] + ) + self.norm_q = torch.nn.RMSNorm(dim_head * heads, eps=eps, elementwise_affine=True) + self.norm_k = torch.nn.RMSNorm(dim_head * heads, eps=eps, elementwise_affine=True) + + self.add_k_proj = self.add_v_proj = None + if added_kv_proj_dim is not None: + self.add_k_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=True) + self.add_v_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=True) + self.norm_added_k = torch.nn.RMSNorm(dim_head * heads, eps=eps) + + if is_cross_attention is not None: + self.is_cross_attention = is_cross_attention + else: + self.is_cross_attention = cross_attention_dim_head is not None + + self.set_processor(processor) + + self.is_amplify_history = is_amplify_history + if is_amplify_history: + if history_scale_mode == "scalar": + self.history_key_scale = nn.Parameter(torch.ones(1)) + elif history_scale_mode == "per_head": + self.history_key_scale = nn.Parameter(torch.ones(heads)) + else: + raise ValueError(f"Unknown history_scale_mode: {history_scale_mode}") + self.history_scale_mode = history_scale_mode + self.max_scale = 10.0 + + def fuse_projections(self): + if getattr(self, "fused_projections", False): + return + + if not self.is_cross_attention: + concatenated_weights = torch.cat([self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data]) + concatenated_bias = torch.cat([self.to_q.bias.data, self.to_k.bias.data, self.to_v.bias.data]) + out_features, in_features = concatenated_weights.shape + with torch.device("meta"): + self.to_qkv = nn.Linear(in_features, out_features, bias=True) + self.to_qkv.load_state_dict( + {"weight": concatenated_weights, "bias": concatenated_bias}, strict=True, assign=True + ) + else: + concatenated_weights = torch.cat([self.to_k.weight.data, self.to_v.weight.data]) + concatenated_bias = torch.cat([self.to_k.bias.data, self.to_v.bias.data]) + out_features, in_features = concatenated_weights.shape + with torch.device("meta"): + self.to_kv = nn.Linear(in_features, out_features, bias=True) + self.to_kv.load_state_dict( + {"weight": concatenated_weights, "bias": concatenated_bias}, strict=True, assign=True + ) + + if self.added_kv_proj_dim is not None: + concatenated_weights = torch.cat([self.add_k_proj.weight.data, self.add_v_proj.weight.data]) + concatenated_bias = torch.cat([self.add_k_proj.bias.data, self.add_v_proj.bias.data]) + out_features, in_features = concatenated_weights.shape + with torch.device("meta"): + self.to_added_kv = nn.Linear(in_features, out_features, bias=True) + self.to_added_kv.load_state_dict( + {"weight": concatenated_weights, "bias": concatenated_bias}, strict=True, assign=True + ) + + self.fused_projections = True + + @torch.no_grad() + def unfuse_projections(self): + if not getattr(self, "fused_projections", False): + return + + if hasattr(self, "to_qkv"): + delattr(self, "to_qkv") + if hasattr(self, "to_kv"): + delattr(self, "to_kv") + if hasattr(self, "to_added_kv"): + delattr(self, "to_added_kv") + + self.fused_projections = False + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + original_context_length: int = None, + **kwargs, + ) -> torch.Tensor: + return self.processor( + self, + hidden_states, + encoder_hidden_states, + attention_mask, + rotary_emb, + original_context_length, + **kwargs, + ) + + +class HeliosTimeTextEmbedding(nn.Module): + def __init__( + self, + dim: int, + time_freq_dim: int, + time_proj_dim: int, + text_embed_dim: int, + ): + super().__init__() + + self.timesteps_proj = Timesteps(num_channels=time_freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0) + self.time_embedder = TimestepEmbedding(in_channels=time_freq_dim, time_embed_dim=dim) + self.act_fn = nn.SiLU() + self.time_proj = nn.Linear(dim, time_proj_dim) + self.text_embedder = PixArtAlphaTextProjection(text_embed_dim, dim, act_fn="gelu_tanh") + + def forward( + self, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + is_return_encoder_hidden_states: bool = True, + ): + timestep = self.timesteps_proj(timestep) + + time_embedder_dtype = next(iter(self.time_embedder.parameters())).dtype + if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8: + timestep = timestep.to(time_embedder_dtype) + temb = self.time_embedder(timestep).type_as(encoder_hidden_states) + timestep_proj = self.time_proj(self.act_fn(temb)) + + if encoder_hidden_states is not None and is_return_encoder_hidden_states: + encoder_hidden_states = self.text_embedder(encoder_hidden_states) + + return temb, timestep_proj, encoder_hidden_states + + +class HeliosRotaryPosEmbed(nn.Module): + def __init__(self, rope_dim, theta): + super().__init__() + self.DT, self.DY, self.DX = rope_dim + self.theta = theta + self.register_buffer("freqs_base_t", self._get_freqs_base(self.DT), persistent=False) + self.register_buffer("freqs_base_y", self._get_freqs_base(self.DY), persistent=False) + self.register_buffer("freqs_base_x", self._get_freqs_base(self.DX), persistent=False) + + def _get_freqs_base(self, dim): + return 1.0 / (self.theta ** (torch.arange(0, dim, 2, dtype=torch.float32)[: (dim // 2)] / dim)) + + @torch.no_grad() + def get_frequency_batched(self, freqs_base, pos): + freqs = torch.einsum("d,bthw->dbthw", freqs_base, pos) + freqs = freqs.repeat_interleave(2, dim=0) + return freqs.cos(), freqs.sin() + + @torch.no_grad() + @lru_cache(maxsize=32) + def _get_spatial_meshgrid(self, height, width, device_str): + device = torch.device(device_str) + grid_y_coords = torch.arange(height, device=device, dtype=torch.float32) + grid_x_coords = torch.arange(width, device=device, dtype=torch.float32) + grid_y, grid_x = torch.meshgrid(grid_y_coords, grid_x_coords, indexing="ij") + return grid_y, grid_x + + @torch.no_grad() + def forward(self, frame_indices, height, width, device): + batch_size = frame_indices.shape[0] + num_frames = frame_indices.shape[1] + + frame_indices = frame_indices.to(device=device, dtype=torch.float32) + grid_y, grid_x = self._get_spatial_meshgrid(height, width, str(device)) + + grid_t = frame_indices[:, :, None, None].expand(batch_size, num_frames, height, width) + grid_y_batch = grid_y[None, None, :, :].expand(batch_size, num_frames, -1, -1) + grid_x_batch = grid_x[None, None, :, :].expand(batch_size, num_frames, -1, -1) + + freqs_cos_t, freqs_sin_t = self.get_frequency_batched(self.freqs_base_t, grid_t) + freqs_cos_y, freqs_sin_y = self.get_frequency_batched(self.freqs_base_y, grid_y_batch) + freqs_cos_x, freqs_sin_x = self.get_frequency_batched(self.freqs_base_x, grid_x_batch) + + result = torch.cat([freqs_cos_t, freqs_cos_y, freqs_cos_x, freqs_sin_t, freqs_sin_y, freqs_sin_x], dim=0) + + return result.permute(1, 0, 2, 3, 4) + + +@maybe_allow_in_graph +class HeliosTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + ffn_dim: int, + num_heads: int, + qk_norm: str = "rms_norm_across_heads", + cross_attn_norm: bool = False, + eps: float = 1e-6, + added_kv_proj_dim: int | None = None, + guidance_cross_attn: bool = False, + is_amplify_history: bool = False, + history_scale_mode: str = "per_head", # [scalar, per_head] + ): + super().__init__() + + # 1. Self-attention + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.attn1 = HeliosAttention( + dim=dim, + heads=num_heads, + dim_head=dim // num_heads, + eps=eps, + cross_attention_dim_head=None, + processor=HeliosAttnProcessor(), + is_amplify_history=is_amplify_history, + history_scale_mode=history_scale_mode, + ) + + # 2. Cross-attention + self.attn2 = HeliosAttention( + dim=dim, + heads=num_heads, + dim_head=dim // num_heads, + eps=eps, + added_kv_proj_dim=added_kv_proj_dim, + cross_attention_dim_head=dim // num_heads, + processor=HeliosAttnProcessor(), + ) + self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() + + # 3. Feed-forward + self.ffn = FeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate") + self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False) + + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + # 4. Guidance cross-attention + self.guidance_cross_attn = guidance_cross_attn + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + rotary_emb: torch.Tensor, + original_context_length: int = None, + ) -> torch.Tensor: + if temb.ndim == 4: + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( + self.scale_shift_table.unsqueeze(0) + temb.float() + ).chunk(6, dim=2) + # batch_size, seq_len, 1, inner_dim + shift_msa = shift_msa.squeeze(2) + scale_msa = scale_msa.squeeze(2) + gate_msa = gate_msa.squeeze(2) + c_shift_msa = c_shift_msa.squeeze(2) + c_scale_msa = c_scale_msa.squeeze(2) + c_gate_msa = c_gate_msa.squeeze(2) + else: + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( + self.scale_shift_table + temb.float() + ).chunk(6, dim=1) + + # 1. Self-attention + norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).type_as(hidden_states) + attn_output = self.attn1( + norm_hidden_states, + None, + None, + rotary_emb, + original_context_length, + ) + hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(hidden_states) + + # 2. Cross-attention + if self.guidance_cross_attn: + history_seq_len = hidden_states.shape[1] - original_context_length + + history_hidden_states, hidden_states = torch.split( + hidden_states, [history_seq_len, original_context_length], dim=1 + ) + norm_hidden_states = self.norm2(hidden_states.float()).type_as(hidden_states) + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states, + None, + None, + original_context_length, + ) + hidden_states = hidden_states + attn_output + hidden_states = torch.cat([history_hidden_states, hidden_states], dim=1) + else: + norm_hidden_states = self.norm2(hidden_states.float()).type_as(hidden_states) + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states, + None, + None, + original_context_length, + ) + hidden_states = hidden_states + attn_output + + # 3. Feed-forward + norm_hidden_states = (self.norm3(hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa).type_as( + hidden_states + ) + ff_output = self.ffn(norm_hidden_states) + hidden_states = (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states) + + return hidden_states + + +class HeliosTransformer3DModel( + ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin, AttentionMixin +): + r""" + A Transformer model for video-like data used in the Helios model. + + Args: + patch_size (`tuple[int]`, defaults to `(1, 2, 2)`): + 3D patch dimensions for video embedding (t_patch, h_patch, w_patch). + num_attention_heads (`int`, defaults to `40`): + Fixed length for text embeddings. + attention_head_dim (`int`, defaults to `128`): + The number of channels in each head. + in_channels (`int`, defaults to `16`): + The number of channels in the input. + out_channels (`int`, defaults to `16`): + The number of channels in the output. + text_dim (`int`, defaults to `512`): + Input dimension for text embeddings. + freq_dim (`int`, defaults to `256`): + Dimension for sinusoidal time embeddings. + ffn_dim (`int`, defaults to `13824`): + Intermediate dimension in feed-forward network. + num_layers (`int`, defaults to `40`): + The number of layers of transformer blocks to use. + window_size (`tuple[int]`, defaults to `(-1, -1)`): + Window size for local attention (-1 indicates global attention). + cross_attn_norm (`bool`, defaults to `True`): + Enable cross-attention normalization. + qk_norm (`bool`, defaults to `True`): + Enable query/key normalization. + eps (`float`, defaults to `1e-6`): + Epsilon value for normalization layers. + add_img_emb (`bool`, defaults to `False`): + Whether to use img_emb. + added_kv_proj_dim (`int`, *optional*, defaults to `None`): + The number of channels to use for the added key and value projections. If `None`, no projection is used. + """ + + _supports_gradient_checkpointing = True + _skip_layerwise_casting_patterns = [ + "patch_embedding", + "patch_short", + "patch_mid", + "patch_long", + "condition_embedder", + "norm", + ] + _no_split_modules = ["HeliosTransformerBlock", "HeliosOutputNorm"] + _keep_in_fp32_modules = [ + "time_embedder", + "scale_shift_table", + "norm1", + "norm2", + "norm3", + "history_key_scale", + ] + _keys_to_ignore_on_load_unexpected = ["norm_added_q"] + _repeated_blocks = ["HeliosTransformerBlock"] + _cp_plan = { + # Input split at attn level and ffn level. + "blocks.*.attn1": { + "hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), + "rotary_emb": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), + }, + "blocks.*.attn2": { + "hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), + }, + "blocks.*.ffn": { + "hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), + }, + # Output gather at attn level and ffn level. + **{f"blocks.{i}.attn1": ContextParallelOutput(gather_dim=1, expected_dims=3) for i in range(40)}, + **{f"blocks.{i}.attn2": ContextParallelOutput(gather_dim=1, expected_dims=3) for i in range(40)}, + **{f"blocks.{i}.ffn": ContextParallelOutput(gather_dim=1, expected_dims=3) for i in range(40)}, + } + + @register_to_config + def __init__( + self, + patch_size: tuple[int, ...] = (1, 2, 2), + num_attention_heads: int = 40, + attention_head_dim: int = 128, + in_channels: int = 16, + out_channels: int = 16, + text_dim: int = 4096, + freq_dim: int = 256, + ffn_dim: int = 13824, + num_layers: int = 40, + cross_attn_norm: bool = True, + qk_norm: str | None = "rms_norm_across_heads", + eps: float = 1e-6, + added_kv_proj_dim: int | None = None, + rope_dim: tuple[int, ...] = (44, 42, 42), + rope_theta: float = 10000.0, + guidance_cross_attn: bool = True, + zero_history_timestep: bool = True, + has_multi_term_memory_patch: bool = True, + is_amplify_history: bool = False, + history_scale_mode: str = "per_head", # [scalar, per_head] + ) -> None: + super().__init__() + + inner_dim = num_attention_heads * attention_head_dim + out_channels = out_channels or in_channels + + # 1. Patch & position embedding + self.rope = HeliosRotaryPosEmbed(rope_dim=rope_dim, theta=rope_theta) + self.patch_embedding = nn.Conv3d(in_channels, inner_dim, kernel_size=patch_size, stride=patch_size) + + # 2. Initial Multi Term Memory Patch + self.zero_history_timestep = zero_history_timestep + if has_multi_term_memory_patch: + self.patch_short = nn.Conv3d(in_channels, inner_dim, kernel_size=patch_size, stride=patch_size) + self.patch_mid = nn.Conv3d( + in_channels, + inner_dim, + kernel_size=tuple(2 * p for p in patch_size), + stride=tuple(2 * p for p in patch_size), + ) + self.patch_long = nn.Conv3d( + in_channels, + inner_dim, + kernel_size=tuple(4 * p for p in patch_size), + stride=tuple(4 * p for p in patch_size), + ) + + # 3. Condition embeddings + self.condition_embedder = HeliosTimeTextEmbedding( + dim=inner_dim, + time_freq_dim=freq_dim, + time_proj_dim=inner_dim * 6, + text_embed_dim=text_dim, + ) + + # 4. Transformer blocks + self.blocks = nn.ModuleList( + [ + HeliosTransformerBlock( + inner_dim, + ffn_dim, + num_attention_heads, + qk_norm, + cross_attn_norm, + eps, + added_kv_proj_dim, + guidance_cross_attn=guidance_cross_attn, + is_amplify_history=is_amplify_history, + history_scale_mode=history_scale_mode, + ) + for _ in range(num_layers) + ] + ) + + # 5. Output norm & projection + self.norm_out = HeliosOutputNorm(inner_dim, eps, elementwise_affine=False) + self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size)) + + self.gradient_checkpointing = False + + @apply_lora_scale("attention_kwargs") + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.LongTensor, + encoder_hidden_states: torch.Tensor, + # ------------ Stage 1 ------------ + indices_hidden_states=None, + indices_latents_history_short=None, + indices_latents_history_mid=None, + indices_latents_history_long=None, + latents_history_short=None, + latents_history_mid=None, + latents_history_long=None, + return_dict: bool = True, + attention_kwargs: dict[str, Any] | None = None, + ) -> torch.Tensor | dict[str, torch.Tensor]: + # 1. Input + batch_size = hidden_states.shape[0] + p_t, p_h, p_w = self.config.patch_size + + # 2. Process noisy latents + hidden_states = self.patch_embedding(hidden_states) + _, _, post_patch_num_frames, post_patch_height, post_patch_width = hidden_states.shape + + if indices_hidden_states is None: + indices_hidden_states = torch.arange(0, post_patch_num_frames).unsqueeze(0).expand(batch_size, -1) + + hidden_states = hidden_states.flatten(2).transpose(1, 2) + rotary_emb = self.rope( + frame_indices=indices_hidden_states, + height=post_patch_height, + width=post_patch_width, + device=hidden_states.device, + ) + rotary_emb = rotary_emb.flatten(2).transpose(1, 2) + original_context_length = hidden_states.shape[1] + + # 3. Process short history latents + if latents_history_short is not None and indices_latents_history_short is not None: + latents_history_short = self.patch_short(latents_history_short) + _, _, _, H1, W1 = latents_history_short.shape + latents_history_short = latents_history_short.flatten(2).transpose(1, 2) + + rotary_emb_history_short = self.rope( + frame_indices=indices_latents_history_short, + height=H1, + width=W1, + device=latents_history_short.device, + ) + rotary_emb_history_short = rotary_emb_history_short.flatten(2).transpose(1, 2) + + hidden_states = torch.cat([latents_history_short, hidden_states], dim=1) + rotary_emb = torch.cat([rotary_emb_history_short, rotary_emb], dim=1) + + # 4. Process mid history latents + if latents_history_mid is not None and indices_latents_history_mid is not None: + latents_history_mid = pad_for_3d_conv(latents_history_mid, (2, 4, 4)) + latents_history_mid = self.patch_mid(latents_history_mid) + latents_history_mid = latents_history_mid.flatten(2).transpose(1, 2) + + rotary_emb_history_mid = self.rope( + frame_indices=indices_latents_history_mid, + height=H1, + width=W1, + device=latents_history_mid.device, + ) + rotary_emb_history_mid = pad_for_3d_conv(rotary_emb_history_mid, (2, 2, 2)) + rotary_emb_history_mid = center_down_sample_3d(rotary_emb_history_mid, (2, 2, 2)) + rotary_emb_history_mid = rotary_emb_history_mid.flatten(2).transpose(1, 2) + + hidden_states = torch.cat([latents_history_mid, hidden_states], dim=1) + rotary_emb = torch.cat([rotary_emb_history_mid, rotary_emb], dim=1) + + # 5. Process long history latents + if latents_history_long is not None and indices_latents_history_long is not None: + latents_history_long = pad_for_3d_conv(latents_history_long, (4, 8, 8)) + latents_history_long = self.patch_long(latents_history_long) + latents_history_long = latents_history_long.flatten(2).transpose(1, 2) + + rotary_emb_history_long = self.rope( + frame_indices=indices_latents_history_long, + height=H1, + width=W1, + device=latents_history_long.device, + ) + rotary_emb_history_long = pad_for_3d_conv(rotary_emb_history_long, (4, 4, 4)) + rotary_emb_history_long = center_down_sample_3d(rotary_emb_history_long, (4, 4, 4)) + rotary_emb_history_long = rotary_emb_history_long.flatten(2).transpose(1, 2) + + hidden_states = torch.cat([latents_history_long, hidden_states], dim=1) + rotary_emb = torch.cat([rotary_emb_history_long, rotary_emb], dim=1) + + history_context_length = hidden_states.shape[1] - original_context_length + + if indices_hidden_states is not None and self.zero_history_timestep: + timestep_t0 = torch.zeros((1), dtype=timestep.dtype, device=timestep.device) + temb_t0, timestep_proj_t0, _ = self.condition_embedder( + timestep_t0, encoder_hidden_states, is_return_encoder_hidden_states=False + ) + temb_t0 = temb_t0.unsqueeze(1).expand(batch_size, history_context_length, -1) + timestep_proj_t0 = ( + timestep_proj_t0.unflatten(-1, (6, -1)) + .view(1, 6, 1, -1) + .expand(batch_size, -1, history_context_length, -1) + ) + + temb, timestep_proj, encoder_hidden_states = self.condition_embedder(timestep, encoder_hidden_states) + timestep_proj = timestep_proj.unflatten(-1, (6, -1)) + + if indices_hidden_states is not None and not self.zero_history_timestep: + main_repeat_size = hidden_states.shape[1] + else: + main_repeat_size = original_context_length + temb = temb.view(batch_size, 1, -1).expand(batch_size, main_repeat_size, -1) + timestep_proj = timestep_proj.view(batch_size, 6, 1, -1).expand(batch_size, 6, main_repeat_size, -1) + + if indices_hidden_states is not None and self.zero_history_timestep: + temb = torch.cat([temb_t0, temb], dim=1) + timestep_proj = torch.cat([timestep_proj_t0, timestep_proj], dim=2) + + if timestep_proj.ndim == 4: + timestep_proj = timestep_proj.permute(0, 2, 1, 3) + + # 6. Transformer blocks + hidden_states = hidden_states.contiguous() + encoder_hidden_states = encoder_hidden_states.contiguous() + rotary_emb = rotary_emb.contiguous() + if torch.is_grad_enabled() and self.gradient_checkpointing: + for block in self.blocks: + hidden_states = self._gradient_checkpointing_func( + block, + hidden_states, + encoder_hidden_states, + timestep_proj, + rotary_emb, + original_context_length, + ) + else: + for block in self.blocks: + hidden_states = block( + hidden_states, + encoder_hidden_states, + timestep_proj, + rotary_emb, + original_context_length, + ) + + # 7. Normalization + hidden_states = self.norm_out(hidden_states, temb, original_context_length) + hidden_states = self.proj_out(hidden_states) + + # 8. Unpatchify + hidden_states = hidden_states.reshape( + batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1 + ) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + if not return_dict: + return (output,) + + return Transformer2DModelOutput(sample=output) diff --git a/src/diffusers/models/transformers/transformer_hidream_image.py b/src/diffusers/models/transformers/transformer_hidream_image.py index 4a5aee29abc4..57c0a16eee6f 100644 --- a/src/diffusers/models/transformers/transformer_hidream_image.py +++ b/src/diffusers/models/transformers/transformer_hidream_image.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -8,7 +8,7 @@ from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...models.modeling_outputs import Transformer2DModelOutput from ...models.modeling_utils import ModelMixin -from ...utils import USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, deprecate, logging from ...utils.torch_utils import maybe_allow_in_graph from ..attention import Attention from ..embeddings import TimestepEmbedding, Timesteps @@ -23,7 +23,7 @@ def __init__( dim: int, hidden_dim: int, multiple_of: int = 256, - ffn_dim_multiplier: Optional[float] = None, + ffn_dim_multiplier: float | None = None, ): super().__init__() hidden_dim = int(2 * hidden_dim / 3) @@ -55,7 +55,7 @@ def __init__(self, hidden_size, frequency_embedding_size=256): self.time_proj = Timesteps(num_channels=frequency_embedding_size, flip_sin_to_cos=True, downscale_freq_shift=0) self.timestep_embedder = TimestepEmbedding(in_channels=frequency_embedding_size, time_embed_dim=hidden_size) - def forward(self, timesteps: torch.Tensor, wdtype: Optional[torch.dtype] = None) -> torch.Tensor: + def forward(self, timesteps: torch.Tensor, wdtype: torch.dtype | None = None) -> torch.Tensor: t_emb = self.time_proj(timesteps).to(dtype=wdtype) t_emb = self.timestep_embedder(t_emb) return t_emb @@ -114,7 +114,7 @@ def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor: class HiDreamImageEmbedND(nn.Module): - def __init__(self, theta: int, axes_dim: List[int]): + def __init__(self, theta: int, axes_dim: list[int]): super().__init__() self.theta = theta self.axes_dim = axes_dim @@ -128,7 +128,7 @@ def forward(self, ids: torch.Tensor) -> torch.Tensor: return emb.unsqueeze(2) -def apply_rope(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: +def apply_rope(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2) xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2) xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1] @@ -205,8 +205,8 @@ def __call__( self, attn: HiDreamAttention, hidden_states: torch.Tensor, - hidden_states_masks: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, + hidden_states_masks: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, image_rotary_emb: torch.Tensor = None, *args, **kwargs, @@ -458,9 +458,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - hidden_states_masks: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - temb: Optional[torch.Tensor] = None, + hidden_states_masks: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + temb: torch.Tensor | None = None, image_rotary_emb: torch.Tensor = None, ) -> torch.Tensor: wtype = hidden_states.dtype @@ -530,11 +530,11 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - hidden_states_masks: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - temb: Optional[torch.Tensor] = None, + hidden_states_masks: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + temb: torch.Tensor | None = None, image_rotary_emb: torch.Tensor = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: wtype = hidden_states.dtype ( shift_msa_i, @@ -581,18 +581,18 @@ def forward( class HiDreamBlock(nn.Module): - def __init__(self, block: Union[HiDreamImageTransformerBlock, HiDreamImageSingleTransformerBlock]): + def __init__(self, block: HiDreamImageTransformerBlock | HiDreamImageSingleTransformerBlock): super().__init__() self.block = block def forward( self, hidden_states: torch.Tensor, - hidden_states_masks: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - temb: Optional[torch.Tensor] = None, + hidden_states_masks: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + temb: torch.Tensor | None = None, image_rotary_emb: torch.Tensor = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: return self.block( hidden_states=hidden_states, hidden_states_masks=hidden_states_masks, @@ -609,20 +609,20 @@ class HiDreamImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, @register_to_config def __init__( self, - patch_size: Optional[int] = None, + patch_size: int | None = None, in_channels: int = 64, - out_channels: Optional[int] = None, + out_channels: int | None = None, num_layers: int = 16, num_single_layers: int = 32, attention_head_dim: int = 128, num_attention_heads: int = 20, - caption_channels: List[int] = None, + caption_channels: list[int] = None, text_emb_dim: int = 2048, num_routed_experts: int = 4, num_activated_experts: int = 2, - axes_dims_rope: Tuple[int, int] = (32, 32), - max_resolution: Tuple[int, int] = (128, 128), - llama_layers: List[int] = None, + axes_dims_rope: tuple[int, int] = (32, 32), + max_resolution: tuple[int, int] = (128, 128), + llama_layers: list[int] = None, force_inference_output: bool = False, ): super().__init__() @@ -681,7 +681,7 @@ def __init__( self.gradient_checkpointing = False - def unpatchify(self, x: torch.Tensor, img_sizes: List[Tuple[int, int]], is_training: bool) -> List[torch.Tensor]: + def unpatchify(self, x: torch.Tensor, img_sizes: list[tuple[int, int]], is_training: bool) -> list[torch.Tensor]: if is_training and not self.config.force_inference_output: B, S, F = x.shape C = F // (self.config.patch_size * self.config.patch_size) @@ -773,6 +773,7 @@ def patchify(self, hidden_states): return hidden_states, hidden_states_masks, img_sizes, img_ids + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, @@ -780,13 +781,13 @@ def forward( encoder_hidden_states_t5: torch.Tensor = None, encoder_hidden_states_llama3: torch.Tensor = None, pooled_embeds: torch.Tensor = None, - img_ids: Optional[torch.Tensor] = None, - img_sizes: Optional[List[Tuple[int, int]]] = None, - hidden_states_masks: Optional[torch.Tensor] = None, - attention_kwargs: Optional[Dict[str, Any]] = None, + img_ids: torch.Tensor | None = None, + img_sizes: list[tuple[int, int]] | None = None, + hidden_states_masks: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, **kwargs, - ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: + ) -> tuple[torch.Tensor] | Transformer2DModelOutput: encoder_hidden_states = kwargs.get("encoder_hidden_states", None) if encoder_hidden_states is not None: @@ -808,21 +809,6 @@ def forward( "if `hidden_states_masks` is passed, `hidden_states` must be a 3D tensors with shape (batch_size, patch_height * patch_width, patch_size * patch_size * channels)" ) - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) - # spatial forward batch_size = hidden_states.shape[0] hidden_states_type = hidden_states.dtype @@ -933,10 +919,6 @@ def forward( if hidden_states_masks is not None: hidden_states_masks = hidden_states_masks[:, :image_tokens_seq_len] - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) return Transformer2DModelOutput(sample=output) diff --git a/src/diffusers/models/transformers/transformer_hunyuan_video.py b/src/diffusers/models/transformers/transformer_hunyuan_video.py index fb0ce1a30ff9..aab593c93643 100644 --- a/src/diffusers/models/transformers/transformer_hunyuan_video.py +++ b/src/diffusers/models/transformers/transformer_hunyuan_video.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -22,7 +22,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, logging from ..attention import AttentionMixin, FeedForward from ..attention_dispatch import dispatch_attention_fn from ..attention_processor import Attention @@ -56,9 +56,9 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: if attn.add_q_proj is None and encoder_hidden_states is not None: hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) @@ -162,7 +162,7 @@ def __call__( class HunyuanVideoPatchEmbed(nn.Module): def __init__( self, - patch_size: Union[int, Tuple[int, int, int]] = 16, + patch_size: int | tuple[int, int, int] = 16, in_chans: int = 3, embed_dim: int = 768, ) -> None: @@ -178,7 +178,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class HunyuanVideoAdaNorm(nn.Module): - def __init__(self, in_features: int, out_features: Optional[int] = None) -> None: + def __init__(self, in_features: int, out_features: int | None = None) -> None: super().__init__() out_features = out_features or 2 * in_features @@ -187,7 +187,7 @@ def __init__(self, in_features: int, out_features: Optional[int] = None) -> None def forward( self, temb: torch.Tensor - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: temb = self.linear(self.nonlinearity(temb)) gate_msa, gate_mlp = temb.chunk(2, dim=1) gate_msa, gate_mlp = gate_msa.unsqueeze(1), gate_mlp.unsqueeze(1) @@ -216,7 +216,7 @@ def forward( emb: torch.Tensor, token_replace_emb: torch.Tensor, first_frame_num_tokens: int, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: emb = self.linear(self.silu(emb)) token_replace_emb = self.linear(self.silu(token_replace_emb)) @@ -267,7 +267,7 @@ def forward( emb: torch.Tensor, token_replace_emb: torch.Tensor, first_frame_num_tokens: int, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: emb = self.linear(self.silu(emb)) token_replace_emb = self.linear(self.silu(token_replace_emb)) @@ -292,7 +292,7 @@ def __init__( embedding_dim: int, pooled_projection_dim: int, guidance_embeds: bool, - image_condition_type: Optional[str] = None, + image_condition_type: str | None = None, ): super().__init__() @@ -307,12 +307,11 @@ def __init__( self.guidance_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) def forward( - self, timestep: torch.Tensor, pooled_projection: torch.Tensor, guidance: Optional[torch.Tensor] = None - ) -> Tuple[torch.Tensor, torch.Tensor]: + self, timestep: torch.Tensor, pooled_projection: torch.Tensor, guidance: torch.Tensor | None = None + ) -> tuple[torch.Tensor, torch.Tensor]: timesteps_proj = self.time_proj(timestep) timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype)) # (N, D) pooled_projections = self.text_embedder(pooled_projection) - conditioning = timesteps_emb + pooled_projections token_replace_emb = None if self.image_condition_type == "token_replace": @@ -324,8 +323,9 @@ def forward( if self.guidance_embedder is not None: guidance_proj = self.time_proj(guidance) guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=pooled_projection.dtype)) - conditioning = conditioning + guidance_emb - + conditioning = timesteps_emb + guidance_emb + pooled_projections + else: + conditioning = timesteps_emb + pooled_projections return conditioning, token_replace_emb @@ -360,7 +360,7 @@ def forward( self, hidden_states: torch.Tensor, temb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: norm_hidden_states = self.norm1(hidden_states) @@ -408,7 +408,7 @@ def forward( self, hidden_states: torch.Tensor, temb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, ) -> None: self_attn_mask = None if attention_mask is not None: @@ -458,7 +458,7 @@ def forward( self, hidden_states: torch.Tensor, timestep: torch.LongTensor, - attention_mask: Optional[torch.LongTensor] = None, + attention_mask: torch.LongTensor | None = None, ) -> torch.Tensor: if attention_mask is None: pooled_projections = hidden_states.mean(dim=1) @@ -476,7 +476,7 @@ def forward( class HunyuanVideoRotaryPosEmbed(nn.Module): - def __init__(self, patch_size: int, patch_size_t: int, rope_dim: List[int], theta: float = 256.0) -> None: + def __init__(self, patch_size: int, patch_size_t: int, rope_dim: list[int], theta: float = 256.0) -> None: super().__init__() self.patch_size = patch_size @@ -544,11 +544,11 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, *args, **kwargs, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: text_seq_length = encoder_hidden_states.shape[1] hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) @@ -624,11 +624,11 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + attention_mask: torch.Tensor | None = None, + freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, *args, **kwargs, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: # 1. Input normalization norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( @@ -699,11 +699,11 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, token_replace_emb: torch.Tensor = None, num_tokens: int = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: text_seq_length = encoder_hidden_states.shape[1] hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) @@ -783,11 +783,11 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + attention_mask: torch.Tensor | None = None, + freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, token_replace_emb: torch.Tensor = None, num_tokens: int = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: # 1. Input normalization ( norm_hidden_states, @@ -875,7 +875,7 @@ class HunyuanVideoTransformer3DModel( The dimension of the pooled projection of the text embeddings. rope_theta (`float`, defaults to `256.0`): The value of theta to use in the RoPE layer. - rope_axes_dim (`Tuple[int]`, defaults to `(16, 56, 56)`): + rope_axes_dim (`tuple[int]`, defaults to `(16, 56, 56)`): The dimensions of the axes to use in the RoPE layer. image_condition_type (`str`, *optional*, defaults to `None`): The type of image conditioning to use. If `None`, no image conditioning is used. If `latent_concat`, the @@ -916,8 +916,8 @@ def __init__( text_embed_dim: int = 4096, pooled_projection_dim: int = 768, rope_theta: float = 256.0, - rope_axes_dim: Tuple[int, ...] = (16, 56, 56), - image_condition_type: Optional[str] = None, + rope_axes_dim: tuple[int, ...] = (16, 56, 56), + image_condition_type: str | None = None, ) -> None: super().__init__() @@ -989,6 +989,7 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, @@ -997,24 +998,9 @@ def forward( encoder_attention_mask: torch.Tensor, pooled_projections: torch.Tensor, guidance: torch.Tensor = None, - attention_kwargs: Optional[Dict[str, Any]] = None, + attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) - + ) -> tuple[torch.Tensor] | Transformer2DModelOutput: batch_size, num_channels, num_frames, height, width = hidden_states.shape p, p_t = self.config.patch_size, self.config.patch_size_t post_patch_num_frames = num_frames // p_t @@ -1104,10 +1090,6 @@ def forward( hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7) hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (hidden_states,) diff --git a/src/diffusers/models/transformers/transformer_hunyuan_video15.py b/src/diffusers/models/transformers/transformer_hunyuan_video15.py index 293ba996ea98..222b0791d650 100644 --- a/src/diffusers/models/transformers/transformer_hunyuan_video15.py +++ b/src/diffusers/models/transformers/transformer_hunyuan_video15.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -22,7 +22,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, logging from ..attention import AttentionMixin, FeedForward from ..attention_dispatch import dispatch_attention_fn from ..attention_processor import Attention @@ -55,9 +55,9 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: # 1. QKV projections query = attn.to_q(hidden_states) @@ -140,7 +140,7 @@ def __call__( class HunyuanVideo15PatchEmbed(nn.Module): def __init__( self, - patch_size: Union[int, Tuple[int, int, int]] = 16, + patch_size: int | tuple[int, int, int] = 16, in_chans: int = 3, embed_dim: int = 768, ) -> None: @@ -156,7 +156,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class HunyuanVideo15AdaNorm(nn.Module): - def __init__(self, in_features: int, out_features: Optional[int] = None) -> None: + def __init__(self, in_features: int, out_features: int | None = None) -> None: super().__init__() out_features = out_features or 2 * in_features @@ -165,7 +165,7 @@ def __init__(self, in_features: int, out_features: Optional[int] = None) -> None def forward( self, temb: torch.Tensor - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: temb = self.linear(self.nonlinearity(temb)) gate_msa, gate_mlp = temb.chunk(2, dim=1) gate_msa, gate_mlp = gate_msa.unsqueeze(1), gate_mlp.unsqueeze(1) @@ -200,7 +200,7 @@ def __init__(self, embedding_dim: int, use_meanflow: bool = False): def forward( self, timestep: torch.Tensor, - timestep_r: Optional[torch.Tensor] = None, + timestep_r: torch.Tensor | None = None, ) -> torch.Tensor: timesteps_proj = self.time_proj(timestep) timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=timestep.dtype)) @@ -244,7 +244,7 @@ def forward( self, hidden_states: torch.Tensor, temb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: norm_hidden_states = self.norm1(hidden_states) @@ -292,7 +292,7 @@ def forward( self, hidden_states: torch.Tensor, temb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, ) -> None: self_attn_mask = None if attention_mask is not None: @@ -341,7 +341,7 @@ def forward( self, hidden_states: torch.Tensor, timestep: torch.LongTensor, - attention_mask: Optional[torch.LongTensor] = None, + attention_mask: torch.LongTensor | None = None, ) -> torch.Tensor: if attention_mask is None: pooled_projections = hidden_states.mean(dim=1) @@ -359,7 +359,7 @@ def forward( class HunyuanVideo15RotaryPosEmbed(nn.Module): - def __init__(self, patch_size: int, patch_size_t: int, rope_dim: List[int], theta: float = 256.0) -> None: + def __init__(self, patch_size: int, patch_size_t: int, rope_dim: list[int], theta: float = 256.0) -> None: super().__init__() self.patch_size = patch_size @@ -468,11 +468,11 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + attention_mask: torch.Tensor | None = None, + freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, *args, **kwargs, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: # 1. Input normalization norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( @@ -542,7 +542,7 @@ class HunyuanVideo15Transformer3DModel( The dimension of the pooled projection of the text embeddings. rope_theta (`float`, defaults to `256.0`): The value of theta to use in the RoPE layer. - rope_axes_dim (`Tuple[int]`, defaults to `(16, 56, 56)`): + rope_axes_dim (`tuple[int]`, defaults to `(16, 56, 56)`): The dimensions of the axes to use in the RoPE layer. """ @@ -576,7 +576,7 @@ def __init__( text_embed_2_dim: int = 1472, image_embed_dim: int = 1152, rope_theta: float = 256.0, - rope_axes_dim: Tuple[int, ...] = (16, 56, 56), + rope_axes_dim: tuple[int, ...] = (16, 56, 56), # YiYi Notes: config based on target_size_config https://github.com/yiyixuxu/hy15/blob/main/hyvideo/pipelines/hunyuan_video_pipeline.py#L205 target_size: int = 640, # did not name sample_size since it is in pixel spaces task_type: str = "i2v", @@ -620,34 +620,20 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, timestep: torch.LongTensor, encoder_hidden_states: torch.Tensor, encoder_attention_mask: torch.Tensor, - timestep_r: Optional[torch.LongTensor] = None, - encoder_hidden_states_2: Optional[torch.Tensor] = None, - encoder_attention_mask_2: Optional[torch.Tensor] = None, - image_embeds: Optional[torch.Tensor] = None, - attention_kwargs: Optional[Dict[str, Any]] = None, + timestep_r: torch.LongTensor | None = None, + encoder_hidden_states_2: torch.Tensor | None = None, + encoder_attention_mask_2: torch.Tensor | None = None, + image_embeds: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) - + ) -> tuple[torch.Tensor] | Transformer2DModelOutput: batch_size, num_channels, num_frames, height, width = hidden_states.shape p_t, p_h, p_w = self.config.patch_size_t, self.config.patch_size, self.config.patch_size post_patch_num_frames = num_frames // p_t @@ -783,10 +769,6 @@ def forward( hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7) hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (hidden_states,) diff --git a/src/diffusers/models/transformers/transformer_hunyuan_video_framepack.py b/src/diffusers/models/transformers/transformer_hunyuan_video_framepack.py index 601ba0f0b472..f005c4d4cd51 100644 --- a/src/diffusers/models/transformers/transformer_hunyuan_video_framepack.py +++ b/src/diffusers/models/transformers/transformer_hunyuan_video_framepack.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -20,7 +20,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, get_logger, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, get_logger from ..cache_utils import CacheMixin from ..embeddings import get_1d_rotary_pos_embed from ..modeling_outputs import Transformer2DModelOutput @@ -39,7 +39,7 @@ class HunyuanVideoFramepackRotaryPosEmbed(nn.Module): - def __init__(self, patch_size: int, patch_size_t: int, rope_dim: List[int], theta: float = 256.0) -> None: + def __init__(self, patch_size: int, patch_size_t: int, rope_dim: list[int], theta: float = 256.0) -> None: super().__init__() self.patch_size = patch_size @@ -91,9 +91,9 @@ def __init__(self, in_channels: int, inner_dim: int): def forward( self, - latents_clean: Optional[torch.Tensor] = None, - latents_clean_2x: Optional[torch.Tensor] = None, - latents_clean_4x: Optional[torch.Tensor] = None, + latents_clean: torch.Tensor | None = None, + latents_clean_2x: torch.Tensor | None = None, + latents_clean_4x: torch.Tensor | None = None, ): if latents_clean is not None: latents_clean = self.proj(latents_clean) @@ -139,8 +139,8 @@ def __init__( text_embed_dim: int = 4096, pooled_projection_dim: int = 768, rope_theta: float = 256.0, - rope_axes_dim: Tuple[int, ...] = (16, 56, 56), - image_condition_type: Optional[str] = None, + rope_axes_dim: tuple[int, ...] = (16, 56, 56), + image_condition_type: str | None = None, has_image_proj: int = False, image_proj_dim: int = 1152, has_clean_x_embedder: int = False, @@ -198,6 +198,7 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, @@ -207,31 +208,16 @@ def forward( pooled_projections: torch.Tensor, image_embeds: torch.Tensor, indices_latents: torch.Tensor, - guidance: Optional[torch.Tensor] = None, - latents_clean: Optional[torch.Tensor] = None, - indices_latents_clean: Optional[torch.Tensor] = None, - latents_history_2x: Optional[torch.Tensor] = None, - indices_latents_history_2x: Optional[torch.Tensor] = None, - latents_history_4x: Optional[torch.Tensor] = None, - indices_latents_history_4x: Optional[torch.Tensor] = None, - attention_kwargs: Optional[Dict[str, Any]] = None, + guidance: torch.Tensor | None = None, + latents_clean: torch.Tensor | None = None, + indices_latents_clean: torch.Tensor | None = None, + latents_history_2x: torch.Tensor | None = None, + indices_latents_history_2x: torch.Tensor | None = None, + latents_history_4x: torch.Tensor | None = None, + indices_latents_history_4x: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) - + ) -> tuple[torch.Tensor] | Transformer2DModelOutput: batch_size, num_channels, num_frames, height, width = hidden_states.shape p, p_t = self.config.patch_size, self.config.patch_size_t post_patch_num_frames = num_frames // p_t @@ -337,10 +323,6 @@ def forward( hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7) hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (hidden_states,) return Transformer2DModelOutput(sample=hidden_states) @@ -348,13 +330,13 @@ def forward( def _pack_history_states( self, hidden_states: torch.Tensor, - latents_clean: Optional[torch.Tensor] = None, - latents_history_2x: Optional[torch.Tensor] = None, - latents_history_4x: Optional[torch.Tensor] = None, - image_rotary_emb: Tuple[torch.Tensor, torch.Tensor] = None, - image_rotary_emb_clean: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - image_rotary_emb_history_2x: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - image_rotary_emb_history_4x: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + latents_clean: torch.Tensor | None = None, + latents_history_2x: torch.Tensor | None = None, + latents_history_4x: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] = None, + image_rotary_emb_clean: tuple[torch.Tensor, torch.Tensor] | None = None, + image_rotary_emb_history_2x: tuple[torch.Tensor, torch.Tensor] | None = None, + image_rotary_emb_history_4x: tuple[torch.Tensor, torch.Tensor] | None = None, height: int = None, width: int = None, ): @@ -381,10 +363,10 @@ def _pack_history_states( def _pad_rotary_emb( self, - image_rotary_emb: Tuple[torch.Tensor], + image_rotary_emb: tuple[torch.Tensor], height: int, width: int, - kernel_size: Tuple[int, int, int], + kernel_size: tuple[int, int, int], ): # freqs_cos, freqs_sin have shape [W * H * T, D / 2], where D is attention head dim freqs_cos, freqs_sin = image_rotary_emb diff --git a/src/diffusers/models/transformers/transformer_hunyuanimage.py b/src/diffusers/models/transformers/transformer_hunyuanimage.py index d626e322ad6f..a2d3d9229963 100644 --- a/src/diffusers/models/transformers/transformer_hunyuanimage.py +++ b/src/diffusers/models/transformers/transformer_hunyuanimage.py @@ -13,7 +13,7 @@ # limitations under the License. import math -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -23,7 +23,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, logging from ...utils.torch_utils import maybe_allow_in_graph from ..attention import AttentionMixin, FeedForward from ..attention_dispatch import dispatch_attention_fn @@ -57,9 +57,9 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: if attn.add_q_proj is None and encoder_hidden_states is not None: hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) @@ -157,7 +157,7 @@ def __call__( class HunyuanImagePatchEmbed(nn.Module): def __init__( self, - patch_size: Union[Tuple[int, int], Tuple[int, int, int]] = (16, 16), + patch_size: tuple[int, int, tuple[int, int, int]] = (16, 16), in_chans: int = 3, embed_dim: int = 768, ) -> None: @@ -198,7 +198,7 @@ def forward(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor: class HunyuanImageAdaNorm(nn.Module): - def __init__(self, in_features: int, out_features: Optional[int] = None) -> None: + def __init__(self, in_features: int, out_features: int | None = None) -> None: super().__init__() out_features = out_features or 2 * in_features @@ -207,7 +207,7 @@ def __init__(self, in_features: int, out_features: Optional[int] = None) -> None def forward( self, temb: torch.Tensor - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: temb = self.linear(self.nonlinearity(temb)) gate_msa, gate_mlp = temb.chunk(2, dim=1) gate_msa, gate_mlp = gate_msa.unsqueeze(1), gate_mlp.unsqueeze(1) @@ -241,9 +241,9 @@ def __init__( def forward( self, timestep: torch.Tensor, - timestep_r: Optional[torch.Tensor] = None, - guidance: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + timestep_r: torch.Tensor | None = None, + guidance: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: timesteps_proj = self.time_proj(timestep) timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=timestep.dtype)) @@ -295,7 +295,7 @@ def forward( self, hidden_states: torch.Tensor, temb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: norm_hidden_states = self.norm1(hidden_states) @@ -343,7 +343,7 @@ def forward( self, hidden_states: torch.Tensor, temb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, ) -> None: self_attn_mask = None if attention_mask is not None: @@ -394,7 +394,7 @@ def forward( self, hidden_states: torch.Tensor, timestep: torch.LongTensor, - attention_mask: Optional[torch.LongTensor] = None, + attention_mask: torch.LongTensor | None = None, ) -> torch.Tensor: if attention_mask is None: pooled_hidden_states = hidden_states.mean(dim=1) @@ -412,9 +412,7 @@ def forward( class HunyuanImageRotaryPosEmbed(nn.Module): - def __init__( - self, patch_size: Union[Tuple, List[int]], rope_dim: Union[Tuple, List[int]], theta: float = 256.0 - ) -> None: + def __init__(self, patch_size: tuple | list[int], rope_dim: tuple | list[int], theta: float = 256.0) -> None: super().__init__() if not isinstance(patch_size, (tuple, list)) or len(patch_size) not in [2, 3]: @@ -496,8 +494,8 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, *args, **kwargs, ) -> torch.Tensor: @@ -577,11 +575,11 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, *args, **kwargs, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: # 1. Input normalization norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( @@ -653,7 +651,7 @@ class HunyuanImageTransformer2DModel( The dimension of the pooled projection of the text embeddings. rope_theta (`float`, defaults to `256.0`): The value of theta to use in the RoPE layer. - rope_axes_dim (`Tuple[int]`, defaults to `(16, 56, 56)`): + rope_axes_dim (`tuple[int]`, defaults to `(16, 56, 56)`): The dimensions of the axes to use in the RoPE layer. image_condition_type (`str`, *optional*, defaults to `None`): The type of image conditioning to use. If `None`, no image conditioning is used. If `latent_concat`, the @@ -682,13 +680,13 @@ def __init__( num_single_layers: int = 40, num_refiner_layers: int = 2, mlp_ratio: float = 4.0, - patch_size: Tuple[int, int] = (1, 1), + patch_size: tuple[int, int] = (1, 1), qk_norm: str = "rms_norm", guidance_embeds: bool = False, text_embed_dim: int = 3584, - text_embed_2_dim: Optional[int] = None, + text_embed_2_dim: int | None = None, rope_theta: float = 256.0, - rope_axes_dim: Tuple[int, ...] = (64, 64), + rope_axes_dim: tuple[int, ...] = (64, 64), use_meanflow: bool = False, ) -> None: super().__init__() @@ -742,34 +740,20 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, timestep: torch.LongTensor, encoder_hidden_states: torch.Tensor, encoder_attention_mask: torch.Tensor, - timestep_r: Optional[torch.LongTensor] = None, - encoder_hidden_states_2: Optional[torch.Tensor] = None, - encoder_attention_mask_2: Optional[torch.Tensor] = None, - guidance: Optional[torch.Tensor] = None, - attention_kwargs: Optional[Dict[str, Any]] = None, + timestep_r: torch.LongTensor | None = None, + encoder_hidden_states_2: torch.Tensor | None = None, + encoder_attention_mask_2: torch.Tensor | None = None, + guidance: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) - + ) -> torch.Tensor | dict[str, torch.Tensor]: if hidden_states.ndim == 4: batch_size, channels, height, width = hidden_states.shape sizes = (height, width) @@ -900,10 +884,6 @@ def forward( ] hidden_states = hidden_states.reshape(*final_dims) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (hidden_states,) diff --git a/src/diffusers/models/transformers/transformer_kandinsky.py b/src/diffusers/models/transformers/transformer_kandinsky.py index 316e79da4fd6..88ef70d546c8 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky.py +++ b/src/diffusers/models/transformers/transformer_kandinsky.py @@ -14,7 +14,7 @@ import inspect import math -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -165,9 +165,8 @@ def __init__(self, model_dim, time_dim, max_period=10000.0): self.activation = nn.SiLU() self.out_layer = nn.Linear(time_dim, time_dim, bias=True) - @torch.autocast(device_type="cuda", dtype=torch.float32) def forward(self, time): - args = torch.outer(time, self.freqs.to(device=time.device)) + args = torch.outer(time.to(torch.float32), self.freqs.to(device=time.device)) time_embed = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) time_embed = self.out_layer(self.activation(self.in_layer(time_embed))) return time_embed @@ -269,7 +268,6 @@ def __init__(self, time_dim, model_dim, num_params): self.out_layer.weight.data.zero_() self.out_layer.bias.data.zero_() - @torch.autocast(device_type="cuda", dtype=torch.float32) def forward(self, x): return self.out_layer(self.activation(x)) @@ -368,9 +366,9 @@ def __init__(self, num_channels, head_dim, processor=None): def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - sparse_params: Optional[torch.Tensor] = None, - rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + encoder_hidden_states: torch.Tensor | None = None, + sparse_params: torch.Tensor | None = None, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, **kwargs, ) -> torch.Tensor: attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) @@ -525,6 +523,7 @@ class Kandinsky5Transformer3DModel( "Kandinsky5TransformerEncoderBlock", "Kandinsky5TransformerDecoderBlock", ] + _keep_in_fp32_modules = ["time_embeddings", "modulation", "visual_modulation", "text_modulation"] _supports_gradient_checkpointing = True @register_to_config @@ -597,12 +596,12 @@ def forward( encoder_hidden_states: torch.Tensor, # text_embed timestep: torch.Tensor, # time pooled_projections: torch.Tensor, # pooled_text_embed - visual_rope_pos: Tuple[int, int, int], + visual_rope_pos: tuple[int, int, int], text_rope_pos: torch.LongTensor, - scale_factor: Tuple[float, float, float] = (1.0, 1.0, 1.0), - sparse_params: Optional[Dict[str, Any]] = None, + scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0), + sparse_params: dict[str, Any] | None = None, return_dict: bool = True, - ) -> Union[Transformer2DModelOutput, torch.FloatTensor]: + ) -> Transformer2DModelOutput | torch.FloatTensor: """ Forward pass of the Kandinsky5 3D Transformer. @@ -611,10 +610,10 @@ def forward( encoder_hidden_states (`torch.FloatTensor`): Text embeddings timestep (`torch.Tensor` or `float` or `int`): Current timestep pooled_projections (`torch.FloatTensor`): Pooled text embeddings - visual_rope_pos (`Tuple[int, int, int]`): Position for visual RoPE + visual_rope_pos (`tuple[int, int, int]`): Position for visual RoPE text_rope_pos (`torch.LongTensor`): Position for text RoPE - scale_factor (`Tuple[float, float, float]`, optional): Scale factor for RoPE - sparse_params (`Dict[str, Any]`, optional): Parameters for sparse attention + scale_factor (`tuple[float, float, float]`, optional): Scale factor for RoPE + sparse_params (`dict[str, Any]`, optional): Parameters for sparse attention return_dict (`bool`, optional): Whether to return a dictionary Returns: diff --git a/src/diffusers/models/transformers/transformer_longcat_image.py b/src/diffusers/models/transformers/transformer_longcat_image.py index 7fbaaa3fee85..7a000fa2b2ce 100644 --- a/src/diffusers/models/transformers/transformer_longcat_image.py +++ b/src/diffusers/models/transformers/transformer_longcat_image.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -21,9 +21,9 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin -from ...utils import is_torch_npu_available, logging +from ...utils import logging from ...utils.torch_utils import maybe_allow_in_graph -from ..attention import AttentionModuleMixin, FeedForward +from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward from ..attention_dispatch import dispatch_attention_fn from ..cache_utils import CacheMixin from ..embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed @@ -78,8 +78,8 @@ def __call__( attn: "LongCatImageAttention", hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections( attn, hidden_states, encoder_hidden_states @@ -145,12 +145,12 @@ def __init__( dim_head: int = 64, dropout: float = 0.0, bias: bool = False, - added_kv_proj_dim: Optional[int] = None, - added_proj_bias: Optional[bool] = True, + added_kv_proj_dim: int | None = None, + added_proj_bias: bool | None = True, out_bias: bool = True, eps: float = 1e-5, out_dim: int = None, - context_pre_only: Optional[bool] = None, + context_pre_only: bool | None = None, pre_only: bool = False, elementwise_affine: bool = True, processor=None, @@ -195,9 +195,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) @@ -238,9 +238,9 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: text_seq_len = encoder_hidden_states.shape[1] hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) @@ -298,9 +298,9 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( @@ -351,7 +351,7 @@ def forward( class LongCatImagePosEmbed(nn.Module): - def __init__(self, theta: int, axes_dim: List[int]): + def __init__(self, theta: int, axes_dim: list[int]): super().__init__() self.theta = theta self.axes_dim = axes_dim @@ -400,12 +400,14 @@ class LongCatImageTransformer2DModel( PeftAdapterMixin, FromOriginalModelMixin, CacheMixin, + AttentionMixin, ): """ The Transformer model introduced in Longcat-Image. """ _supports_gradient_checkpointing = True + _repeated_blocks = ["LongCatImageTransformerBlock", "LongCatImageSingleTransformerBlock"] @register_to_config def __init__( @@ -418,7 +420,7 @@ def __init__( num_attention_heads: int = 24, joint_attention_dim: int = 3584, pooled_projection_dim: int = 3584, - axes_dims_rope: List[int] = [16, 56, 56], + axes_dims_rope: list[int] = [16, 56, 56], ): super().__init__() self.out_channels = in_channels @@ -470,7 +472,7 @@ def forward( txt_ids: torch.Tensor = None, guidance: torch.Tensor = None, return_dict: bool = True, - ) -> Union[torch.FloatTensor, Transformer2DModelOutput]: + ) -> torch.FloatTensor | Transformer2DModelOutput: """ The forward method. @@ -499,11 +501,7 @@ def forward( encoder_hidden_states = self.context_embedder(encoder_hidden_states) ids = torch.cat((txt_ids, img_ids), dim=0) - if is_torch_npu_available(): - freqs_cos, freqs_sin = self.pos_embed(ids.cpu()) - image_rotary_emb = (freqs_cos.npu(), freqs_sin.npu()) - else: - image_rotary_emb = self.pos_embed(ids) + image_rotary_emb = self.pos_embed(ids) for index_block, block in enumerate(self.transformer_blocks): if torch.is_grad_enabled() and self.gradient_checkpointing and self.use_checkpoint[index_block]: diff --git a/src/diffusers/models/transformers/transformer_ltx.py b/src/diffusers/models/transformers/transformer_ltx.py index 685c73c07c75..0034d636761b 100644 --- a/src/diffusers/models/transformers/transformer_ltx.py +++ b/src/diffusers/models/transformers/transformer_ltx.py @@ -15,14 +15,14 @@ import inspect import math -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, deprecate, is_torch_version, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, deprecate, is_torch_version, logging from ...utils.torch_utils import maybe_allow_in_graph from .._modeling_parallel import ContextParallelInput, ContextParallelOutput from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward @@ -64,9 +64,9 @@ def __call__( self, attn: "LTXAttention", hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: batch_size, sequence_length, _ = ( hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape @@ -124,7 +124,7 @@ def __init__( dim_head: int = 64, dropout: float = 0.0, bias: bool = True, - cross_attention_dim: Optional[int] = None, + cross_attention_dim: int | None = None, out_bias: bool = True, qk_norm: str = "rms_norm_across_heads", processor=None, @@ -161,9 +161,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) @@ -203,7 +203,7 @@ def _prepare_video_coords( num_frames: int, height: int, width: int, - rope_interpolation_scale: Tuple[torch.Tensor, float, float], + rope_interpolation_scale: tuple[torch.Tensor, float, float], device: torch.device, ) -> torch.Tensor: # Always compute rope in fp32 @@ -226,12 +226,12 @@ def _prepare_video_coords( def forward( self, hidden_states: torch.Tensor, - num_frames: Optional[int] = None, - height: Optional[int] = None, - width: Optional[int] = None, - rope_interpolation_scale: Optional[Tuple[torch.Tensor, float, float]] = None, - video_coords: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + num_frames: int | None = None, + height: int | None = None, + width: int | None = None, + rope_interpolation_scale: tuple[torch.Tensor, float, float] | None = None, + video_coords: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: batch_size = hidden_states.size(0) if video_coords is None: @@ -346,8 +346,8 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + encoder_attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: batch_size = hidden_states.size(0) norm_hidden_states = self.norm1(hidden_states) @@ -491,35 +491,21 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, timestep: torch.LongTensor, encoder_attention_mask: torch.Tensor, - num_frames: Optional[int] = None, - height: Optional[int] = None, - width: Optional[int] = None, - rope_interpolation_scale: Optional[Union[Tuple[float, float, float], torch.Tensor]] = None, - video_coords: Optional[torch.Tensor] = None, - attention_kwargs: Optional[Dict[str, Any]] = None, + num_frames: int | None = None, + height: int | None = None, + width: int | None = None, + rope_interpolation_scale: tuple[float, float, float] | torch.Tensor | None = None, + video_coords: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, ) -> torch.Tensor: - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) - image_rotary_emb = self.rope(hidden_states, num_frames, height, width, rope_interpolation_scale, video_coords) # convert encoder_attention_mask to a bias the same way we do for attention_mask @@ -568,10 +554,6 @@ def forward( hidden_states = hidden_states * (1 + scale) + shift output = self.proj_out(hidden_states) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) return Transformer2DModelOutput(sample=output) diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py new file mode 100644 index 000000000000..db949ca34a1f --- /dev/null +++ b/src/diffusers/models/transformers/transformer_ltx2.py @@ -0,0 +1,1323 @@ +# Copyright 2025 The Lightricks team and The HuggingFace Team. +# All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin +from ...utils import BaseOutput, apply_lora_scale, is_torch_version, logging +from .._modeling_parallel import ContextParallelInput, ContextParallelOutput +from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward +from ..attention_dispatch import dispatch_attention_fn +from ..cache_utils import CacheMixin +from ..embeddings import PixArtAlphaCombinedTimestepSizeEmbeddings, PixArtAlphaTextProjection +from ..modeling_utils import ModelMixin +from ..normalization import RMSNorm + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def apply_interleaved_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor: + cos, sin = freqs + x_real, x_imag = x.unflatten(2, (-1, 2)).unbind(-1) # [B, S, C // 2] + x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2) + out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) + return out + + +def apply_split_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor: + cos, sin = freqs + + x_dtype = x.dtype + needs_reshape = False + if x.ndim != 4 and cos.ndim == 4: + # cos is (b, h, t, r) -> reshape x to (b, h, t, dim_per_head) + b, h, t, _ = cos.shape + x = x.reshape(b, t, h, -1).swapaxes(1, 2) + needs_reshape = True + + # Split last dim (2*r) into (d=2, r) + last = x.shape[-1] + if last % 2 != 0: + raise ValueError(f"Expected x.shape[-1] to be even for split rotary, got {last}.") + r = last // 2 + + # (..., 2, r) + split_x = x.reshape(*x.shape[:-1], 2, r).float() # Explicitly upcast to float + first_x = split_x[..., :1, :] # (..., 1, r) + second_x = split_x[..., 1:, :] # (..., 1, r) + + cos_u = cos.unsqueeze(-2) # broadcast to (..., 1, r) against (..., 2, r) + sin_u = sin.unsqueeze(-2) + + out = split_x * cos_u + first_out = out[..., :1, :] + second_out = out[..., 1:, :] + + first_out.addcmul_(-sin_u, second_x) + second_out.addcmul_(sin_u, first_x) + + out = out.reshape(*out.shape[:-2], last) + + if needs_reshape: + out = out.swapaxes(1, 2).reshape(b, t, -1) + + out = out.to(dtype=x_dtype) + return out + + +@dataclass +class AudioVisualModelOutput(BaseOutput): + r""" + Holds the output of an audiovisual model which produces both visual (e.g. video) and audio outputs. + + Args: + sample (`torch.Tensor` of shape `(batch_size, num_channels, num_frames, height, width)`): + The hidden states output conditioned on the `encoder_hidden_states` input, representing the visual output + of the model. This is typically a video (spatiotemporal) output. + audio_sample (`torch.Tensor` of shape `(batch_size, TODO)`): + The audio output of the audiovisual model. + """ + + sample: "torch.Tensor" # noqa: F821 + audio_sample: "torch.Tensor" # noqa: F821 + + +class LTX2AdaLayerNormSingle(nn.Module): + r""" + Norm layer adaptive layer norm single (adaLN-single). + + As proposed in PixArt-Alpha (see: https://huggingface.co/papers/2310.00426; Section 2.3) and adapted by the LTX-2.0 + model. In particular, the number of modulation parameters to be calculated is now configurable. + + Parameters: + embedding_dim (`int`): The size of each embedding vector. + num_mod_params (`int`, *optional*, defaults to `6`): + The number of modulation parameters which will be calculated in the first return argument. The default of 6 + is standard, but sometimes we may want to have a different (usually smaller) number of modulation + parameters. + use_additional_conditions (`bool`, *optional*, defaults to `False`): + Whether to use additional conditions for normalization or not. + """ + + def __init__(self, embedding_dim: int, num_mod_params: int = 6, use_additional_conditions: bool = False): + super().__init__() + self.num_mod_params = num_mod_params + + self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings( + embedding_dim, size_emb_dim=embedding_dim // 3, use_additional_conditions=use_additional_conditions + ) + + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, self.num_mod_params * embedding_dim, bias=True) + + def forward( + self, + timestep: torch.Tensor, + added_cond_kwargs: dict[str, torch.Tensor] | None = None, + batch_size: int | None = None, + hidden_dtype: torch.dtype | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + # No modulation happening here. + added_cond_kwargs = added_cond_kwargs or {"resolution": None, "aspect_ratio": None} + embedded_timestep = self.emb(timestep, **added_cond_kwargs, batch_size=batch_size, hidden_dtype=hidden_dtype) + return self.linear(self.silu(embedded_timestep)), embedded_timestep + + +class LTX2AudioVideoAttnProcessor: + r""" + Processor for implementing attention (SDPA is used by default if you're using PyTorch 2.0) for the LTX-2.0 model. + Compared to the LTX-1.0 model, we allow the RoPE embeddings for the queries and keys to be separate so that we can + support audio-to-video (a2v) and video-to-audio (v2a) cross attention. + """ + + _attention_backend = None + _parallel_config = None + + def __init__(self): + if is_torch_version("<", "2.0"): + raise ValueError( + "LTX attention processors require a minimum PyTorch version of 2.0. Please upgrade your PyTorch installation." + ) + + def __call__( + self, + attn: "LTX2Attention", + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + query_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + key_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + batch_size, sequence_length, _ = ( + hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + ) + + if attention_mask is not None: + attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) + + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + + query = attn.to_q(hidden_states) + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + query = attn.norm_q(query) + key = attn.norm_k(key) + + if query_rotary_emb is not None: + if attn.rope_type == "interleaved": + query = apply_interleaved_rotary_emb(query, query_rotary_emb) + key = apply_interleaved_rotary_emb( + key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb + ) + elif attn.rope_type == "split": + query = apply_split_rotary_emb(query, query_rotary_emb) + key = apply_split_rotary_emb(key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb) + + query = query.unflatten(2, (attn.heads, -1)) + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(query.dtype) + + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + return hidden_states + + +class LTX2Attention(torch.nn.Module, AttentionModuleMixin): + r""" + Attention class for all LTX-2.0 attention layers. Compared to LTX-1.0, this supports specifying the query and key + RoPE embeddings separately for audio-to-video (a2v) and video-to-audio (v2a) cross-attention. + """ + + _default_processor_cls = LTX2AudioVideoAttnProcessor + _available_processors = [LTX2AudioVideoAttnProcessor] + + def __init__( + self, + query_dim: int, + heads: int = 8, + kv_heads: int = 8, + dim_head: int = 64, + dropout: float = 0.0, + bias: bool = True, + cross_attention_dim: int | None = None, + out_bias: bool = True, + qk_norm: str = "rms_norm_across_heads", + norm_eps: float = 1e-6, + norm_elementwise_affine: bool = True, + rope_type: str = "interleaved", + processor=None, + ): + super().__init__() + if qk_norm != "rms_norm_across_heads": + raise NotImplementedError("Only 'rms_norm_across_heads' is supported as a valid value for `qk_norm`.") + + self.head_dim = dim_head + self.inner_dim = dim_head * heads + self.inner_kv_dim = self.inner_dim if kv_heads is None else dim_head * kv_heads + self.query_dim = query_dim + self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim + self.use_bias = bias + self.dropout = dropout + self.out_dim = query_dim + self.heads = heads + self.rope_type = rope_type + + self.norm_q = torch.nn.RMSNorm(dim_head * heads, eps=norm_eps, elementwise_affine=norm_elementwise_affine) + self.norm_k = torch.nn.RMSNorm(dim_head * kv_heads, eps=norm_eps, elementwise_affine=norm_elementwise_affine) + self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=bias) + self.to_k = torch.nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias) + self.to_v = torch.nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias) + self.to_out = torch.nn.ModuleList([]) + self.to_out.append(torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias)) + self.to_out.append(torch.nn.Dropout(dropout)) + + if processor is None: + processor = self._default_processor_cls() + self.set_processor(processor) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + query_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + key_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs, + ) -> torch.Tensor: + attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) + unused_kwargs = [k for k, _ in kwargs.items() if k not in attn_parameters] + if len(unused_kwargs) > 0: + logger.warning( + f"attention_kwargs {unused_kwargs} are not expected by {self.processor.__class__.__name__} and will be ignored." + ) + kwargs = {k: w for k, w in kwargs.items() if k in attn_parameters} + hidden_states = self.processor( + self, hidden_states, encoder_hidden_states, attention_mask, query_rotary_emb, key_rotary_emb, **kwargs + ) + return hidden_states + + +class LTX2VideoTransformerBlock(nn.Module): + r""" + Transformer block used in [LTX-2.0](https://huggingface.co/Lightricks/LTX-Video). + + Args: + dim (`int`): + The number of channels in the input and output. + num_attention_heads (`int`): + The number of heads to use for multi-head attention. + attention_head_dim (`int`): + The number of channels in each head. + qk_norm (`str`, defaults to `"rms_norm"`): + The normalization layer to use. + activation_fn (`str`, defaults to `"gelu-approximate"`): + Activation function to use in feed-forward. + eps (`float`, defaults to `1e-6`): + Epsilon value for normalization layers. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + cross_attention_dim: int, + audio_dim: int, + audio_num_attention_heads: int, + audio_attention_head_dim, + audio_cross_attention_dim: int, + qk_norm: str = "rms_norm_across_heads", + activation_fn: str = "gelu-approximate", + attention_bias: bool = True, + attention_out_bias: bool = True, + eps: float = 1e-6, + elementwise_affine: bool = False, + rope_type: str = "interleaved", + ): + super().__init__() + + # 1. Self-Attention (video and audio) + self.norm1 = RMSNorm(dim, eps=eps, elementwise_affine=elementwise_affine) + self.attn1 = LTX2Attention( + query_dim=dim, + heads=num_attention_heads, + kv_heads=num_attention_heads, + dim_head=attention_head_dim, + bias=attention_bias, + cross_attention_dim=None, + out_bias=attention_out_bias, + qk_norm=qk_norm, + rope_type=rope_type, + ) + + self.audio_norm1 = RMSNorm(audio_dim, eps=eps, elementwise_affine=elementwise_affine) + self.audio_attn1 = LTX2Attention( + query_dim=audio_dim, + heads=audio_num_attention_heads, + kv_heads=audio_num_attention_heads, + dim_head=audio_attention_head_dim, + bias=attention_bias, + cross_attention_dim=None, + out_bias=attention_out_bias, + qk_norm=qk_norm, + rope_type=rope_type, + ) + + # 2. Prompt Cross-Attention + self.norm2 = RMSNorm(dim, eps=eps, elementwise_affine=elementwise_affine) + self.attn2 = LTX2Attention( + query_dim=dim, + cross_attention_dim=cross_attention_dim, + heads=num_attention_heads, + kv_heads=num_attention_heads, + dim_head=attention_head_dim, + bias=attention_bias, + out_bias=attention_out_bias, + qk_norm=qk_norm, + rope_type=rope_type, + ) + + self.audio_norm2 = RMSNorm(audio_dim, eps=eps, elementwise_affine=elementwise_affine) + self.audio_attn2 = LTX2Attention( + query_dim=audio_dim, + cross_attention_dim=audio_cross_attention_dim, + heads=audio_num_attention_heads, + kv_heads=audio_num_attention_heads, + dim_head=audio_attention_head_dim, + bias=attention_bias, + out_bias=attention_out_bias, + qk_norm=qk_norm, + rope_type=rope_type, + ) + + # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention + # Audio-to-Video (a2v) Attention --> Q: Video; K,V: Audio + self.audio_to_video_norm = RMSNorm(dim, eps=eps, elementwise_affine=elementwise_affine) + self.audio_to_video_attn = LTX2Attention( + query_dim=dim, + cross_attention_dim=audio_dim, + heads=audio_num_attention_heads, + kv_heads=audio_num_attention_heads, + dim_head=audio_attention_head_dim, + bias=attention_bias, + out_bias=attention_out_bias, + qk_norm=qk_norm, + rope_type=rope_type, + ) + + # Video-to-Audio (v2a) Attention --> Q: Audio; K,V: Video + self.video_to_audio_norm = RMSNorm(audio_dim, eps=eps, elementwise_affine=elementwise_affine) + self.video_to_audio_attn = LTX2Attention( + query_dim=audio_dim, + cross_attention_dim=dim, + heads=audio_num_attention_heads, + kv_heads=audio_num_attention_heads, + dim_head=audio_attention_head_dim, + bias=attention_bias, + out_bias=attention_out_bias, + qk_norm=qk_norm, + rope_type=rope_type, + ) + + # 4. Feedforward layers + self.norm3 = RMSNorm(dim, eps=eps, elementwise_affine=elementwise_affine) + self.ff = FeedForward(dim, activation_fn=activation_fn) + + self.audio_norm3 = RMSNorm(audio_dim, eps=eps, elementwise_affine=elementwise_affine) + self.audio_ff = FeedForward(audio_dim, activation_fn=activation_fn) + + # 5. Per-Layer Modulation Parameters + # Self-Attention / Feedforward AdaLayerNorm-Zero mod params + self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) + self.audio_scale_shift_table = nn.Parameter(torch.randn(6, audio_dim) / audio_dim**0.5) + + # Per-layer a2v, v2a Cross-Attention mod params + self.video_a2v_cross_attn_scale_shift_table = nn.Parameter(torch.randn(5, dim)) + self.audio_a2v_cross_attn_scale_shift_table = nn.Parameter(torch.randn(5, audio_dim)) + + def forward( + self, + hidden_states: torch.Tensor, + audio_hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + audio_encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + temb_audio: torch.Tensor, + temb_ca_scale_shift: torch.Tensor, + temb_ca_audio_scale_shift: torch.Tensor, + temb_ca_gate: torch.Tensor, + temb_ca_audio_gate: torch.Tensor, + video_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + audio_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + ca_video_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + ca_audio_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + encoder_attention_mask: torch.Tensor | None = None, + audio_encoder_attention_mask: torch.Tensor | None = None, + a2v_cross_attention_mask: torch.Tensor | None = None, + v2a_cross_attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + batch_size = hidden_states.size(0) + + # 1. Video and Audio Self-Attention + norm_hidden_states = self.norm1(hidden_states) + + num_ada_params = self.scale_shift_table.shape[0] + ada_values = self.scale_shift_table[None, None].to(temb.device) + temb.reshape( + batch_size, temb.size(1), num_ada_params, -1 + ) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada_values.unbind(dim=2) + norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa + + attn_hidden_states = self.attn1( + hidden_states=norm_hidden_states, + encoder_hidden_states=None, + query_rotary_emb=video_rotary_emb, + ) + hidden_states = hidden_states + attn_hidden_states * gate_msa + + norm_audio_hidden_states = self.audio_norm1(audio_hidden_states) + + num_audio_ada_params = self.audio_scale_shift_table.shape[0] + audio_ada_values = self.audio_scale_shift_table[None, None].to(temb_audio.device) + temb_audio.reshape( + batch_size, temb_audio.size(1), num_audio_ada_params, -1 + ) + audio_shift_msa, audio_scale_msa, audio_gate_msa, audio_shift_mlp, audio_scale_mlp, audio_gate_mlp = ( + audio_ada_values.unbind(dim=2) + ) + norm_audio_hidden_states = norm_audio_hidden_states * (1 + audio_scale_msa) + audio_shift_msa + + attn_audio_hidden_states = self.audio_attn1( + hidden_states=norm_audio_hidden_states, + encoder_hidden_states=None, + query_rotary_emb=audio_rotary_emb, + ) + audio_hidden_states = audio_hidden_states + attn_audio_hidden_states * audio_gate_msa + + # 2. Video and Audio Cross-Attention with the text embeddings + norm_hidden_states = self.norm2(hidden_states) + attn_hidden_states = self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + query_rotary_emb=None, + attention_mask=encoder_attention_mask, + ) + hidden_states = hidden_states + attn_hidden_states + + norm_audio_hidden_states = self.audio_norm2(audio_hidden_states) + attn_audio_hidden_states = self.audio_attn2( + norm_audio_hidden_states, + encoder_hidden_states=audio_encoder_hidden_states, + query_rotary_emb=None, + attention_mask=audio_encoder_attention_mask, + ) + audio_hidden_states = audio_hidden_states + attn_audio_hidden_states + + # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention + norm_hidden_states = self.audio_to_video_norm(hidden_states) + norm_audio_hidden_states = self.video_to_audio_norm(audio_hidden_states) + + # Combine global and per-layer cross attention modulation parameters + # Video + video_per_layer_ca_scale_shift = self.video_a2v_cross_attn_scale_shift_table[:4, :] + video_per_layer_ca_gate = self.video_a2v_cross_attn_scale_shift_table[4:, :] + + video_ca_scale_shift_table = ( + video_per_layer_ca_scale_shift[:, :, ...].to(temb_ca_scale_shift.dtype) + + temb_ca_scale_shift.reshape(batch_size, temb_ca_scale_shift.shape[1], 4, -1) + ).unbind(dim=2) + video_ca_gate = ( + video_per_layer_ca_gate[:, :, ...].to(temb_ca_gate.dtype) + + temb_ca_gate.reshape(batch_size, temb_ca_gate.shape[1], 1, -1) + ).unbind(dim=2) + + video_a2v_ca_scale, video_a2v_ca_shift, video_v2a_ca_scale, video_v2a_ca_shift = video_ca_scale_shift_table + a2v_gate = video_ca_gate[0].squeeze(2) + + # Audio + audio_per_layer_ca_scale_shift = self.audio_a2v_cross_attn_scale_shift_table[:4, :] + audio_per_layer_ca_gate = self.audio_a2v_cross_attn_scale_shift_table[4:, :] + + audio_ca_scale_shift_table = ( + audio_per_layer_ca_scale_shift[:, :, ...].to(temb_ca_audio_scale_shift.dtype) + + temb_ca_audio_scale_shift.reshape(batch_size, temb_ca_audio_scale_shift.shape[1], 4, -1) + ).unbind(dim=2) + audio_ca_gate = ( + audio_per_layer_ca_gate[:, :, ...].to(temb_ca_audio_gate.dtype) + + temb_ca_audio_gate.reshape(batch_size, temb_ca_audio_gate.shape[1], 1, -1) + ).unbind(dim=2) + + audio_a2v_ca_scale, audio_a2v_ca_shift, audio_v2a_ca_scale, audio_v2a_ca_shift = audio_ca_scale_shift_table + v2a_gate = audio_ca_gate[0].squeeze(2) + + # Audio-to-Video Cross Attention: Q: Video; K,V: Audio + mod_norm_hidden_states = norm_hidden_states * (1 + video_a2v_ca_scale.squeeze(2)) + video_a2v_ca_shift.squeeze( + 2 + ) + mod_norm_audio_hidden_states = norm_audio_hidden_states * ( + 1 + audio_a2v_ca_scale.squeeze(2) + ) + audio_a2v_ca_shift.squeeze(2) + + a2v_attn_hidden_states = self.audio_to_video_attn( + mod_norm_hidden_states, + encoder_hidden_states=mod_norm_audio_hidden_states, + query_rotary_emb=ca_video_rotary_emb, + key_rotary_emb=ca_audio_rotary_emb, + attention_mask=a2v_cross_attention_mask, + ) + + hidden_states = hidden_states + a2v_gate * a2v_attn_hidden_states + + # Video-to-Audio Cross Attention: Q: Audio; K,V: Video + mod_norm_hidden_states = norm_hidden_states * (1 + video_v2a_ca_scale.squeeze(2)) + video_v2a_ca_shift.squeeze( + 2 + ) + mod_norm_audio_hidden_states = norm_audio_hidden_states * ( + 1 + audio_v2a_ca_scale.squeeze(2) + ) + audio_v2a_ca_shift.squeeze(2) + + v2a_attn_hidden_states = self.video_to_audio_attn( + mod_norm_audio_hidden_states, + encoder_hidden_states=mod_norm_hidden_states, + query_rotary_emb=ca_audio_rotary_emb, + key_rotary_emb=ca_video_rotary_emb, + attention_mask=v2a_cross_attention_mask, + ) + + audio_hidden_states = audio_hidden_states + v2a_gate * v2a_attn_hidden_states + + # 4. Feedforward + norm_hidden_states = self.norm3(hidden_states) * (1 + scale_mlp) + shift_mlp + ff_output = self.ff(norm_hidden_states) + hidden_states = hidden_states + ff_output * gate_mlp + + norm_audio_hidden_states = self.audio_norm3(audio_hidden_states) * (1 + audio_scale_mlp) + audio_shift_mlp + audio_ff_output = self.audio_ff(norm_audio_hidden_states) + audio_hidden_states = audio_hidden_states + audio_ff_output * audio_gate_mlp + + return hidden_states, audio_hidden_states + + +class LTX2AudioVideoRotaryPosEmbed(nn.Module): + """ + Video and audio rotary positional embeddings (RoPE) for the LTX-2.0 model. + + Args: + causal_offset (`int`, *optional*, defaults to `1`): + Offset in the temporal axis for causal VAE modeling. This is typically 1 (for causal modeling where the VAE + treats the very first frame differently), but could also be 0 (for non-causal modeling). + """ + + def __init__( + self, + dim: int, + patch_size: int = 1, + patch_size_t: int = 1, + base_num_frames: int = 20, + base_height: int = 2048, + base_width: int = 2048, + sampling_rate: int = 16000, + hop_length: int = 160, + scale_factors: tuple[int, ...] = (8, 32, 32), + theta: float = 10000.0, + causal_offset: int = 1, + modality: str = "video", + double_precision: bool = True, + rope_type: str = "interleaved", + num_attention_heads: int = 32, + ) -> None: + super().__init__() + + self.dim = dim + self.patch_size = patch_size + self.patch_size_t = patch_size_t + + if rope_type not in ["interleaved", "split"]: + raise ValueError(f"{rope_type=} not supported. Choose between 'interleaved' and 'split'.") + self.rope_type = rope_type + + self.base_num_frames = base_num_frames + self.num_attention_heads = num_attention_heads + + # Video-specific + self.base_height = base_height + self.base_width = base_width + + # Audio-specific + self.sampling_rate = sampling_rate + self.hop_length = hop_length + self.audio_latents_per_second = float(sampling_rate) / float(hop_length) / float(scale_factors[0]) + + self.scale_factors = scale_factors + self.theta = theta + self.causal_offset = causal_offset + + self.modality = modality + if self.modality not in ["video", "audio"]: + raise ValueError(f"Modality {modality} is not supported. Supported modalities are `video` and `audio`.") + self.double_precision = double_precision + + def prepare_video_coords( + self, + batch_size: int, + num_frames: int, + height: int, + width: int, + device: torch.device, + fps: float = 24.0, + ) -> torch.Tensor: + """ + Create per-dimension bounds [inclusive start, exclusive end) for each patch with respect to the original pixel + space video grid (num_frames, height, width). This will ultimately have shape (batch_size, 3, num_patches, 2) + where + - axis 1 (size 3) enumerates (frame, height, width) dimensions (e.g. idx 0 corresponds to frames) + - axis 3 (size 2) stores `[start, end)` indices within each dimension + + Args: + batch_size (`int`): + Batch size of the video latents. + num_frames (`int`): + Number of latent frames in the video latents. + height (`int`): + Latent height of the video latents. + width (`int`): + Latent width of the video latents. + device (`torch.device`): + Device on which to create the video grid. + + Returns: + `torch.Tensor`: + Per-dimension patch boundaries tensor of shape [batch_size, 3, num_patches, 2]. + """ + + # 1. Generate grid coordinates for each spatiotemporal dimension (frames, height, width) + # Always compute rope in fp32 + grid_f = torch.arange(start=0, end=num_frames, step=self.patch_size_t, dtype=torch.float32, device=device) + grid_h = torch.arange(start=0, end=height, step=self.patch_size, dtype=torch.float32, device=device) + grid_w = torch.arange(start=0, end=width, step=self.patch_size, dtype=torch.float32, device=device) + # indexing='ij' ensures that the dimensions are kept in order as (frames, height, width) + grid = torch.meshgrid(grid_f, grid_h, grid_w, indexing="ij") + grid = torch.stack(grid, dim=0) # [3, N_F, N_H, N_W], where e.g. N_F is the number of temporal patches + + # 2. Get the patch boundaries with respect to the latent video grid + patch_size = (self.patch_size_t, self.patch_size, self.patch_size) + patch_size_delta = torch.tensor(patch_size, dtype=grid.dtype, device=grid.device) + patch_ends = grid + patch_size_delta.view(3, 1, 1, 1) + + # Combine the start (grid) and end (patch_ends) coordinates along new trailing dimension + latent_coords = torch.stack([grid, patch_ends], dim=-1) # [3, N_F, N_H, N_W, 2] + # Reshape to (batch_size, 3, num_patches, 2) + latent_coords = latent_coords.flatten(1, 3) + latent_coords = latent_coords.unsqueeze(0).repeat(batch_size, 1, 1, 1) + + # 3. Calculate the pixel space patch boundaries from the latent boundaries. + scale_tensor = torch.tensor(self.scale_factors, device=latent_coords.device) + # Broadcast the VAE scale factors such that they are compatible with latent_coords's shape + broadcast_shape = [1] * latent_coords.ndim + broadcast_shape[1] = -1 # This is the (frame, height, width) dim + # Apply per-axis scaling to convert latent coordinates to pixel space coordinates + pixel_coords = latent_coords * scale_tensor.view(*broadcast_shape) + + # As the VAE temporal stride for the first frame is 1 instead of self.vae_scale_factors[0], we need to shift + # and clamp to keep the first-frame timestamps causal and non-negative. + pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + self.causal_offset - self.scale_factors[0]).clamp(min=0) + + # Scale the temporal coordinates by the video FPS + pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / fps + + return pixel_coords + + def prepare_audio_coords( + self, + batch_size: int, + num_frames: int, + device: torch.device, + shift: int = 0, + ) -> torch.Tensor: + """ + Create per-dimension bounds [inclusive start, exclusive end) of start and end timestamps for each latent frame. + This will ultimately have shape (batch_size, 3, num_patches, 2) where + - axis 1 (size 1) represents the temporal dimension + - axis 3 (size 2) stores `[start, end)` indices within each dimension + + Args: + batch_size (`int`): + Batch size of the audio latents. + num_frames (`int`): + Number of latent frames in the audio latents. + device (`torch.device`): + Device on which to create the audio grid. + shift (`int`, *optional*, defaults to `0`): + Offset on the latent indices. Different shift values correspond to different overlapping windows with + respect to the same underlying latent grid. + + Returns: + `torch.Tensor`: + Per-dimension patch boundaries tensor of shape [batch_size, 1, num_patches, 2]. + """ + + # 1. Generate coordinates in the frame (time) dimension. + # Always compute rope in fp32 + grid_f = torch.arange( + start=shift, end=num_frames + shift, step=self.patch_size_t, dtype=torch.float32, device=device + ) + + # 2. Calculate start timstamps in seconds with respect to the original spectrogram grid + audio_scale_factor = self.scale_factors[0] + # Scale back to mel spectrogram space + grid_start_mel = grid_f * audio_scale_factor + # Handle first frame causal offset, ensuring non-negative timestamps + grid_start_mel = (grid_start_mel + self.causal_offset - audio_scale_factor).clip(min=0) + # Convert mel bins back into seconds + grid_start_s = grid_start_mel * self.hop_length / self.sampling_rate + + # 3. Calculate start timstamps in seconds with respect to the original spectrogram grid + grid_end_mel = (grid_f + self.patch_size_t) * audio_scale_factor + grid_end_mel = (grid_end_mel + self.causal_offset - audio_scale_factor).clip(min=0) + grid_end_s = grid_end_mel * self.hop_length / self.sampling_rate + + audio_coords = torch.stack([grid_start_s, grid_end_s], dim=-1) # [num_patches, 2] + audio_coords = audio_coords.unsqueeze(0).expand(batch_size, -1, -1) # [batch_size, num_patches, 2] + audio_coords = audio_coords.unsqueeze(1) # [batch_size, 1, num_patches, 2] + return audio_coords + + def prepare_coords(self, *args, **kwargs): + if self.modality == "video": + return self.prepare_video_coords(*args, **kwargs) + elif self.modality == "audio": + return self.prepare_audio_coords(*args, **kwargs) + + def forward( + self, coords: torch.Tensor, device: str | torch.device | None = None + ) -> tuple[torch.Tensor, torch.Tensor]: + device = device or coords.device + + # Number of spatiotemporal dimensions (3 for video, 1 (temporal) for audio and cross attn) + num_pos_dims = coords.shape[1] + + # 1. If the coords are patch boundaries [start, end), use the midpoint of these boundaries as the patch + # position index + if coords.ndim == 4: + coords_start, coords_end = coords.chunk(2, dim=-1) + coords = (coords_start + coords_end) / 2.0 + coords = coords.squeeze(-1) # [B, num_pos_dims, num_patches] + + # 2. Get coordinates as a fraction of the base data shape + if self.modality == "video": + max_positions = (self.base_num_frames, self.base_height, self.base_width) + elif self.modality == "audio": + max_positions = (self.base_num_frames,) + # [B, num_pos_dims, num_patches] --> [B, num_patches, num_pos_dims] + grid = torch.stack([coords[:, i] / max_positions[i] for i in range(num_pos_dims)], dim=-1).to(device) + # Number of spatiotemporal dimensions (3 for video, 1 for audio and cross attn) times 2 for cos, sin + num_rope_elems = num_pos_dims * 2 + + # 3. Create a 1D grid of frequencies for RoPE + freqs_dtype = torch.float64 if self.double_precision else torch.float32 + pow_indices = torch.pow( + self.theta, + torch.linspace(start=0.0, end=1.0, steps=self.dim // num_rope_elems, dtype=freqs_dtype, device=device), + ) + freqs = (pow_indices * torch.pi / 2.0).to(dtype=torch.float32) + + # 4. Tensor-vector outer product between pos ids tensor of shape (B, 3, num_patches) and freqs vector of shape + # (self.dim // num_elems,) + freqs = (grid.unsqueeze(-1) * 2 - 1) * freqs # [B, num_patches, num_pos_dims, self.dim // num_elems] + freqs = freqs.transpose(-1, -2).flatten(2) # [B, num_patches, self.dim // 2] + + # 5. Get real, interleaved (cos, sin) frequencies, padded to self.dim + # TODO: consider implementing this as a utility and reuse in `connectors.py`. + # src/diffusers/pipelines/ltx2/connectors.py + if self.rope_type == "interleaved": + cos_freqs = freqs.cos().repeat_interleave(2, dim=-1) + sin_freqs = freqs.sin().repeat_interleave(2, dim=-1) + + if self.dim % num_rope_elems != 0: + cos_padding = torch.ones_like(cos_freqs[:, :, : self.dim % num_rope_elems]) + sin_padding = torch.zeros_like(cos_freqs[:, :, : self.dim % num_rope_elems]) + cos_freqs = torch.cat([cos_padding, cos_freqs], dim=-1) + sin_freqs = torch.cat([sin_padding, sin_freqs], dim=-1) + + elif self.rope_type == "split": + expected_freqs = self.dim // 2 + current_freqs = freqs.shape[-1] + pad_size = expected_freqs - current_freqs + cos_freq = freqs.cos() + sin_freq = freqs.sin() + + if pad_size != 0: + cos_padding = torch.ones_like(cos_freq[:, :, :pad_size]) + sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size]) + + cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1) + sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1) + + # Reshape freqs to be compatible with multi-head attention + b = cos_freq.shape[0] + t = cos_freq.shape[1] + + cos_freq = cos_freq.reshape(b, t, self.num_attention_heads, -1) + sin_freq = sin_freq.reshape(b, t, self.num_attention_heads, -1) + + cos_freqs = torch.swapaxes(cos_freq, 1, 2) # (B,H,T,D//2) + sin_freqs = torch.swapaxes(sin_freq, 1, 2) # (B,H,T,D//2) + + return cos_freqs, sin_freqs + + +class LTX2VideoTransformer3DModel( + ModelMixin, ConfigMixin, AttentionMixin, FromOriginalModelMixin, PeftAdapterMixin, CacheMixin +): + r""" + A Transformer model for video-like data used in [LTX](https://huggingface.co/Lightricks/LTX-Video). + + Args: + in_channels (`int`, defaults to `128`): + The number of channels in the input. + out_channels (`int`, defaults to `128`): + The number of channels in the output. + patch_size (`int`, defaults to `1`): + The size of the spatial patches to use in the patch embedding layer. + patch_size_t (`int`, defaults to `1`): + The size of the tmeporal patches to use in the patch embedding layer. + num_attention_heads (`int`, defaults to `32`): + The number of heads to use for multi-head attention. + attention_head_dim (`int`, defaults to `64`): + The number of channels in each head. + cross_attention_dim (`int`, defaults to `2048 `): + The number of channels for cross attention heads. + num_layers (`int`, defaults to `28`): + The number of layers of Transformer blocks to use. + activation_fn (`str`, defaults to `"gelu-approximate"`): + Activation function to use in feed-forward. + qk_norm (`str`, defaults to `"rms_norm_across_heads"`): + The normalization layer to use. + """ + + _supports_gradient_checkpointing = True + _skip_layerwise_casting_patterns = ["norm"] + _repeated_blocks = ["LTX2VideoTransformerBlock"] + _cp_plan = { + "": { + "hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), + "encoder_hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), + "encoder_attention_mask": ContextParallelInput(split_dim=1, expected_dims=2, split_output=False), + }, + "rope": { + 0: ContextParallelInput(split_dim=1, expected_dims=3, split_output=True), + 1: ContextParallelInput(split_dim=1, expected_dims=3, split_output=True), + }, + "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3), + } + + @register_to_config + def __init__( + self, + in_channels: int = 128, # Video Arguments + out_channels: int | None = 128, + patch_size: int = 1, + patch_size_t: int = 1, + num_attention_heads: int = 32, + attention_head_dim: int = 128, + cross_attention_dim: int = 4096, + vae_scale_factors: tuple[int, int, int] = (8, 32, 32), + pos_embed_max_pos: int = 20, + base_height: int = 2048, + base_width: int = 2048, + audio_in_channels: int = 128, # Audio Arguments + audio_out_channels: int | None = 128, + audio_patch_size: int = 1, + audio_patch_size_t: int = 1, + audio_num_attention_heads: int = 32, + audio_attention_head_dim: int = 64, + audio_cross_attention_dim: int = 2048, + audio_scale_factor: int = 4, + audio_pos_embed_max_pos: int = 20, + audio_sampling_rate: int = 16000, + audio_hop_length: int = 160, + num_layers: int = 48, # Shared arguments + activation_fn: str = "gelu-approximate", + qk_norm: str = "rms_norm_across_heads", + norm_elementwise_affine: bool = False, + norm_eps: float = 1e-6, + caption_channels: int = 3840, + attention_bias: bool = True, + attention_out_bias: bool = True, + rope_theta: float = 10000.0, + rope_double_precision: bool = True, + causal_offset: int = 1, + timestep_scale_multiplier: int = 1000, + cross_attn_timestep_scale_multiplier: int = 1000, + rope_type: str = "interleaved", + ) -> None: + super().__init__() + + out_channels = out_channels or in_channels + audio_out_channels = audio_out_channels or audio_in_channels + inner_dim = num_attention_heads * attention_head_dim + audio_inner_dim = audio_num_attention_heads * audio_attention_head_dim + + # 1. Patchification input projections + self.proj_in = nn.Linear(in_channels, inner_dim) + self.audio_proj_in = nn.Linear(audio_in_channels, audio_inner_dim) + + # 2. Prompt embeddings + self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim) + self.audio_caption_projection = PixArtAlphaTextProjection( + in_features=caption_channels, hidden_size=audio_inner_dim + ) + + # 3. Timestep Modulation Params and Embedding + # 3.1. Global Timestep Modulation Parameters (except for cross-attention) and timestep + size embedding + # time_embed and audio_time_embed calculate both the timestep embedding and (global) modulation parameters + self.time_embed = LTX2AdaLayerNormSingle(inner_dim, num_mod_params=6, use_additional_conditions=False) + self.audio_time_embed = LTX2AdaLayerNormSingle( + audio_inner_dim, num_mod_params=6, use_additional_conditions=False + ) + + # 3.2. Global Cross Attention Modulation Parameters + # Used in the audio-to-video and video-to-audio cross attention layers as a global set of modulation params, + # which are then further modified by per-block modulaton params in each transformer block. + # There are 2 sets of scale/shift parameters for each modality, 1 each for audio-to-video (a2v) and + # video-to-audio (v2a) cross attention + self.av_cross_attn_video_scale_shift = LTX2AdaLayerNormSingle( + inner_dim, num_mod_params=4, use_additional_conditions=False + ) + self.av_cross_attn_audio_scale_shift = LTX2AdaLayerNormSingle( + audio_inner_dim, num_mod_params=4, use_additional_conditions=False + ) + # Gate param for audio-to-video (a2v) cross attn (where the video is the queries (Q) and the audio is the keys + # and values (KV)) + self.av_cross_attn_video_a2v_gate = LTX2AdaLayerNormSingle( + inner_dim, num_mod_params=1, use_additional_conditions=False + ) + # Gate param for video-to-audio (v2a) cross attn (where the audio is the queries (Q) and the video is the keys + # and values (KV)) + self.av_cross_attn_audio_v2a_gate = LTX2AdaLayerNormSingle( + audio_inner_dim, num_mod_params=1, use_additional_conditions=False + ) + + # 3.3. Output Layer Scale/Shift Modulation parameters + self.scale_shift_table = nn.Parameter(torch.randn(2, inner_dim) / inner_dim**0.5) + self.audio_scale_shift_table = nn.Parameter(torch.randn(2, audio_inner_dim) / audio_inner_dim**0.5) + + # 4. Rotary Positional Embeddings (RoPE) + # Self-Attention + self.rope = LTX2AudioVideoRotaryPosEmbed( + dim=inner_dim, + patch_size=patch_size, + patch_size_t=patch_size_t, + base_num_frames=pos_embed_max_pos, + base_height=base_height, + base_width=base_width, + scale_factors=vae_scale_factors, + theta=rope_theta, + causal_offset=causal_offset, + modality="video", + double_precision=rope_double_precision, + rope_type=rope_type, + num_attention_heads=num_attention_heads, + ) + self.audio_rope = LTX2AudioVideoRotaryPosEmbed( + dim=audio_inner_dim, + patch_size=audio_patch_size, + patch_size_t=audio_patch_size_t, + base_num_frames=audio_pos_embed_max_pos, + sampling_rate=audio_sampling_rate, + hop_length=audio_hop_length, + scale_factors=[audio_scale_factor], + theta=rope_theta, + causal_offset=causal_offset, + modality="audio", + double_precision=rope_double_precision, + rope_type=rope_type, + num_attention_heads=audio_num_attention_heads, + ) + + # Audio-to-Video, Video-to-Audio Cross-Attention + cross_attn_pos_embed_max_pos = max(pos_embed_max_pos, audio_pos_embed_max_pos) + self.cross_attn_rope = LTX2AudioVideoRotaryPosEmbed( + dim=audio_cross_attention_dim, + patch_size=patch_size, + patch_size_t=patch_size_t, + base_num_frames=cross_attn_pos_embed_max_pos, + base_height=base_height, + base_width=base_width, + theta=rope_theta, + causal_offset=causal_offset, + modality="video", + double_precision=rope_double_precision, + rope_type=rope_type, + num_attention_heads=num_attention_heads, + ) + self.cross_attn_audio_rope = LTX2AudioVideoRotaryPosEmbed( + dim=audio_cross_attention_dim, + patch_size=audio_patch_size, + patch_size_t=audio_patch_size_t, + base_num_frames=cross_attn_pos_embed_max_pos, + sampling_rate=audio_sampling_rate, + hop_length=audio_hop_length, + theta=rope_theta, + causal_offset=causal_offset, + modality="audio", + double_precision=rope_double_precision, + rope_type=rope_type, + num_attention_heads=audio_num_attention_heads, + ) + + # 5. Transformer Blocks + self.transformer_blocks = nn.ModuleList( + [ + LTX2VideoTransformerBlock( + dim=inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + cross_attention_dim=cross_attention_dim, + audio_dim=audio_inner_dim, + audio_num_attention_heads=audio_num_attention_heads, + audio_attention_head_dim=audio_attention_head_dim, + audio_cross_attention_dim=audio_cross_attention_dim, + qk_norm=qk_norm, + activation_fn=activation_fn, + attention_bias=attention_bias, + attention_out_bias=attention_out_bias, + eps=norm_eps, + elementwise_affine=norm_elementwise_affine, + rope_type=rope_type, + ) + for _ in range(num_layers) + ] + ) + + # 6. Output layers + self.norm_out = nn.LayerNorm(inner_dim, eps=1e-6, elementwise_affine=False) + self.proj_out = nn.Linear(inner_dim, out_channels) + + self.audio_norm_out = nn.LayerNorm(audio_inner_dim, eps=1e-6, elementwise_affine=False) + self.audio_proj_out = nn.Linear(audio_inner_dim, audio_out_channels) + + self.gradient_checkpointing = False + + @apply_lora_scale("attention_kwargs") + def forward( + self, + hidden_states: torch.Tensor, + audio_hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + audio_encoder_hidden_states: torch.Tensor, + timestep: torch.LongTensor, + audio_timestep: torch.LongTensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + audio_encoder_attention_mask: torch.Tensor | None = None, + num_frames: int | None = None, + height: int | None = None, + width: int | None = None, + fps: float = 24.0, + audio_num_frames: int | None = None, + video_coords: torch.Tensor | None = None, + audio_coords: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, + return_dict: bool = True, + ) -> torch.Tensor: + """ + Forward pass for LTX-2.0 audiovisual video transformer. + + Args: + hidden_states (`torch.Tensor`): + Input patchified video latents of shape `(batch_size, num_video_tokens, in_channels)`. + audio_hidden_states (`torch.Tensor`): + Input patchified audio latents of shape `(batch_size, num_audio_tokens, audio_in_channels)`. + encoder_hidden_states (`torch.Tensor`): + Input video text embeddings of shape `(batch_size, text_seq_len, self.config.caption_channels)`. + audio_encoder_hidden_states (`torch.Tensor`): + Input audio text embeddings of shape `(batch_size, text_seq_len, self.config.caption_channels)`. + timestep (`torch.Tensor`): + Input timestep of shape `(batch_size, num_video_tokens)`. These should already be scaled by + `self.config.timestep_scale_multiplier`. + audio_timestep (`torch.Tensor`, *optional*): + Input timestep of shape `(batch_size,)` or `(batch_size, num_audio_tokens)` for audio modulation + params. This is only used by certain pipelines such as the I2V pipeline. + encoder_attention_mask (`torch.Tensor`, *optional*): + Optional multiplicative text attention mask of shape `(batch_size, text_seq_len)`. + audio_encoder_attention_mask (`torch.Tensor`, *optional*): + Optional multiplicative text attention mask of shape `(batch_size, text_seq_len)` for audio modeling. + num_frames (`int`, *optional*): + The number of latent video frames. Used if calculating the video coordinates for RoPE. + height (`int`, *optional*): + The latent video height. Used if calculating the video coordinates for RoPE. + width (`int`, *optional*): + The latent video width. Used if calculating the video coordinates for RoPE. + fps: (`float`, *optional*, defaults to `24.0`): + The desired frames per second of the generated video. Used if calculating the video coordinates for + RoPE. + audio_num_frames: (`int`, *optional*): + The number of latent audio frames. Used if calculating the audio coordinates for RoPE. + video_coords (`torch.Tensor`, *optional*): + The video coordinates to be used when calculating the rotary positional embeddings (RoPE) of shape + `(batch_size, 3, num_video_tokens, 2)`. If not supplied, this will be calculated inside `forward`. + audio_coords (`torch.Tensor`, *optional*): + The audio coordinates to be used when calculating the rotary positional embeddings (RoPE) of shape + `(batch_size, 1, num_audio_tokens, 2)`. If not supplied, this will be calculated inside `forward`. + attention_kwargs (`dict[str, Any]`, *optional*): + Optional dict of keyword args to be passed to the attention processor. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a dict-like structured output of type `AudioVisualModelOutput` or a tuple. + + Returns: + `AudioVisualModelOutput` or `tuple`: + If `return_dict` is `True`, returns a structured output of type `AudioVisualModelOutput`, otherwise a + `tuple` is returned where the first element is the denoised video latent patch sequence and the second + element is the denoised audio latent patch sequence. + """ + # Determine timestep for audio. + audio_timestep = audio_timestep if audio_timestep is not None else timestep + + # convert encoder_attention_mask to a bias the same way we do for attention_mask + if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2: + encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0 + encoder_attention_mask = encoder_attention_mask.unsqueeze(1) + + if audio_encoder_attention_mask is not None and audio_encoder_attention_mask.ndim == 2: + audio_encoder_attention_mask = (1 - audio_encoder_attention_mask.to(audio_hidden_states.dtype)) * -10000.0 + audio_encoder_attention_mask = audio_encoder_attention_mask.unsqueeze(1) + + batch_size = hidden_states.size(0) + + # 1. Prepare RoPE positional embeddings + if video_coords is None: + video_coords = self.rope.prepare_video_coords( + batch_size, num_frames, height, width, hidden_states.device, fps=fps + ) + if audio_coords is None: + audio_coords = self.audio_rope.prepare_audio_coords( + batch_size, audio_num_frames, audio_hidden_states.device + ) + + video_rotary_emb = self.rope(video_coords, device=hidden_states.device) + audio_rotary_emb = self.audio_rope(audio_coords, device=audio_hidden_states.device) + + video_cross_attn_rotary_emb = self.cross_attn_rope(video_coords[:, 0:1, :], device=hidden_states.device) + audio_cross_attn_rotary_emb = self.cross_attn_audio_rope( + audio_coords[:, 0:1, :], device=audio_hidden_states.device + ) + + # 2. Patchify input projections + hidden_states = self.proj_in(hidden_states) + audio_hidden_states = self.audio_proj_in(audio_hidden_states) + + # 3. Prepare timestep embeddings and modulation parameters + timestep_cross_attn_gate_scale_factor = ( + self.config.cross_attn_timestep_scale_multiplier / self.config.timestep_scale_multiplier + ) + + # 3.1. Prepare global modality (video and audio) timestep embedding and modulation parameters + # temb is used in the transformer blocks (as expected), while embedded_timestep is used for the output layer + # modulation with scale_shift_table (and similarly for audio) + temb, embedded_timestep = self.time_embed( + timestep.flatten(), + batch_size=batch_size, + hidden_dtype=hidden_states.dtype, + ) + temb = temb.view(batch_size, -1, temb.size(-1)) + embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.size(-1)) + + temb_audio, audio_embedded_timestep = self.audio_time_embed( + audio_timestep.flatten(), + batch_size=batch_size, + hidden_dtype=audio_hidden_states.dtype, + ) + temb_audio = temb_audio.view(batch_size, -1, temb_audio.size(-1)) + audio_embedded_timestep = audio_embedded_timestep.view(batch_size, -1, audio_embedded_timestep.size(-1)) + + # 3.2. Prepare global modality cross attention modulation parameters + video_cross_attn_scale_shift, _ = self.av_cross_attn_video_scale_shift( + timestep.flatten(), + batch_size=batch_size, + hidden_dtype=hidden_states.dtype, + ) + video_cross_attn_a2v_gate, _ = self.av_cross_attn_video_a2v_gate( + timestep.flatten() * timestep_cross_attn_gate_scale_factor, + batch_size=batch_size, + hidden_dtype=hidden_states.dtype, + ) + video_cross_attn_scale_shift = video_cross_attn_scale_shift.view( + batch_size, -1, video_cross_attn_scale_shift.shape[-1] + ) + video_cross_attn_a2v_gate = video_cross_attn_a2v_gate.view(batch_size, -1, video_cross_attn_a2v_gate.shape[-1]) + + audio_cross_attn_scale_shift, _ = self.av_cross_attn_audio_scale_shift( + audio_timestep.flatten(), + batch_size=batch_size, + hidden_dtype=audio_hidden_states.dtype, + ) + audio_cross_attn_v2a_gate, _ = self.av_cross_attn_audio_v2a_gate( + audio_timestep.flatten() * timestep_cross_attn_gate_scale_factor, + batch_size=batch_size, + hidden_dtype=audio_hidden_states.dtype, + ) + audio_cross_attn_scale_shift = audio_cross_attn_scale_shift.view( + batch_size, -1, audio_cross_attn_scale_shift.shape[-1] + ) + audio_cross_attn_v2a_gate = audio_cross_attn_v2a_gate.view(batch_size, -1, audio_cross_attn_v2a_gate.shape[-1]) + + # 4. Prepare prompt embeddings + encoder_hidden_states = self.caption_projection(encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1)) + + audio_encoder_hidden_states = self.audio_caption_projection(audio_encoder_hidden_states) + audio_encoder_hidden_states = audio_encoder_hidden_states.view(batch_size, -1, audio_hidden_states.size(-1)) + + # 5. Run transformer blocks + for block in self.transformer_blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states, audio_hidden_states = self._gradient_checkpointing_func( + block, + hidden_states, + audio_hidden_states, + encoder_hidden_states, + audio_encoder_hidden_states, + temb, + temb_audio, + video_cross_attn_scale_shift, + audio_cross_attn_scale_shift, + video_cross_attn_a2v_gate, + audio_cross_attn_v2a_gate, + video_rotary_emb, + audio_rotary_emb, + video_cross_attn_rotary_emb, + audio_cross_attn_rotary_emb, + encoder_attention_mask, + audio_encoder_attention_mask, + ) + else: + hidden_states, audio_hidden_states = block( + hidden_states=hidden_states, + audio_hidden_states=audio_hidden_states, + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + temb=temb, + temb_audio=temb_audio, + temb_ca_scale_shift=video_cross_attn_scale_shift, + temb_ca_audio_scale_shift=audio_cross_attn_scale_shift, + temb_ca_gate=video_cross_attn_a2v_gate, + temb_ca_audio_gate=audio_cross_attn_v2a_gate, + video_rotary_emb=video_rotary_emb, + audio_rotary_emb=audio_rotary_emb, + ca_video_rotary_emb=video_cross_attn_rotary_emb, + ca_audio_rotary_emb=audio_cross_attn_rotary_emb, + encoder_attention_mask=encoder_attention_mask, + audio_encoder_attention_mask=audio_encoder_attention_mask, + ) + + # 6. Output layers (including unpatchification) + scale_shift_values = self.scale_shift_table[None, None] + embedded_timestep[:, :, None] + shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1] + + hidden_states = self.norm_out(hidden_states) + hidden_states = hidden_states * (1 + scale) + shift + output = self.proj_out(hidden_states) + + audio_scale_shift_values = self.audio_scale_shift_table[None, None] + audio_embedded_timestep[:, :, None] + audio_shift, audio_scale = audio_scale_shift_values[:, :, 0], audio_scale_shift_values[:, :, 1] + + audio_hidden_states = self.audio_norm_out(audio_hidden_states) + audio_hidden_states = audio_hidden_states * (1 + audio_scale) + audio_shift + audio_output = self.audio_proj_out(audio_hidden_states) + + if not return_dict: + return (output, audio_output) + return AudioVisualModelOutput(sample=output, audio_sample=audio_output) diff --git a/src/diffusers/models/transformers/transformer_lumina2.py b/src/diffusers/models/transformers/transformer_lumina2.py index 77121edb9fc9..03e2841f8bcb 100644 --- a/src/diffusers/models/transformers/transformer_lumina2.py +++ b/src/diffusers/models/transformers/transformer_lumina2.py @@ -13,7 +13,7 @@ # limitations under the License. import math -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -22,7 +22,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin from ...loaders.single_file_model import FromOriginalModelMixin -from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, logging from ..attention import LuminaFeedForward from ..attention_processor import Attention from ..embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed @@ -58,7 +58,7 @@ def __init__( def forward( self, hidden_states: torch.Tensor, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: timestep_proj = self.time_proj(timestep).type_as(hidden_states) time_embed = self.timestep_embedder(timestep_proj) caption_embed = self.caption_embedder(encoder_hidden_states) @@ -80,9 +80,9 @@ def __call__( attn: Attention, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, - base_sequence_length: Optional[int] = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, + base_sequence_length: int | None = None, ) -> torch.Tensor: batch_size, sequence_length, _ = hidden_states.shape @@ -202,7 +202,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.Tensor, image_rotary_emb: torch.Tensor, - temb: Optional[torch.Tensor] = None, + temb: torch.Tensor | None = None, ) -> torch.Tensor: if self.modulation: norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb) @@ -231,7 +231,7 @@ def forward( class Lumina2RotaryPosEmbed(nn.Module): - def __init__(self, theta: int, axes_dim: List[int], axes_lens: List[int] = (300, 512, 512), patch_size: int = 2): + def __init__(self, theta: int, axes_dim: list[int], axes_lens: list[int] = (300, 512, 512), patch_size: int = 2): super().__init__() self.theta = theta self.axes_dim = axes_dim @@ -240,7 +240,7 @@ def __init__(self, theta: int, axes_dim: List[int], axes_lens: List[int] = (300, self.freqs_cis = self._precompute_freqs_cis(axes_dim, axes_lens, theta) - def _precompute_freqs_cis(self, axes_dim: List[int], axes_lens: List[int], theta: int) -> List[torch.Tensor]: + def _precompute_freqs_cis(self, axes_dim: list[int], axes_lens: list[int], theta: int) -> list[torch.Tensor]: freqs_cis = [] freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64 for i, (d, e) in enumerate(zip(axes_dim, axes_lens)): @@ -368,18 +368,18 @@ def __init__( sample_size: int = 128, patch_size: int = 2, in_channels: int = 16, - out_channels: Optional[int] = None, + out_channels: int | None = None, hidden_size: int = 2304, num_layers: int = 26, num_refiner_layers: int = 2, num_attention_heads: int = 24, num_kv_heads: int = 8, multiple_of: int = 256, - ffn_dim_multiplier: Optional[float] = None, + ffn_dim_multiplier: float | None = None, norm_eps: float = 1e-5, scaling_factor: float = 1.0, - axes_dim_rope: Tuple[int, int, int] = (32, 32, 32), - axes_lens: Tuple[int, int, int] = (300, 512, 512), + axes_dim_rope: tuple[int, int, int] = (32, 32, 32), + axes_lens: tuple[int, int, int] = (300, 512, 512), cap_feat_dim: int = 1024, ) -> None: super().__init__() @@ -455,30 +455,16 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, encoder_attention_mask: torch.Tensor, - attention_kwargs: Optional[Dict[str, Any]] = None, + attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - ) -> Union[torch.Tensor, Transformer2DModelOutput]: - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) - + ) -> torch.Tensor | Transformer2DModelOutput: # 1. Condition, positional & patch embedding batch_size, _, height, width = hidden_states.shape @@ -539,10 +525,6 @@ def forward( ) output = torch.stack(output, dim=0) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) return Transformer2DModelOutput(sample=output) diff --git a/src/diffusers/models/transformers/transformer_mochi.py b/src/diffusers/models/transformers/transformer_mochi.py index 63911fe7c10d..31106e3a0476 100644 --- a/src/diffusers/models/transformers/transformer_mochi.py +++ b/src/diffusers/models/transformers/transformer_mochi.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, Optional, Tuple +from typing import Any import torch import torch.nn as nn @@ -21,7 +21,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin from ...loaders.single_file_model import FromOriginalModelMixin -from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, logging from ...utils.torch_utils import maybe_allow_in_graph from ..attention import FeedForward from ..attention_processor import MochiAttention, MochiAttnProcessor2_0 @@ -104,7 +104,7 @@ def __init__( def forward( self, hidden_states: torch.Tensor, emb: torch.Tensor - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: hidden_states_dtype = hidden_states.dtype emb = self.linear(self.silu(emb)) @@ -205,8 +205,8 @@ def forward( encoder_hidden_states: torch.Tensor, temb: torch.Tensor, encoder_attention_mask: torch.Tensor, - image_rotary_emb: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + image_rotary_emb: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb) if not self.context_pre_only: @@ -268,8 +268,8 @@ def _get_positions( num_frames: int, height: int, width: int, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ) -> torch.Tensor: scale = (self.target_area / (height * width)) ** 0.5 @@ -297,9 +297,9 @@ def forward( num_frames: int, height: int, width: int, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: pos = self._get_positions(num_frames, height, width, device, dtype) rope_cos, rope_sin = self._create_rope(pos_frequencies, pos) return rope_cos, rope_sin @@ -348,7 +348,7 @@ def __init__( num_layers: int = 48, pooled_projection_dim: int = 1536, in_channels: int = 12, - out_channels: Optional[int] = None, + out_channels: int | None = None, qk_norm: str = "rms_norm", text_embed_dim: int = 4096, time_embed_dim: int = 256, @@ -404,30 +404,16 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, timestep: torch.LongTensor, encoder_attention_mask: torch.Tensor, - attention_kwargs: Optional[Dict[str, Any]] = None, + attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, ) -> torch.Tensor: - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) - batch_size, num_channels, num_frames, height, width = hidden_states.shape p = self.config.patch_size @@ -479,10 +465,6 @@ def forward( hidden_states = hidden_states.permute(0, 6, 1, 2, 4, 3, 5) output = hidden_states.reshape(batch_size, -1, num_frames, height, width) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) return Transformer2DModelOutput(sample=output) diff --git a/src/diffusers/models/transformers/transformer_omnigen.py b/src/diffusers/models/transformers/transformer_omnigen.py index 6939cac0a3a7..bd8bb107e25c 100644 --- a/src/diffusers/models/transformers/transformer_omnigen.py +++ b/src/diffusers/models/transformers/transformer_omnigen.py @@ -13,7 +13,6 @@ # limitations under the License. import math -from typing import Dict, List, Optional, Tuple, Union import torch import torch.nn as nn @@ -200,8 +199,8 @@ def __call__( attn: Attention, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: batch_size, sequence_length, _ = hidden_states.shape @@ -308,7 +307,7 @@ class OmniGenTransformer2DModel(ModelMixin, ConfigMixin): The size of the vocabulary of the embedding vocabulary. rope_base (`int`, default to `10000`): The default theta value to use when creating RoPE. - rope_scaling (`Dict`, optional): + rope_scaling (`dict`, optional): The scaling factors for the RoPE. Must contain `short_factor` and `long_factor`. pos_embed_max_size (`int`, default to `192`): The maximum size of the positional embeddings. @@ -342,7 +341,7 @@ def __init__( max_position_embeddings: int = 131072, original_max_position_embeddings: int = 4096, rope_base: int = 10000, - rope_scaling: Dict = None, + rope_scaling: dict = None, pos_embed_max_size: int = 192, time_step_dim: int = 256, flip_sin_to_cos: bool = True, @@ -387,8 +386,8 @@ def __init__( self.gradient_checkpointing = False def _get_multimodal_embeddings( - self, input_ids: torch.Tensor, input_img_latents: List[torch.Tensor], input_image_sizes: Dict - ) -> Optional[torch.Tensor]: + self, input_ids: torch.Tensor, input_img_latents: list[torch.Tensor], input_image_sizes: dict + ) -> torch.Tensor | None: if input_ids is None: return None @@ -408,14 +407,14 @@ def _get_multimodal_embeddings( def forward( self, hidden_states: torch.Tensor, - timestep: Union[int, float, torch.FloatTensor], + timestep: int | float | torch.FloatTensor, input_ids: torch.Tensor, - input_img_latents: List[torch.Tensor], - input_image_sizes: Dict[int, List[int]], + input_img_latents: list[torch.Tensor], + input_image_sizes: dict[int, list[int]], attention_mask: torch.Tensor, position_ids: torch.Tensor, return_dict: bool = True, - ) -> Union[Transformer2DModelOutput, Tuple[torch.Tensor]]: + ) -> Transformer2DModelOutput | tuple[torch.Tensor]: batch_size, num_channels, height, width = hidden_states.shape p = self.config.patch_size post_patch_height, post_patch_width = height // p, width // p diff --git a/src/diffusers/models/transformers/transformer_ovis_image.py b/src/diffusers/models/transformers/transformer_ovis_image.py index 0a09aa720b3f..8280a4871f10 100644 --- a/src/diffusers/models/transformers/transformer_ovis_image.py +++ b/src/diffusers/models/transformers/transformer_ovis_image.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -21,7 +21,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin -from ...utils import is_torch_npu_available, logging +from ...utils import logging from ...utils.torch_utils import maybe_allow_in_graph from ..attention import AttentionModuleMixin, FeedForward from ..attention_dispatch import dispatch_attention_fn @@ -78,8 +78,8 @@ def __call__( attn: "OvisImageAttention", hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections( attn, hidden_states, encoder_hidden_states @@ -145,12 +145,12 @@ def __init__( dim_head: int = 64, dropout: float = 0.0, bias: bool = False, - added_kv_proj_dim: Optional[int] = None, - added_proj_bias: Optional[bool] = True, + added_kv_proj_dim: int | None = None, + added_proj_bias: bool | None = True, out_bias: bool = True, eps: float = 1e-5, out_dim: int = None, - context_pre_only: Optional[bool] = None, + context_pre_only: bool | None = None, pre_only: bool = False, elementwise_affine: bool = True, processor=None, @@ -195,9 +195,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) @@ -238,9 +238,9 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: text_seq_len = encoder_hidden_states.shape[1] hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) @@ -301,9 +301,9 @@ def forward( hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( @@ -354,7 +354,7 @@ def forward( class OvisImagePosEmbed(nn.Module): - def __init__(self, theta: int, axes_dim: List[int]): + def __init__(self, theta: int, axes_dim: list[int]): super().__init__() self.theta = theta self.axes_dim = axes_dim @@ -413,7 +413,7 @@ class OvisImageTransformer2DModel( joint_attention_dim (`int`, defaults to `2048`): The number of dimensions to use for the joint attention (embedding/channel dimension of `encoder_hidden_states`). - axes_dims_rope (`Tuple[int]`, defaults to `(16, 56, 56)`): + axes_dims_rope (`tuple[int]`, defaults to `(16, 56, 56)`): The dimensions to use for the rotary positional embeddings. """ @@ -427,13 +427,13 @@ def __init__( self, patch_size: int = 1, in_channels: int = 64, - out_channels: Optional[int] = 64, + out_channels: int | None = 64, num_layers: int = 6, num_single_layers: int = 27, attention_head_dim: int = 128, num_attention_heads: int = 24, joint_attention_dim: int = 2048, - axes_dims_rope: Tuple[int, int, int] = (16, 56, 56), + axes_dims_rope: tuple[int, int, int] = (16, 56, 56), ): super().__init__() self.out_channels = out_channels or in_channels @@ -483,7 +483,7 @@ def forward( img_ids: torch.Tensor = None, txt_ids: torch.Tensor = None, return_dict: bool = True, - ) -> Union[torch.Tensor, Transformer2DModelOutput]: + ) -> torch.Tensor | Transformer2DModelOutput: """ The [`OvisImageTransformer2DModel`] forward method. @@ -530,11 +530,7 @@ def forward( img_ids = img_ids[0] ids = torch.cat((txt_ids, img_ids), dim=0) - if is_torch_npu_available(): - freqs_cos, freqs_sin = self.pos_embed(ids.cpu()) - image_rotary_emb = (freqs_cos.npu(), freqs_sin.npu()) - else: - image_rotary_emb = self.pos_embed(ids) + image_rotary_emb = self.pos_embed(ids) for index_block, block in enumerate(self.transformer_blocks): if torch.is_grad_enabled() and self.gradient_checkpointing: diff --git a/src/diffusers/models/transformers/transformer_prx.py b/src/diffusers/models/transformers/transformer_prx.py index a87c120fdcd7..3c8e8ae4e2c9 100644 --- a/src/diffusers/models/transformers/transformer_prx.py +++ b/src/diffusers/models/transformers/transformer_prx.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch from torch import nn @@ -96,9 +96,9 @@ def __call__( self, attn: "PRXAttention", hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: """ @@ -234,9 +234,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: return self.processor( @@ -263,10 +263,10 @@ class PRXEmbedND(nn.Module): theta (int): Scaling factor that controls the frequency spectrum of the rotary embeddings. axes_dim (list[int]): - List of embedding dimensions for each axis (each must be even). + list of embedding dimensions for each axis (each must be even). """ - def __init__(self, dim: int, theta: int, axes_dim: List[int]): + def __init__(self, dim: int, theta: int, axes_dim: list[int]): super().__init__() self.dim = dim self.theta = theta @@ -346,7 +346,7 @@ def __init__(self, dim: int): def forward( self, vec: torch.Tensor - ) -> Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + ) -> tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(6, dim=-1) return tuple(out[:3]), tuple(out[3:]) @@ -389,7 +389,7 @@ def __init__( hidden_size: int, num_heads: int, mlp_ratio: float = 4.0, - qk_scale: Optional[float] = None, + qk_scale: float | None = None, ): super().__init__() @@ -430,8 +430,8 @@ def forward( encoder_hidden_states: torch.Tensor, temb: torch.Tensor, image_rotary_emb: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - **kwargs: Dict[str, Any], + attention_mask: torch.Tensor | None = None, + **kwargs: dict[str, Any], ) -> torch.Tensor: r""" Runs modulation-gated cross-attention and MLP, with residual connections. @@ -607,7 +607,7 @@ class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin): depth (`int`, *optional*, defaults to 16): Number of transformer blocks. axes_dim (`list[int]`, *optional*): - List of dimensions for each positional embedding axis. Defaults to `[32, 32]`. + list of dimensions for each positional embedding axis. Defaults to `[32, 32]`. theta (`int`, *optional*, defaults to 10000): Frequency scaling factor for rotary embeddings. time_factor (`float`, *optional*, defaults to 1000.0): @@ -728,10 +728,10 @@ def forward( hidden_states: torch.Tensor, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - attention_kwargs: Optional[Dict[str, Any]] = None, + attention_mask: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - ) -> Union[Tuple[torch.Tensor, ...], Transformer2DModelOutput]: + ) -> tuple[torch.Tensor, ...] | Transformer2DModelOutput: r""" Forward pass of the PRXTransformer2DModel. diff --git a/src/diffusers/models/transformers/transformer_qwenimage.py b/src/diffusers/models/transformers/transformer_qwenimage.py index 1229bab169b2..a54cb3b8e092 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage.py +++ b/src/diffusers/models/transformers/transformer_qwenimage.py @@ -15,7 +15,7 @@ import functools import math from math import prod -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import numpy as np import torch @@ -24,7 +24,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, deprecate, logging from ...utils.torch_utils import maybe_allow_in_graph from .._modeling_parallel import ContextParallelInput, ContextParallelOutput from ..attention import AttentionMixin, FeedForward @@ -96,10 +96,10 @@ def get_timestep_embedding( def apply_rotary_emb_qwen( x: torch.Tensor, - freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]], + freqs_cis: torch.Tensor | tuple[torch.Tensor], use_real: bool = True, use_real_unbind_dim: int = -1, -) -> Tuple[torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor]: """ Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are @@ -109,10 +109,10 @@ def apply_rotary_emb_qwen( Args: x (`torch.Tensor`): Query or key tensor to apply rotary embeddings. [B, S, H, D] xk (torch.Tensor): Key tensor to apply - freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],) + freqs_cis (`tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],) Returns: - Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings. + tuple[torch.Tensor, torch.Tensor]: tuple of modified query tensor and key tensor with rotary embeddings. """ if use_real: cos, sin = freqs_cis # [S, D] @@ -142,6 +142,36 @@ def apply_rotary_emb_qwen( return x_out.type_as(x) +def compute_text_seq_len_from_mask( + encoder_hidden_states: torch.Tensor, encoder_hidden_states_mask: torch.Tensor | None +) -> tuple[int, torch.Tensor | None, torch.Tensor | None]: + """ + Compute text sequence length without assuming contiguous masks. Returns length for RoPE and a normalized bool mask. + """ + batch_size, text_seq_len = encoder_hidden_states.shape[:2] + if encoder_hidden_states_mask is None: + return text_seq_len, None, None + + if encoder_hidden_states_mask.shape[:2] != (batch_size, text_seq_len): + raise ValueError( + f"`encoder_hidden_states_mask` shape {encoder_hidden_states_mask.shape} must match " + f"(batch_size, text_seq_len)=({batch_size}, {text_seq_len})." + ) + + if encoder_hidden_states_mask.dtype != torch.bool: + encoder_hidden_states_mask = encoder_hidden_states_mask.to(torch.bool) + + position_ids = torch.arange(text_seq_len, device=encoder_hidden_states.device, dtype=torch.long) + active_positions = torch.where(encoder_hidden_states_mask, position_ids, position_ids.new_zeros(())) + has_active = encoder_hidden_states_mask.any(dim=1) + per_sample_len = torch.where( + has_active, + active_positions.max(dim=1).values + 1, + torch.as_tensor(text_seq_len, device=encoder_hidden_states.device), + ) + return text_seq_len, per_sample_len, encoder_hidden_states_mask + + class QwenTimestepProjEmbeddings(nn.Module): def __init__(self, embedding_dim, use_additional_t_cond=False): super().__init__() @@ -168,7 +198,7 @@ def forward(self, timestep, hidden_states, addition_t_cond=None): class QwenEmbedRope(nn.Module): - def __init__(self, theta: int, axes_dim: List[int], scale_rope=False): + def __init__(self, theta: int, axes_dim: list[int], scale_rope=False): super().__init__() self.theta = theta self.axes_dim = axes_dim @@ -206,22 +236,51 @@ def rope_params(self, index, dim, theta=10000): def forward( self, - video_fhw: Union[Tuple[int, int, int], List[Tuple[int, int, int]]], - txt_seq_lens: List[int], - device: torch.device, - ) -> Tuple[torch.Tensor, torch.Tensor]: + video_fhw: tuple[int, int, int, list[tuple[int, int, int]]], + txt_seq_lens: list[int] | None = None, + device: torch.device = None, + max_txt_seq_len: int | torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: """ Args: - video_fhw (`Tuple[int, int, int]` or `List[Tuple[int, int, int]]`): + video_fhw (`tuple[int, int, int]` or `list[tuple[int, int, int]]`): A list of 3 integers [frame, height, width] representing the shape of the video. - txt_seq_lens (`List[int]`): - A list of integers of length batch_size representing the length of each text prompt. - device: (`torch.device`): + txt_seq_lens (`list[int]`, *optional*, **Deprecated**): + Deprecated parameter. Use `max_txt_seq_len` instead. If provided, the maximum value will be used. + device: (`torch.device`, *optional*): The device on which to perform the RoPE computation. + max_txt_seq_len (`int` or `torch.Tensor`, *optional*): + The maximum text sequence length for RoPE computation. This should match the encoder hidden states + sequence length. Can be either an int or a scalar tensor (for torch.compile compatibility). """ - if self.pos_freqs.device != device: - self.pos_freqs = self.pos_freqs.to(device) - self.neg_freqs = self.neg_freqs.to(device) + # Handle deprecated txt_seq_lens parameter + if txt_seq_lens is not None: + deprecate( + "txt_seq_lens", + "0.39.0", + "Passing `txt_seq_lens` is deprecated and will be removed in version 0.39.0. " + "Please use `max_txt_seq_len` instead. " + "The new parameter accepts a single int or tensor value representing the maximum text sequence length.", + standard_warn=False, + ) + if max_txt_seq_len is None: + # Use max of txt_seq_lens for backward compatibility + max_txt_seq_len = max(txt_seq_lens) if isinstance(txt_seq_lens, list) else txt_seq_lens + + if max_txt_seq_len is None: + raise ValueError("Either `max_txt_seq_len` or `txt_seq_lens` (deprecated) must be provided.") + + # Validate batch inference with variable-sized images + if isinstance(video_fhw, list) and len(video_fhw) > 1: + # Check if all instances have the same size + first_fhw = video_fhw[0] + if not all(fhw == first_fhw for fhw in video_fhw): + logger.warning( + "Batch inference with variable-sized images is not currently supported in QwenEmbedRope. " + "All images in the batch should have the same dimensions (frame, height, width). " + f"Detected sizes: {video_fhw}. Using the first image's dimensions {first_fhw} " + "for RoPE computation, which may lead to incorrect results for other images in the batch." + ) if isinstance(video_fhw, list): video_fhw = video_fhw[0] @@ -233,8 +292,7 @@ def forward( for idx, fhw in enumerate(video_fhw): frame, height, width = fhw # RoPE frequencies are cached via a lru_cache decorator on _compute_video_freqs - video_freq = self._compute_video_freqs(frame, height, width, idx) - video_freq = video_freq.to(device) + video_freq = self._compute_video_freqs(frame, height, width, idx, device) vid_freqs.append(video_freq) if self.scale_rope: @@ -242,17 +300,23 @@ def forward( else: max_vid_index = max(height, width, max_vid_index) - max_len = max(txt_seq_lens) - txt_freqs = self.pos_freqs[max_vid_index : max_vid_index + max_len, ...] + max_txt_seq_len_int = int(max_txt_seq_len) + # Create device-specific copy for text freqs without modifying self.pos_freqs + txt_freqs = self.pos_freqs.to(device)[max_vid_index : max_vid_index + max_txt_seq_len_int, ...] vid_freqs = torch.cat(vid_freqs, dim=0) return vid_freqs, txt_freqs @functools.lru_cache(maxsize=128) - def _compute_video_freqs(self, frame: int, height: int, width: int, idx: int = 0) -> torch.Tensor: + def _compute_video_freqs( + self, frame: int, height: int, width: int, idx: int = 0, device: torch.device = None + ) -> torch.Tensor: seq_lens = frame * height * width - freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1) - freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1) + pos_freqs = self.pos_freqs.to(device) if device is not None else self.pos_freqs + neg_freqs = self.neg_freqs.to(device) if device is not None else self.neg_freqs + + freqs_pos = pos_freqs.split([x // 2 for x in self.axes_dim], dim=1) + freqs_neg = neg_freqs.split([x // 2 for x in self.axes_dim], dim=1) freqs_frame = freqs_pos[0][idx : idx + frame].view(frame, 1, 1, -1).expand(frame, height, width, -1) if self.scale_rope: @@ -269,7 +333,7 @@ def _compute_video_freqs(self, frame: int, height: int, width: int, idx: int = 0 class QwenEmbedLayer3DRope(nn.Module): - def __init__(self, theta: int, axes_dim: List[int], scale_rope=False): + def __init__(self, theta: int, axes_dim: list[int], scale_rope=False): super().__init__() self.theta = theta self.axes_dim = axes_dim @@ -304,14 +368,35 @@ def rope_params(self, index, dim, theta=10000): freqs = torch.polar(torch.ones_like(freqs), freqs) return freqs - def forward(self, video_fhw, txt_seq_lens, device): + def forward( + self, + video_fhw: tuple[int, int, int, list[tuple[int, int, int]]], + max_txt_seq_len: int | torch.Tensor, + device: torch.device = None, + ) -> tuple[torch.Tensor, torch.Tensor]: """ - Args: video_fhw: [frame, height, width] a list of 3 integers representing the shape of the video Args: - txt_length: [bs] a list of 1 integers representing the length of the text + Args: + video_fhw (`tuple[int, int, int]` or `list[tuple[int, int, int]]`): + A list of 3 integers [frame, height, width] representing the shape of the video, or a list of layer + structures. + max_txt_seq_len (`int` or `torch.Tensor`): + The maximum text sequence length for RoPE computation. This should match the encoder hidden states + sequence length. Can be either an int or a scalar tensor (for torch.compile compatibility). + device: (`torch.device`, *optional*): + The device on which to perform the RoPE computation. """ - if self.pos_freqs.device != device: - self.pos_freqs = self.pos_freqs.to(device) - self.neg_freqs = self.neg_freqs.to(device) + # Validate batch inference with variable-sized images + # In Layer3DRope, the outer list represents batch, inner list/tuple represents layers + if isinstance(video_fhw, list) and len(video_fhw) > 1: + # Check if this is batch inference (list of layer lists/tuples) + first_entry = video_fhw[0] + if not all(entry == first_entry for entry in video_fhw): + logger.warning( + "Batch inference with variable-sized images is not currently supported in QwenEmbedLayer3DRope. " + "All images in the batch should have the same layer structure. " + f"Detected sizes: {video_fhw}. Using the first image's layer structure {first_entry} " + "for RoPE computation, which may lead to incorrect results for other images in the batch." + ) if isinstance(video_fhw, list): video_fhw = video_fhw[0] @@ -324,11 +409,10 @@ def forward(self, video_fhw, txt_seq_lens, device): for idx, fhw in enumerate(video_fhw): frame, height, width = fhw if idx != layer_num: - video_freq = self._compute_video_freqs(frame, height, width, idx) + video_freq = self._compute_video_freqs(frame, height, width, idx, device) else: ### For the condition image, we set the layer index to -1 - video_freq = self._compute_condition_freqs(frame, height, width) - video_freq = video_freq.to(device) + video_freq = self._compute_condition_freqs(frame, height, width, device) vid_freqs.append(video_freq) if self.scale_rope: @@ -337,17 +421,21 @@ def forward(self, video_fhw, txt_seq_lens, device): max_vid_index = max(height, width, max_vid_index) max_vid_index = max(max_vid_index, layer_num) - max_len = max(txt_seq_lens) - txt_freqs = self.pos_freqs[max_vid_index : max_vid_index + max_len, ...] + max_txt_seq_len_int = int(max_txt_seq_len) + # Create device-specific copy for text freqs without modifying self.pos_freqs + txt_freqs = self.pos_freqs.to(device)[max_vid_index : max_vid_index + max_txt_seq_len_int, ...] vid_freqs = torch.cat(vid_freqs, dim=0) return vid_freqs, txt_freqs @functools.lru_cache(maxsize=None) - def _compute_video_freqs(self, frame, height, width, idx=0): + def _compute_video_freqs(self, frame, height, width, idx=0, device: torch.device = None): seq_lens = frame * height * width - freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1) - freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1) + pos_freqs = self.pos_freqs.to(device) if device is not None else self.pos_freqs + neg_freqs = self.neg_freqs.to(device) if device is not None else self.neg_freqs + + freqs_pos = pos_freqs.split([x // 2 for x in self.axes_dim], dim=1) + freqs_neg = neg_freqs.split([x // 2 for x in self.axes_dim], dim=1) freqs_frame = freqs_pos[0][idx : idx + frame].view(frame, 1, 1, -1).expand(frame, height, width, -1) if self.scale_rope: @@ -363,10 +451,13 @@ def _compute_video_freqs(self, frame, height, width, idx=0): return freqs.clone().contiguous() @functools.lru_cache(maxsize=None) - def _compute_condition_freqs(self, frame, height, width): + def _compute_condition_freqs(self, frame, height, width, device: torch.device = None): seq_lens = frame * height * width - freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1) - freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1) + pos_freqs = self.pos_freqs.to(device) if device is not None else self.pos_freqs + neg_freqs = self.neg_freqs.to(device) if device is not None else self.neg_freqs + + freqs_pos = pos_freqs.split([x // 2 for x in self.axes_dim], dim=1) + freqs_neg = neg_freqs.split([x // 2 for x in self.axes_dim], dim=1) freqs_frame = freqs_neg[0][-1:].view(frame, 1, 1, -1).expand(frame, height, width, -1) if self.scale_rope: @@ -403,8 +494,8 @@ def __call__( hidden_states: torch.FloatTensor, # Image stream encoder_hidden_states: torch.FloatTensor = None, # Text stream encoder_hidden_states_mask: torch.FloatTensor = None, - attention_mask: Optional[torch.FloatTensor] = None, - image_rotary_emb: Optional[torch.Tensor] = None, + attention_mask: torch.FloatTensor | None = None, + image_rotary_emb: torch.Tensor | None = None, ) -> torch.FloatTensor: if encoder_hidden_states is None: raise ValueError("QwenDoubleStreamAttnProcessor2_0 requires encoder_hidden_states (text stream)") @@ -454,7 +545,6 @@ def __call__( joint_key = torch.cat([txt_key, img_key], dim=1) joint_value = torch.cat([txt_value, img_value], dim=1) - # Compute joint attention joint_hidden_states = dispatch_attention_fn( joint_query, joint_key, @@ -475,11 +565,11 @@ def __call__( img_attn_output = joint_hidden_states[:, seq_txt:, :] # Image part # Apply output projections - img_attn_output = attn.to_out[0](img_attn_output) + img_attn_output = attn.to_out[0](img_attn_output.contiguous()) if len(attn.to_out) > 1: img_attn_output = attn.to_out[1](img_attn_output) # dropout - txt_attn_output = attn.to_add_out(txt_attn_output) + txt_attn_output = attn.to_add_out(txt_attn_output.contiguous()) return img_attn_output, txt_attn_output @@ -577,10 +667,10 @@ def forward( encoder_hidden_states: torch.Tensor, encoder_hidden_states_mask: torch.Tensor, temb: torch.Tensor, - image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - modulate_index: Optional[List[int]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, + modulate_index: list[int] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: # Get modulation parameters for both streams img_mod_params = self.img_mod(temb) # [B, 6*dim] @@ -667,7 +757,7 @@ class QwenImageTransformer2DModel( `encoder_hidden_states`). guidance_embeds (`bool`, defaults to `False`): Whether to use guidance embeddings for guidance-distilled variant of the model. - axes_dims_rope (`Tuple[int]`, defaults to `(16, 56, 56)`): + axes_dims_rope (`tuple[int]`, defaults to `(16, 56, 56)`): The dimensions to use for the rotary positional embeddings. """ @@ -675,11 +765,14 @@ class QwenImageTransformer2DModel( _no_split_modules = ["QwenImageTransformerBlock"] _skip_layerwise_casting_patterns = ["pos_embed", "norm"] _repeated_blocks = ["QwenImageTransformerBlock"] + # Make CP plan compatible with https://github.com/huggingface/diffusers/pull/12702 _cp_plan = { - "": { + "transformer_blocks.0": { "hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), "encoder_hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), - "encoder_hidden_states_mask": ContextParallelInput(split_dim=1, expected_dims=2, split_output=False), + }, + "transformer_blocks.*": { + "modulate_index": ContextParallelInput(split_dim=1, expected_dims=2, split_output=False), }, "pos_embed": { 0: ContextParallelInput(split_dim=0, expected_dims=2, split_output=True), @@ -693,13 +786,13 @@ def __init__( self, patch_size: int = 2, in_channels: int = 64, - out_channels: Optional[int] = 16, + out_channels: int | None = 16, num_layers: int = 60, attention_head_dim: int = 128, num_attention_heads: int = 24, joint_attention_dim: int = 3584, guidance_embeds: bool = False, # TODO: this should probably be removed - axes_dims_rope: Tuple[int, int, int] = (16, 56, 56), + axes_dims_rope: tuple[int, int, int] = (16, 56, 56), zero_cond_t: bool = False, use_additional_t_cond: bool = False, use_layer3d_rope: bool = False, @@ -740,20 +833,21 @@ def __init__( self.gradient_checkpointing = False self.zero_cond_t = zero_cond_t + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor = None, encoder_hidden_states_mask: torch.Tensor = None, timestep: torch.LongTensor = None, - img_shapes: Optional[List[Tuple[int, int, int]]] = None, - txt_seq_lens: Optional[List[int]] = None, + img_shapes: list[tuple[int, int, int]] | None = None, + txt_seq_lens: list[int] | None = None, guidance: torch.Tensor = None, # TODO: this should probably be removed - attention_kwargs: Optional[Dict[str, Any]] = None, + attention_kwargs: dict[str, Any] | None = None, controlnet_block_samples=None, additional_t_cond=None, return_dict: bool = True, - ) -> Union[torch.Tensor, Transformer2DModelOutput]: + ) -> torch.Tensor | Transformer2DModelOutput: """ The [`QwenTransformer2DModel`] forward method. @@ -762,14 +856,25 @@ def forward( Input `hidden_states`. encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`): Conditional embeddings (embeddings computed from the input conditions such as prompts) to use. - encoder_hidden_states_mask (`torch.Tensor` of shape `(batch_size, text_sequence_length)`): - Mask of the input conditions. + encoder_hidden_states_mask (`torch.Tensor` of shape `(batch_size, text_sequence_length)`, *optional*): + Mask for the encoder hidden states. Expected to have 1.0 for valid tokens and 0.0 for padding tokens. + Used in the attention processor to prevent attending to padding tokens. The mask can have any pattern + (not just contiguous valid tokens followed by padding) since it's applied element-wise in attention. timestep ( `torch.LongTensor`): Used to indicate denoising step. + img_shapes (`list[tuple[int, int, int]]`, *optional*): + Image shapes for RoPE computation. + txt_seq_lens (`list[int]`, *optional*, **Deprecated**): + Deprecated parameter. Use `encoder_hidden_states_mask` instead. If provided, the maximum value will be + used to compute RoPE sequence length. + guidance (`torch.Tensor`, *optional*): + Guidance tensor for conditional generation. attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + controlnet_block_samples (*optional*): + ControlNet block samples to add to the transformer blocks. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain tuple. @@ -778,20 +883,15 @@ def forward( If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a `tuple` where the first element is the sample tensor. """ - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective." - ) + if txt_seq_lens is not None: + deprecate( + "txt_seq_lens", + "0.39.0", + "Passing `txt_seq_lens` is deprecated and will be removed in version 0.39.0. " + "Please use `encoder_hidden_states_mask` instead. " + "The mask-based approach is more flexible and supports variable-length sequences.", + standard_warn=False, + ) hidden_states = self.img_in(hidden_states) @@ -810,6 +910,11 @@ def forward( encoder_hidden_states = self.txt_norm(encoder_hidden_states) encoder_hidden_states = self.txt_in(encoder_hidden_states) + # Use the encoder_hidden_states sequence length for RoPE computation and normalize mask + text_seq_len, _, encoder_hidden_states_mask = compute_text_seq_len_from_mask( + encoder_hidden_states, encoder_hidden_states_mask + ) + if guidance is not None: guidance = guidance.to(hidden_states.dtype) * 1000 @@ -819,7 +924,17 @@ def forward( else self.time_text_embed(timestep, guidance, hidden_states, additional_t_cond) ) - image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device) + image_rotary_emb = self.pos_embed(img_shapes, max_txt_seq_len=text_seq_len, device=hidden_states.device) + + # Construct joint attention mask once to avoid reconstructing in every block + # This eliminates 60 GPU syncs during training while maintaining torch.compile compatibility + block_attention_kwargs = attention_kwargs.copy() if attention_kwargs is not None else {} + if encoder_hidden_states_mask is not None: + # Build joint mask: [text_mask, all_ones_for_image] + batch_size, image_seq_len = hidden_states.shape[:2] + image_mask = torch.ones((batch_size, image_seq_len), dtype=torch.bool, device=hidden_states.device) + joint_attention_mask = torch.cat([encoder_hidden_states_mask, image_mask], dim=1) + block_attention_kwargs["attention_mask"] = joint_attention_mask for index_block, block in enumerate(self.transformer_blocks): if torch.is_grad_enabled() and self.gradient_checkpointing: @@ -827,10 +942,10 @@ def forward( block, hidden_states, encoder_hidden_states, - encoder_hidden_states_mask, + None, # Don't pass encoder_hidden_states_mask (using attention_mask instead) temb, image_rotary_emb, - attention_kwargs, + block_attention_kwargs, modulate_index, ) @@ -838,10 +953,10 @@ def forward( encoder_hidden_states, hidden_states = block( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, - encoder_hidden_states_mask=encoder_hidden_states_mask, + encoder_hidden_states_mask=None, # Don't pass (using attention_mask instead) temb=temb, image_rotary_emb=image_rotary_emb, - joint_attention_kwargs=attention_kwargs, + joint_attention_kwargs=block_attention_kwargs, modulate_index=modulate_index, ) @@ -857,10 +972,6 @@ def forward( hidden_states = self.norm_out(hidden_states, temb) output = self.proj_out(hidden_states) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) diff --git a/src/diffusers/models/transformers/transformer_sana_video.py b/src/diffusers/models/transformers/transformer_sana_video.py index a4f90342631a..f833c0e842c3 100644 --- a/src/diffusers/models/transformers/transformer_sana_video.py +++ b/src/diffusers/models/transformers/transformer_sana_video.py @@ -13,7 +13,7 @@ # limitations under the License. import math -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any import torch import torch.nn.functional as F @@ -21,7 +21,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, logging from ..attention import AttentionMixin from ..attention_dispatch import dispatch_attention_fn from ..attention_processor import Attention @@ -40,7 +40,7 @@ def __init__( in_channels: int, out_channels: int, expand_ratio: float = 4, - norm_type: Optional[str] = None, + norm_type: str | None = None, residual_connection: bool = True, ) -> None: super().__init__() @@ -103,9 +103,9 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - rotary_emb: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: original_dtype = hidden_states.dtype @@ -176,7 +176,7 @@ class WanRotaryPosEmbed(nn.Module): def __init__( self, attention_head_dim: int, - patch_size: Tuple[int, int, int], + patch_size: tuple[int, int, int], max_seq_len: int, theta: float = 10000.0, ): @@ -290,8 +290,8 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: batch_size, sequence_length, _ = ( hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape @@ -358,15 +358,15 @@ def __init__( num_attention_heads: int = 20, attention_head_dim: int = 112, dropout: float = 0.0, - num_cross_attention_heads: Optional[int] = 20, - cross_attention_head_dim: Optional[int] = 112, - cross_attention_dim: Optional[int] = 2240, + num_cross_attention_heads: int | None = 20, + cross_attention_head_dim: int | None = 112, + cross_attention_dim: int | None = 2240, attention_bias: bool = True, norm_elementwise_affine: bool = False, norm_eps: float = 1e-6, attention_out_bias: bool = True, mlp_ratio: float = 3.0, - qk_norm: Optional[str] = "rms_norm_across_heads", + qk_norm: str | None = "rms_norm_across_heads", rope_max_seq_len: int = 1024, ) -> None: super().__init__() @@ -409,14 +409,14 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - timestep: Optional[torch.LongTensor] = None, + attention_mask: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + timestep: torch.LongTensor | None = None, frames: int = None, height: int = None, width: int = None, - rotary_emb: Optional[torch.Tensor] = None, + rotary_emb: torch.Tensor | None = None, ) -> torch.Tensor: batch_size = hidden_states.shape[0] @@ -503,25 +503,25 @@ class SanaVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Fro def __init__( self, in_channels: int = 16, - out_channels: Optional[int] = 16, + out_channels: int | None = 16, num_attention_heads: int = 20, attention_head_dim: int = 112, num_layers: int = 20, - num_cross_attention_heads: Optional[int] = 20, - cross_attention_head_dim: Optional[int] = 112, - cross_attention_dim: Optional[int] = 2240, + num_cross_attention_heads: int | None = 20, + cross_attention_head_dim: int | None = 112, + cross_attention_dim: int | None = 2240, caption_channels: int = 2304, mlp_ratio: float = 2.5, dropout: float = 0.0, attention_bias: bool = False, sample_size: int = 30, - patch_size: Tuple[int, int, int] = (1, 2, 2), + patch_size: tuple[int, int, int] = (1, 2, 2), norm_elementwise_affine: bool = False, norm_eps: float = 1e-6, - interpolation_scale: Optional[int] = None, + interpolation_scale: int | None = None, guidance_embeds: bool = False, guidance_embeds_scale: float = 0.1, - qk_norm: Optional[str] = "rms_norm_across_heads", + qk_norm: str | None = "rms_norm_across_heads", rope_max_seq_len: int = 1024, ) -> None: super().__init__() @@ -570,33 +570,19 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, timestep: torch.Tensor, - guidance: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_block_samples: Optional[Tuple[torch.Tensor]] = None, + guidance: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, + controlnet_block_samples: tuple[torch.Tensor] | None = None, return_dict: bool = True, - ) -> Union[Tuple[torch.Tensor, ...], Transformer2DModelOutput]: - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) - + ) -> tuple[torch.Tensor, ...] | Transformer2DModelOutput: # ensure attention_mask is a bias, and give it a singleton query_tokens dimension. # we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward. # we can tell by counting dims; if ndim == 2: it's a mask rather than a bias. @@ -695,10 +681,6 @@ def forward( hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) diff --git a/src/diffusers/models/transformers/transformer_sd3.py b/src/diffusers/models/transformers/transformer_sd3.py index 05391e047b7a..ead657d0cfd2 100644 --- a/src/diffusers/models/transformers/transformer_sd3.py +++ b/src/diffusers/models/transformers/transformer_sd3.py @@ -11,14 +11,14 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin, SD3Transformer2DLoadersMixin -from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, logging from ...utils.torch_utils import maybe_allow_in_graph from ..attention import AttentionMixin, FeedForward, JointTransformerBlock from ..attention_processor import ( @@ -106,7 +106,7 @@ class SD3Transformer2DModel( The number of latent channels in the output. pos_embed_max_size (`int`, defaults to `96`): The maximum latent height/width of positional embeddings. - dual_attention_layers (`Tuple[int, ...]`, defaults to `()`): + dual_attention_layers (`tuple[int, ...]`, defaults to `()`): The number of dual-stream transformer blocks to use. qk_norm (`str`, *optional*, defaults to `None`): The normalization to use for query and key in the attention layer. If `None`, no normalization is used. @@ -130,10 +130,10 @@ def __init__( pooled_projection_dim: int = 2048, out_channels: int = 16, pos_embed_max_size: int = 96, - dual_attention_layers: Tuple[ + dual_attention_layers: tuple[ int, ... ] = (), # () for sd3.0; (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12) for sd3.5 - qk_norm: Optional[str] = None, + qk_norm: str | None = None, ): super().__init__() self.out_channels = out_channels if out_channels is not None else in_channels @@ -172,7 +172,7 @@ def __init__( self.gradient_checkpointing = False # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking - def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None: + def enable_forward_chunking(self, chunk_size: int | None = None, dim: int = 0) -> None: """ Sets the attention processor to use [feed forward chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers). @@ -245,17 +245,18 @@ def unfuse_qkv_projections(self): if self.original_attn_processors is not None: self.set_attn_processor(self.original_attn_processors) + @apply_lora_scale("joint_attention_kwargs") def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor = None, pooled_projections: torch.Tensor = None, timestep: torch.LongTensor = None, - block_controlnet_hidden_states: List = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, + block_controlnet_hidden_states: list = None, + joint_attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - skip_layers: Optional[List[int]] = None, - ) -> Union[torch.Tensor, Transformer2DModelOutput]: + skip_layers: list[int] | None = None, + ) -> torch.Tensor | Transformer2DModelOutput: """ The [`SD3Transformer2DModel`] forward method. @@ -284,20 +285,6 @@ def forward( If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a `tuple` where the first element is the sample tensor. """ - if joint_attention_kwargs is not None: - joint_attention_kwargs = joint_attention_kwargs.copy() - lora_scale = joint_attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective." - ) height, width = hidden_states.shape[-2:] @@ -352,10 +339,6 @@ def forward( shape=(hidden_states.shape[0], self.out_channels, height * patch_size, width * patch_size) ) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) diff --git a/src/diffusers/models/transformers/transformer_skyreels_v2.py b/src/diffusers/models/transformers/transformer_skyreels_v2.py index 2b9fc5b8d9fb..9067e32ea5c3 100644 --- a/src/diffusers/models/transformers/transformer_skyreels_v2.py +++ b/src/diffusers/models/transformers/transformer_skyreels_v2.py @@ -13,7 +13,7 @@ # limitations under the License. import math -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -21,7 +21,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, deprecate, logging from ...utils.torch_utils import maybe_allow_in_graph from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward from ..attention_dispatch import dispatch_attention_fn @@ -85,9 +85,9 @@ def __call__( self, attn: "SkyReelsV2Attention", hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: encoder_hidden_states_img = None if attn.add_k_proj is not None: @@ -188,8 +188,8 @@ def __init__( dim_head: int = 64, eps: float = 1e-5, dropout: float = 0.0, - added_kv_proj_dim: Optional[int] = None, - cross_attention_dim_head: Optional[int] = None, + added_kv_proj_dim: int | None = None, + cross_attention_dim_head: int | None = None, processor=None, is_cross_attention=None, ): @@ -275,9 +275,9 @@ def unfuse_projections(self): def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, **kwargs, ) -> torch.Tensor: return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, rotary_emb, **kwargs) @@ -335,8 +335,8 @@ def __init__( time_freq_dim: int, time_proj_dim: int, text_embed_dim: int, - image_embed_dim: Optional[int] = None, - pos_embed_seq_len: Optional[int] = None, + image_embed_dim: int | None = None, + pos_embed_seq_len: int | None = None, ): super().__init__() @@ -354,7 +354,7 @@ def forward( self, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, - encoder_hidden_states_image: Optional[torch.Tensor] = None, + encoder_hidden_states_image: torch.Tensor | None = None, ): timestep = self.timesteps_proj(timestep) @@ -375,7 +375,7 @@ class SkyReelsV2RotaryPosEmbed(nn.Module): def __init__( self, attention_head_dim: int, - patch_size: Tuple[int, int, int], + patch_size: tuple[int, int, int], max_seq_len: int, theta: float = 10000.0, ): @@ -445,7 +445,7 @@ def __init__( qk_norm: str = "rms_norm_across_heads", cross_attn_norm: bool = False, eps: float = 1e-6, - added_kv_proj_dim: Optional[int] = None, + added_kv_proj_dim: int | None = None, ): super().__init__() @@ -522,7 +522,7 @@ class SkyReelsV2Transformer3DModel( A Transformer model for video-like data used in the Wan-based SkyReels-V2 model. Args: - patch_size (`Tuple[int]`, defaults to `(1, 2, 2)`): + patch_size (`tuple[int]`, defaults to `(1, 2, 2)`): 3D patch dimensions for video embedding (t_patch, h_patch, w_patch). num_attention_heads (`int`, defaults to `16`): Fixed length for text embeddings. @@ -540,7 +540,7 @@ class SkyReelsV2Transformer3DModel( Intermediate dimension in feed-forward network. num_layers (`int`, defaults to `32`): The number of layers of transformer blocks to use. - window_size (`Tuple[int]`, defaults to `(-1, -1)`): + window_size (`tuple[int]`, defaults to `(-1, -1)`): Window size for local attention (-1 indicates global attention). cross_attn_norm (`bool`, defaults to `True`): Enable cross-attention normalization. @@ -570,7 +570,7 @@ class SkyReelsV2Transformer3DModel( @register_to_config def __init__( self, - patch_size: Tuple[int, ...] = (1, 2, 2), + patch_size: tuple[int] = (1, 2, 2), num_attention_heads: int = 16, attention_head_dim: int = 128, in_channels: int = 16, @@ -580,12 +580,12 @@ def __init__( ffn_dim: int = 8192, num_layers: int = 32, cross_attn_norm: bool = True, - qk_norm: Optional[str] = "rms_norm_across_heads", + qk_norm: str | None = "rms_norm_across_heads", eps: float = 1e-6, - image_dim: Optional[int] = None, - added_kv_proj_dim: Optional[int] = None, + image_dim: int | None = None, + added_kv_proj_dim: int | None = None, rope_max_seq_len: int = 1024, - pos_embed_seq_len: Optional[int] = None, + pos_embed_seq_len: int | None = None, inject_sample_info: bool = False, num_frame_per_block: int = 1, ) -> None: @@ -630,32 +630,18 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, timestep: torch.LongTensor, encoder_hidden_states: torch.Tensor, - encoder_hidden_states_image: Optional[torch.Tensor] = None, + encoder_hidden_states_image: torch.Tensor | None = None, enable_diffusion_forcing: bool = False, - fps: Optional[torch.Tensor] = None, + fps: torch.Tensor | None = None, return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) - + attention_kwargs: dict[str, Any] | None = None, + ) -> torch.Tensor | dict[str, torch.Tensor]: batch_size, num_channels, num_frames, height, width = hidden_states.shape p_t, p_h, p_w = self.config.patch_size post_patch_num_frames = num_frames // p_t @@ -771,10 +757,6 @@ def forward( hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) diff --git a/src/diffusers/models/transformers/transformer_temporal.py b/src/diffusers/models/transformers/transformer_temporal.py index ffaf31d04570..b6fedcb26cc8 100644 --- a/src/diffusers/models/transformers/transformer_temporal.py +++ b/src/diffusers/models/transformers/transformer_temporal.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from dataclasses import dataclass -from typing import Any, Dict, Optional +from typing import Any import torch from torch import nn @@ -74,19 +74,19 @@ def __init__( self, num_attention_heads: int = 16, attention_head_dim: int = 88, - in_channels: Optional[int] = None, - out_channels: Optional[int] = None, + in_channels: int | None = None, + out_channels: int | None = None, num_layers: int = 1, dropout: float = 0.0, norm_num_groups: int = 32, - cross_attention_dim: Optional[int] = None, + cross_attention_dim: int | None = None, attention_bias: bool = False, - sample_size: Optional[int] = None, + sample_size: int | None = None, activation_fn: str = "geglu", norm_elementwise_affine: bool = True, double_self_attention: bool = True, - positional_embeddings: Optional[str] = None, - num_positional_embeddings: Optional[int] = None, + positional_embeddings: str | None = None, + num_positional_embeddings: int | None = None, ): super().__init__() self.num_attention_heads = num_attention_heads @@ -123,11 +123,11 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.LongTensor] = None, - timestep: Optional[torch.LongTensor] = None, + encoder_hidden_states: torch.LongTensor | None = None, + timestep: torch.LongTensor | None = None, class_labels: torch.LongTensor = None, num_frames: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, ) -> TransformerTemporalModelOutput: """ @@ -222,9 +222,9 @@ def __init__( num_attention_heads: int = 16, attention_head_dim: int = 88, in_channels: int = 320, - out_channels: Optional[int] = None, + out_channels: int | None = None, num_layers: int = 1, - cross_attention_dim: Optional[int] = None, + cross_attention_dim: int | None = None, ): super().__init__() self.num_attention_heads = num_attention_heads @@ -280,8 +280,8 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - image_only_indicator: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + image_only_indicator: torch.Tensor | None = None, return_dict: bool = True, ): """ diff --git a/src/diffusers/models/transformers/transformer_wan.py b/src/diffusers/models/transformers/transformer_wan.py index f7693ec5d3ac..5926bbb8e713 100644 --- a/src/diffusers/models/transformers/transformer_wan.py +++ b/src/diffusers/models/transformers/transformer_wan.py @@ -13,7 +13,7 @@ # limitations under the License. import math -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -21,7 +21,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, deprecate, logging from ...utils.torch_utils import maybe_allow_in_graph from .._modeling_parallel import ContextParallelInput, ContextParallelOutput from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward @@ -42,7 +42,7 @@ def _get_qkv_projections(attn: "WanAttention", hidden_states: torch.Tensor, enco encoder_hidden_states = hidden_states if attn.fused_projections: - if attn.cross_attention_dim_head is None: + if not attn.is_cross_attention: # In self-attention layers, we can fuse the entire QKV projection into a single linear query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1) else: @@ -79,9 +79,9 @@ def __call__( self, attn: "WanAttention", hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: encoder_hidden_states_img = None if attn.add_k_proj is not None: @@ -134,7 +134,8 @@ def apply_rotary_emb( dropout_p=0.0, is_causal=False, backend=self._attention_backend, - parallel_config=self._parallel_config, + # Reference: https://github.com/huggingface/diffusers/pull/12909 + parallel_config=None, ) hidden_states_img = hidden_states_img.flatten(2, 3) hidden_states_img = hidden_states_img.type_as(query) @@ -147,7 +148,8 @@ def apply_rotary_emb( dropout_p=0.0, is_causal=False, backend=self._attention_backend, - parallel_config=self._parallel_config, + # Reference: https://github.com/huggingface/diffusers/pull/12909 + parallel_config=(self._parallel_config if encoder_hidden_states is None else None), ) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.type_as(query) @@ -181,8 +183,8 @@ def __init__( dim_head: int = 64, eps: float = 1e-5, dropout: float = 0.0, - added_kv_proj_dim: Optional[int] = None, - cross_attention_dim_head: Optional[int] = None, + added_kv_proj_dim: int | None = None, + cross_attention_dim_head: int | None = None, processor=None, is_cross_attention=None, ): @@ -212,7 +214,10 @@ def __init__( self.add_v_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=True) self.norm_added_k = torch.nn.RMSNorm(dim_head * heads, eps=eps) - self.is_cross_attention = cross_attention_dim_head is not None + if is_cross_attention is not None: + self.is_cross_attention = is_cross_attention + else: + self.is_cross_attention = cross_attention_dim_head is not None self.set_processor(processor) @@ -220,7 +225,7 @@ def fuse_projections(self): if getattr(self, "fused_projections", False): return - if self.cross_attention_dim_head is None: + if not self.is_cross_attention: concatenated_weights = torch.cat([self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data]) concatenated_bias = torch.cat([self.to_q.bias.data, self.to_k.bias.data, self.to_v.bias.data]) out_features, in_features = concatenated_weights.shape @@ -268,9 +273,9 @@ def unfuse_projections(self): def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, **kwargs, ) -> torch.Tensor: return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, rotary_emb, **kwargs) @@ -307,8 +312,8 @@ def __init__( time_freq_dim: int, time_proj_dim: int, text_embed_dim: int, - image_embed_dim: Optional[int] = None, - pos_embed_seq_len: Optional[int] = None, + image_embed_dim: int | None = None, + pos_embed_seq_len: int | None = None, ): super().__init__() @@ -326,8 +331,8 @@ def forward( self, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, - encoder_hidden_states_image: Optional[torch.Tensor] = None, - timestep_seq_len: Optional[int] = None, + encoder_hidden_states_image: torch.Tensor | None = None, + timestep_seq_len: int | None = None, ): timestep = self.timesteps_proj(timestep) if timestep_seq_len is not None: @@ -350,7 +355,7 @@ class WanRotaryPosEmbed(nn.Module): def __init__( self, attention_head_dim: int, - patch_size: Tuple[int, int, int], + patch_size: tuple[int, int, int], max_seq_len: int, theta: float = 10000.0, ): @@ -421,7 +426,7 @@ def __init__( qk_norm: str = "rms_norm_across_heads", cross_attn_norm: bool = False, eps: float = 1e-6, - added_kv_proj_dim: Optional[int] = None, + added_kv_proj_dim: int | None = None, ): super().__init__() @@ -506,7 +511,7 @@ class WanTransformer3DModel( A Transformer model for video-like data used in the Wan model. Args: - patch_size (`Tuple[int]`, defaults to `(1, 2, 2)`): + patch_size (`tuple[int]`, defaults to `(1, 2, 2)`): 3D patch dimensions for video embedding (t_patch, h_patch, w_patch). num_attention_heads (`int`, defaults to `40`): Fixed length for text embeddings. @@ -524,7 +529,7 @@ class WanTransformer3DModel( Intermediate dimension in feed-forward network. num_layers (`int`, defaults to `40`): The number of layers of transformer blocks to use. - window_size (`Tuple[int]`, defaults to `(-1, -1)`): + window_size (`tuple[int]`, defaults to `(-1, -1)`): Window size for local attention (-1 indicates global attention). cross_attn_norm (`bool`, defaults to `True`): Enable cross-attention normalization. @@ -552,9 +557,11 @@ class WanTransformer3DModel( "blocks.0": { "hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), }, - "blocks.*": { - "encoder_hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), - }, + # Reference: https://github.com/huggingface/diffusers/pull/12909 + # We need to disable the splitting of encoder_hidden_states because the image_encoder + # (Wan 2.1 I2V) consistently generates 257 tokens for image_embed. This causes the shape + # of encoder_hidden_states—whose token count is always 769 (512 + 257) after concatenation + # —to be indivisible by the number of devices in the CP. "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3), "": { "timestep": ContextParallelInput(split_dim=1, expected_dims=2, split_output=False), @@ -564,7 +571,7 @@ class WanTransformer3DModel( @register_to_config def __init__( self, - patch_size: Tuple[int, ...] = (1, 2, 2), + patch_size: tuple[int, ...] = (1, 2, 2), num_attention_heads: int = 40, attention_head_dim: int = 128, in_channels: int = 16, @@ -574,12 +581,12 @@ def __init__( ffn_dim: int = 13824, num_layers: int = 40, cross_attn_norm: bool = True, - qk_norm: Optional[str] = "rms_norm_across_heads", + qk_norm: str | None = "rms_norm_across_heads", eps: float = 1e-6, - image_dim: Optional[int] = None, - added_kv_proj_dim: Optional[int] = None, + image_dim: int | None = None, + added_kv_proj_dim: int | None = None, rope_max_seq_len: int = 1024, - pos_embed_seq_len: Optional[int] = None, + pos_embed_seq_len: int | None = None, ) -> None: super().__init__() @@ -618,30 +625,16 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, timestep: torch.LongTensor, encoder_hidden_states: torch.Tensor, - encoder_hidden_states_image: Optional[torch.Tensor] = None, + encoder_hidden_states_image: torch.Tensor | None = None, return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) - + attention_kwargs: dict[str, Any] | None = None, + ) -> torch.Tensor | dict[str, torch.Tensor]: batch_size, num_channels, num_frames, height, width = hidden_states.shape p_t, p_h, p_w = self.config.patch_size post_patch_num_frames = num_frames // p_t @@ -709,10 +702,6 @@ def forward( hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) diff --git a/src/diffusers/models/transformers/transformer_wan_animate.py b/src/diffusers/models/transformers/transformer_wan_animate.py index 6a47a67385a3..b5b15fe06099 100644 --- a/src/diffusers/models/transformers/transformer_wan_animate.py +++ b/src/diffusers/models/transformers/transformer_wan_animate.py @@ -13,7 +13,7 @@ # limitations under the License. import math -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -21,7 +21,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, logging from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward from ..attention_dispatch import dispatch_attention_fn from ..cache_utils import CacheMixin @@ -54,7 +54,7 @@ def _get_qkv_projections(attn: "WanAttention", hidden_states: torch.Tensor, enco encoder_hidden_states = hidden_states if attn.fused_projections: - if attn.cross_attention_dim_head is None: + if not attn.is_cross_attention: # In self-attention layers, we can fuse the entire QKV projection into a single linear query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1) else: @@ -83,7 +83,7 @@ class FusedLeakyReLU(nn.Module): Fused LeakyRelu with scale factor and channel-wise bias. """ - def __init__(self, negative_slope: float = 0.2, scale: float = 2**0.5, bias_channels: Optional[int] = None): + def __init__(self, negative_slope: float = 0.2, scale: float = 2**0.5, bias_channels: int | None = None): super().__init__() self.negative_slope = negative_slope self.scale = scale @@ -117,7 +117,7 @@ def __init__( stride: int = 1, padding: int = 0, bias: bool = True, - blur_kernel: Optional[Tuple[int, ...]] = None, + blur_kernel: tuple[int, ...] | None = None, blur_upsample_factor: int = 1, use_activation: bool = True, ): @@ -166,9 +166,11 @@ def forward(self, x: torch.Tensor, channel_dim: int = 1) -> torch.Tensor: # NOTE: the original implementation uses a 2D upfirdn operation with the upsampling and downsampling rates # set to 1, which should be equivalent to a 2D convolution expanded_kernel = self.blur_kernel[None, None, :, :].expand(self.in_channels, 1, -1, -1) + x = x.to(expanded_kernel.dtype) x = F.conv2d(x, expanded_kernel, padding=self.blur_padding, groups=self.in_channels) # Main Conv2D with scaling + x = x.to(self.weight.dtype) x = F.conv2d(x, self.weight * self.scale, bias=self.bias, stride=self.stride, padding=self.padding) # Activation with fused bias, if using @@ -229,7 +231,7 @@ def __init__( out_channels: int, kernel_size: int = 3, kernel_size_skip: int = 1, - blur_kernel: Tuple[int, ...] = (1, 3, 3, 1), + blur_kernel: tuple[int, ...] = (1, 3, 3, 1), downsample_factor: int = 2, ): super().__init__() @@ -286,7 +288,7 @@ def __init__( motion_dim: int = 20, out_dim: int = 512, motion_blocks: int = 5, - channels: Optional[Dict[str, int]] = None, + channels: dict[str, int] | None = None, ): super().__init__() self.size = size @@ -433,8 +435,8 @@ def __call__( self, attn: "WanAnimateFaceBlockCrossAttention", hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: # encoder_hidden_states corresponds to the motion vec # attention_mask corresponds to the motion mask (if any) @@ -499,14 +501,17 @@ def __init__( heads: int = 8, dim_head: int = 64, eps: float = 1e-6, - cross_attention_dim_head: Optional[int] = None, + cross_attention_dim_head: int | None = None, + bias: bool = True, processor=None, ): super().__init__() self.inner_dim = dim_head * heads self.heads = heads - self.cross_attention_head_dim = cross_attention_dim_head + self.cross_attention_dim_head = cross_attention_dim_head self.kv_inner_dim = self.inner_dim if cross_attention_dim_head is None else cross_attention_dim_head * heads + self.use_bias = bias + self.is_cross_attention = cross_attention_dim_head is not None # 1. Pre-Attention Norms for the hidden_states (video latents) and encoder_hidden_states (motion vector). # NOTE: this is not used in "vanilla" WanAttention @@ -514,10 +519,10 @@ def __init__( self.pre_norm_kv = nn.LayerNorm(dim, eps, elementwise_affine=False) # 2. QKV and Output Projections - self.to_q = torch.nn.Linear(dim, self.inner_dim, bias=True) - self.to_k = torch.nn.Linear(dim, self.kv_inner_dim, bias=True) - self.to_v = torch.nn.Linear(dim, self.kv_inner_dim, bias=True) - self.to_out = torch.nn.Linear(self.inner_dim, dim, bias=True) + self.to_q = torch.nn.Linear(dim, self.inner_dim, bias=bias) + self.to_k = torch.nn.Linear(dim, self.kv_inner_dim, bias=bias) + self.to_v = torch.nn.Linear(dim, self.kv_inner_dim, bias=bias) + self.to_out = torch.nn.Linear(self.inner_dim, dim, bias=bias) # 3. QK Norm # NOTE: this is applied after the reshape, so only over dim_head rather than dim_head * heads @@ -532,8 +537,8 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: return self.processor(self, hidden_states, encoder_hidden_states, attention_mask) @@ -554,9 +559,9 @@ def __call__( self, attn: "WanAttention", hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: encoder_hidden_states_img = None if attn.add_k_proj is not None: @@ -609,7 +614,8 @@ def apply_rotary_emb( dropout_p=0.0, is_causal=False, backend=self._attention_backend, - parallel_config=self._parallel_config, + # Reference: https://github.com/huggingface/diffusers/pull/12909 + parallel_config=None, ) hidden_states_img = hidden_states_img.flatten(2, 3) hidden_states_img = hidden_states_img.type_as(query) @@ -622,7 +628,8 @@ def apply_rotary_emb( dropout_p=0.0, is_causal=False, backend=self._attention_backend, - parallel_config=self._parallel_config, + # Reference: https://github.com/huggingface/diffusers/pull/12909 + parallel_config=(self._parallel_config if encoder_hidden_states is None else None), ) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.type_as(query) @@ -647,8 +654,8 @@ def __init__( dim_head: int = 64, eps: float = 1e-5, dropout: float = 0.0, - added_kv_proj_dim: Optional[int] = None, - cross_attention_dim_head: Optional[int] = None, + added_kv_proj_dim: int | None = None, + cross_attention_dim_head: int | None = None, processor=None, is_cross_attention=None, ): @@ -678,7 +685,10 @@ def __init__( self.add_v_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=True) self.norm_added_k = torch.nn.RMSNorm(dim_head * heads, eps=eps) - self.is_cross_attention = cross_attention_dim_head is not None + if is_cross_attention is not None: + self.is_cross_attention = is_cross_attention + else: + self.is_cross_attention = cross_attention_dim_head is not None self.set_processor(processor) @@ -686,7 +696,7 @@ def fuse_projections(self): if getattr(self, "fused_projections", False): return - if self.cross_attention_dim_head is None: + if not self.is_cross_attention: concatenated_weights = torch.cat([self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data]) concatenated_bias = torch.cat([self.to_q.bias.data, self.to_k.bias.data, self.to_v.bias.data]) out_features, in_features = concatenated_weights.shape @@ -734,9 +744,9 @@ def unfuse_projections(self): def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, **kwargs, ) -> torch.Tensor: return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, rotary_emb, **kwargs) @@ -767,7 +777,7 @@ def forward(self, encoder_hidden_states_image: torch.Tensor) -> torch.Tensor: return hidden_states -# Copied from diffusers.models.transformers.transformer_wan.WanTimeTextImageEmbedding +# Modified from diffusers.models.transformers.transformer_wan.WanTimeTextImageEmbedding class WanTimeTextImageEmbedding(nn.Module): def __init__( self, @@ -775,8 +785,8 @@ def __init__( time_freq_dim: int, time_proj_dim: int, text_embed_dim: int, - image_embed_dim: Optional[int] = None, - pos_embed_seq_len: Optional[int] = None, + image_embed_dim: int | None = None, + pos_embed_seq_len: int | None = None, ): super().__init__() @@ -794,17 +804,19 @@ def forward( self, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, - encoder_hidden_states_image: Optional[torch.Tensor] = None, - timestep_seq_len: Optional[int] = None, + encoder_hidden_states_image: torch.Tensor | None = None, + timestep_seq_len: int | None = None, ): timestep = self.timesteps_proj(timestep) if timestep_seq_len is not None: timestep = timestep.unflatten(0, (-1, timestep_seq_len)) - time_embedder_dtype = next(iter(self.time_embedder.parameters())).dtype - if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8: - timestep = timestep.to(time_embedder_dtype) - temb = self.time_embedder(timestep).type_as(encoder_hidden_states) + if self.time_embedder.linear_1.weight.dtype.is_floating_point: + time_embedder_dtype = self.time_embedder.linear_1.weight.dtype + else: + time_embedder_dtype = encoder_hidden_states.dtype + + temb = self.time_embedder(timestep.to(time_embedder_dtype)).type_as(encoder_hidden_states) timestep_proj = self.time_proj(self.act_fn(temb)) encoder_hidden_states = self.text_embedder(encoder_hidden_states) @@ -819,7 +831,7 @@ class WanRotaryPosEmbed(nn.Module): def __init__( self, attention_head_dim: int, - patch_size: Tuple[int, int, int], + patch_size: tuple[int, int, int], max_seq_len: int, theta: float = 10000.0, ): @@ -890,7 +902,7 @@ def __init__( qk_norm: str = "rms_norm_across_heads", cross_attn_norm: bool = False, eps: float = 1e-6, - added_kv_proj_dim: Optional[int] = None, + added_kv_proj_dim: int | None = None, ): super().__init__() @@ -975,7 +987,7 @@ class WanAnimateTransformer3DModel( A Transformer model for video-like data used in the WanAnimate model. Args: - patch_size (`Tuple[int]`, defaults to `(1, 2, 2)`): + patch_size (`tuple[int]`, defaults to `(1, 2, 2)`): 3D patch dimensions for video embedding (t_patch, h_patch, w_patch). num_attention_heads (`int`, defaults to `40`): Fixed length for text embeddings. @@ -993,7 +1005,7 @@ class WanAnimateTransformer3DModel( Intermediate dimension in feed-forward network. num_layers (`int`, defaults to `40`): The number of layers of transformer blocks to use. - window_size (`Tuple[int]`, defaults to `(-1, -1)`): + window_size (`tuple[int]`, defaults to `(-1, -1)`): Window size for local attention (-1 indicates global attention). cross_attn_norm (`bool`, defaults to `True`): Enable cross-attention normalization. @@ -1024,24 +1036,24 @@ class WanAnimateTransformer3DModel( @register_to_config def __init__( self, - patch_size: Tuple[int] = (1, 2, 2), + patch_size: tuple[int] = (1, 2, 2), num_attention_heads: int = 40, attention_head_dim: int = 128, - in_channels: Optional[int] = 36, - latent_channels: Optional[int] = 16, - out_channels: Optional[int] = 16, + in_channels: int | None = 36, + latent_channels: int | None = 16, + out_channels: int | None = 16, text_dim: int = 4096, freq_dim: int = 256, ffn_dim: int = 13824, num_layers: int = 40, cross_attn_norm: bool = True, - qk_norm: Optional[str] = "rms_norm_across_heads", + qk_norm: str | None = "rms_norm_across_heads", eps: float = 1e-6, - image_dim: Optional[int] = 1280, - added_kv_proj_dim: Optional[int] = None, + image_dim: int | None = 1280, + added_kv_proj_dim: int | None = None, rope_max_seq_len: int = 1024, - pos_embed_seq_len: Optional[int] = None, - motion_encoder_channel_sizes: Optional[Dict[str, int]] = None, # Start of Wan Animate-specific args + pos_embed_seq_len: int | None = None, + motion_encoder_channel_sizes: dict[str, int] | None = None, # Start of Wan Animate-specific args motion_encoder_size: int = 512, motion_style_dim: int = 512, motion_dim: int = 20, @@ -1135,18 +1147,19 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, timestep: torch.LongTensor, encoder_hidden_states: torch.Tensor, - encoder_hidden_states_image: Optional[torch.Tensor] = None, - pose_hidden_states: Optional[torch.Tensor] = None, - face_pixel_values: Optional[torch.Tensor] = None, - motion_encode_batch_size: Optional[int] = None, + encoder_hidden_states_image: torch.Tensor | None = None, + pose_hidden_states: torch.Tensor | None = None, + face_pixel_values: torch.Tensor | None = None, + motion_encode_batch_size: int | None = None, return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + attention_kwargs: dict[str, Any] | None = None, + ) -> torch.Tensor | dict[str, torch.Tensor]: """ Forward pass of Wan2.2-Animate transformer model. @@ -1173,21 +1186,6 @@ def forward( Whether to return the output as a dict or tuple. """ - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) - # Check that shapes match up if pose_hidden_states is not None and pose_hidden_states.shape[2] + 1 != hidden_states.shape[2]: raise ValueError( @@ -1288,10 +1286,6 @@ def forward( hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) diff --git a/src/diffusers/models/transformers/transformer_wan_vace.py b/src/diffusers/models/transformers/transformer_wan_vace.py index 1be4f73e33e2..7c2e205ee3ed 100644 --- a/src/diffusers/models/transformers/transformer_wan_vace.py +++ b/src/diffusers/models/transformers/transformer_wan_vace.py @@ -13,14 +13,14 @@ # limitations under the License. import math -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin -from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ...utils import apply_lora_scale, logging from ..attention import AttentionMixin, FeedForward from ..cache_utils import CacheMixin from ..modeling_outputs import Transformer2DModelOutput @@ -47,7 +47,7 @@ def __init__( qk_norm: str = "rms_norm_across_heads", cross_attn_norm: bool = False, eps: float = 1e-6, - added_kv_proj_dim: Optional[int] = None, + added_kv_proj_dim: int | None = None, apply_input_projection: bool = False, apply_output_projection: bool = False, ): @@ -76,6 +76,7 @@ def __init__( eps=eps, added_kv_proj_dim=added_kv_proj_dim, processor=WanAttnProcessor(), + is_cross_attention=True, ) self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() @@ -141,7 +142,7 @@ class WanVACETransformer3DModel( A Transformer model for video-like data used in the Wan model. Args: - patch_size (`Tuple[int]`, defaults to `(1, 2, 2)`): + patch_size (`tuple[int]`, defaults to `(1, 2, 2)`): 3D patch dimensions for video embedding (t_patch, h_patch, w_patch). num_attention_heads (`int`, defaults to `40`): Fixed length for text embeddings. @@ -159,7 +160,7 @@ class WanVACETransformer3DModel( Intermediate dimension in feed-forward network. num_layers (`int`, defaults to `40`): The number of layers of transformer blocks to use. - window_size (`Tuple[int]`, defaults to `(-1, -1)`): + window_size (`tuple[int]`, defaults to `(-1, -1)`): Window size for local attention (-1 indicates global attention). cross_attn_norm (`bool`, defaults to `True`): Enable cross-attention normalization. @@ -178,11 +179,12 @@ class WanVACETransformer3DModel( _no_split_modules = ["WanTransformerBlock", "WanVACETransformerBlock"] _keep_in_fp32_modules = ["time_embedder", "scale_shift_table", "norm1", "norm2", "norm3"] _keys_to_ignore_on_load_unexpected = ["norm_added_q"] + _repeated_blocks = ["WanTransformerBlock", "WanVACETransformerBlock"] @register_to_config def __init__( self, - patch_size: Tuple[int, ...] = (1, 2, 2), + patch_size: tuple[int, ...] = (1, 2, 2), num_attention_heads: int = 40, attention_head_dim: int = 128, in_channels: int = 16, @@ -192,13 +194,13 @@ def __init__( ffn_dim: int = 13824, num_layers: int = 40, cross_attn_norm: bool = True, - qk_norm: Optional[str] = "rms_norm_across_heads", + qk_norm: str | None = "rms_norm_across_heads", eps: float = 1e-6, - image_dim: Optional[int] = None, - added_kv_proj_dim: Optional[int] = None, + image_dim: int | None = None, + added_kv_proj_dim: int | None = None, rope_max_seq_len: int = 1024, - pos_embed_seq_len: Optional[int] = None, - vace_layers: List[int] = [0, 5, 10, 15, 20, 25, 30, 35], + pos_embed_seq_len: int | None = None, + vace_layers: list[int] = [0, 5, 10, 15, 20, 25, 30, 35], vace_in_channels: int = 96, ) -> None: super().__init__() @@ -261,32 +263,18 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, timestep: torch.LongTensor, encoder_hidden_states: torch.Tensor, - encoder_hidden_states_image: Optional[torch.Tensor] = None, + encoder_hidden_states_image: torch.Tensor | None = None, control_hidden_states: torch.Tensor = None, control_hidden_states_scale: torch.Tensor = None, return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: - if attention_kwargs is not None: - attention_kwargs = attention_kwargs.copy() - lora_scale = attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - else: - if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None: - logger.warning( - "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." - ) - + attention_kwargs: dict[str, Any] | None = None, + ) -> torch.Tensor | dict[str, torch.Tensor]: batch_size, num_channels, num_frames, height, width = hidden_states.shape p_t, p_h, p_w = self.config.patch_size post_patch_num_frames = num_frames // p_t @@ -379,10 +367,6 @@ def forward( hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (output,) diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py index 5983c34ab640..3bbf78bc5e01 100644 --- a/src/diffusers/models/transformers/transformer_z_image.py +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -13,7 +13,6 @@ # limitations under the License. import math -from typing import Dict, List, Optional, Tuple, Union import torch import torch.nn as nn @@ -92,9 +91,9 @@ def __call__( self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - freqs_cis: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + freqs_cis: torch.Tensor | None = None, ) -> torch.Tensor: query = attn.to_q(hidden_states) key = attn.to_k(hidden_states) @@ -229,10 +228,10 @@ def forward( x: torch.Tensor, attn_mask: torch.Tensor, freqs_cis: torch.Tensor, - adaln_input: Optional[torch.Tensor] = None, - noise_mask: Optional[torch.Tensor] = None, - adaln_noisy: Optional[torch.Tensor] = None, - adaln_clean: Optional[torch.Tensor] = None, + adaln_input: torch.Tensor | None = None, + noise_mask: torch.Tensor | None = None, + adaln_noisy: torch.Tensor | None = None, + adaln_clean: torch.Tensor | None = None, ): if self.modulation: seq_len = x.shape[1] @@ -315,8 +314,8 @@ class RopeEmbedder: def __init__( self, theta: float = 256.0, - axes_dims: List[int] = (16, 56, 56), - axes_lens: List[int] = (64, 128, 128), + axes_dims: list[int] = (16, 56, 56), + axes_lens: list[int] = (64, 128, 128), ): self.theta = theta self.axes_dims = axes_dims @@ -325,7 +324,7 @@ def __init__( self.freqs_cis = None @staticmethod - def precompute_freqs_cis(dim: List[int], end: List[int], theta: float = 256.0): + def precompute_freqs_cis(dim: list[int], end: list[int], theta: float = 256.0): with torch.device("cpu"): freqs_cis = [] for i, (d, e) in enumerate(zip(dim, end)): @@ -482,12 +481,12 @@ def __init__( def unpatchify( self, - x: List[torch.Tensor], - size: List[Tuple], + x: list[torch.Tensor], + size: list[tuple], patch_size, f_patch_size, - x_pos_offsets: Optional[List[Tuple[int, int]]] = None, - ) -> List[torch.Tensor]: + x_pos_offsets: list[tuple[int, int]] | None = None, + ) -> list[torch.Tensor]: pH = pW = patch_size pF = f_patch_size bsz = len(x) @@ -552,10 +551,10 @@ def _patchify_image(self, image: torch.Tensor, patch_size: int, f_patch_size: in def _pad_with_ids( self, feat: torch.Tensor, - pos_grid_size: Tuple, - pos_start: Tuple, + pos_grid_size: tuple, + pos_start: tuple, device: torch.device, - noise_mask_val: Optional[int] = None, + noise_mask_val: int | None = None, ): """Pad feature to SEQ_MULTI_OF, create position IDs and pad mask.""" ori_len = len(feat) @@ -587,7 +586,7 @@ def _pad_with_ids( return padded_feat, pos_ids, pad_mask, total_len, noise_mask def patchify_and_embed( - self, all_image: List[torch.Tensor], all_cap_feats: List[torch.Tensor], patch_size: int, f_patch_size: int + self, all_image: list[torch.Tensor], all_cap_feats: list[torch.Tensor], patch_size: int, f_patch_size: int ): """Patchify for basic mode: single image per batch item.""" device = all_image[0].device @@ -625,12 +624,12 @@ def patchify_and_embed( def patchify_and_embed_omni( self, - all_x: List[List[torch.Tensor]], - all_cap_feats: List[List[torch.Tensor]], - all_siglip_feats: List[List[torch.Tensor]], + all_x: list[list[torch.Tensor]], + all_cap_feats: list[list[torch.Tensor]], + all_siglip_feats: list[list[torch.Tensor]], patch_size: int, f_patch_size: int, - images_noise_mask: List[List[int]], + images_noise_mask: list[list[int]], ): """Patchify for omni mode: multiple images per batch item with noise masks.""" bsz = len(all_x) @@ -764,11 +763,11 @@ def patchify_and_embed_omni( def _prepare_sequence( self, - feats: List[torch.Tensor], - pos_ids: List[torch.Tensor], - inner_pad_mask: List[torch.Tensor], + feats: list[torch.Tensor], + pos_ids: list[torch.Tensor], + inner_pad_mask: list[torch.Tensor], pad_token: torch.nn.Parameter, - noise_mask: Optional[List[List[int]]] = None, + noise_mask: list[list[int]] | None = None, device: torch.device = None, ): """Prepare sequence: apply pad token, RoPE embed, pad to batch, create attention mask.""" @@ -808,16 +807,16 @@ def _build_unified_sequence( self, x: torch.Tensor, x_freqs: torch.Tensor, - x_seqlens: List[int], - x_noise_mask: Optional[List[List[int]]], + x_seqlens: list[int], + x_noise_mask: list[list[int]] | None, cap: torch.Tensor, cap_freqs: torch.Tensor, - cap_seqlens: List[int], - cap_noise_mask: Optional[List[List[int]]], - siglip: Optional[torch.Tensor], - siglip_freqs: Optional[torch.Tensor], - siglip_seqlens: Optional[List[int]], - siglip_noise_mask: Optional[List[List[int]]], + cap_seqlens: list[int], + cap_noise_mask: list[list[int]] | None, + siglip: torch.Tensor | None, + siglip_freqs: torch.Tensor | None, + siglip_seqlens: list[int] | None, + siglip_noise_mask: list[list[int]] | None, omni_mode: bool, device: torch.device, ): @@ -887,13 +886,13 @@ def _build_unified_sequence( def forward( self, - x: Union[List[torch.Tensor], List[List[torch.Tensor]]], + x: list[torch.Tensor, list[list[torch.Tensor]]], t, - cap_feats: Union[List[torch.Tensor], List[List[torch.Tensor]]], + cap_feats: list[torch.Tensor, list[list[torch.Tensor]]], return_dict: bool = True, - controlnet_block_samples: Optional[Dict[int, torch.Tensor]] = None, - siglip_feats: Optional[List[List[torch.Tensor]]] = None, - image_noise_mask: Optional[List[List[int]]] = None, + controlnet_block_samples: dict[int, torch.Tensor] | None = None, + siglip_feats: list[list[torch.Tensor]] | None = None, + image_noise_mask: list[list[int]] | None = None, patch_size: int = 2, f_patch_size: int = 1, ): diff --git a/src/diffusers/models/unets/unet_1d.py b/src/diffusers/models/unets/unet_1d.py index a027c553ed06..83ffe1f6f8cb 100644 --- a/src/diffusers/models/unets/unet_1d.py +++ b/src/diffusers/models/unets/unet_1d.py @@ -13,7 +13,6 @@ # limitations under the License. from dataclasses import dataclass -from typing import Optional, Tuple, Union import torch import torch.nn as nn @@ -56,12 +55,12 @@ class UNet1DModel(ModelMixin, ConfigMixin): freq_shift (`float`, *optional*, defaults to 0.0): Frequency shift for Fourier time embedding. flip_sin_to_cos (`bool`, *optional*, defaults to `False`): Whether to flip sin to cos for Fourier time embedding. - down_block_types (`Tuple[str]`, *optional*, defaults to `("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D")`): - Tuple of downsample block types. - up_block_types (`Tuple[str]`, *optional*, defaults to `("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip")`): - Tuple of upsample block types. - block_out_channels (`Tuple[int]`, *optional*, defaults to `(32, 32, 64)`): - Tuple of block output channels. + down_block_types (`tuple[str]`, *optional*, defaults to `("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D")`): + tuple of downsample block types. + up_block_types (`tuple[str]`, *optional*, defaults to `("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip")`): + tuple of upsample block types. + block_out_channels (`tuple[int]`, *optional*, defaults to `(32, 32, 64)`): + tuple of block output channels. mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock1D"`): Block type for middle of UNet. out_block_type (`str`, *optional*, defaults to `None`): Optional output processing block of UNet. act_fn (`str`, *optional*, defaults to `None`): Optional activation function in UNet blocks. @@ -77,20 +76,20 @@ class UNet1DModel(ModelMixin, ConfigMixin): def __init__( self, sample_size: int = 65536, - sample_rate: Optional[int] = None, + sample_rate: int | None = None, in_channels: int = 2, out_channels: int = 2, extra_in_channels: int = 0, time_embedding_type: str = "fourier", - time_embedding_dim: Optional[int] = None, + time_embedding_dim: int | None = None, flip_sin_to_cos: bool = True, use_timestep_embedding: bool = False, freq_shift: float = 0.0, - down_block_types: Tuple[str, ...] = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D"), - up_block_types: Tuple[str, ...] = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip"), + down_block_types: tuple[str, ...] = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D"), + up_block_types: tuple[str, ...] = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip"), mid_block_type: str = "UNetMidBlock1D", out_block_type: str = None, - block_out_channels: Tuple[int, ...] = (32, 32, 64), + block_out_channels: tuple[int, ...] = (32, 32, 64), act_fn: str = None, norm_num_groups: int = 8, layers_per_block: int = 1, @@ -206,9 +205,9 @@ def __init__( def forward( self, sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], + timestep: torch.Tensor | float | int, return_dict: bool = True, - ) -> Union[UNet1DOutput, Tuple]: + ) -> UNet1DOutput | tuple: r""" The [`UNet1DModel`] forward method. diff --git a/src/diffusers/models/unets/unet_1d_blocks.py b/src/diffusers/models/unets/unet_1d_blocks.py index 58cbdfd005b6..a21521eaed4c 100644 --- a/src/diffusers/models/unets/unet_1d_blocks.py +++ b/src/diffusers/models/unets/unet_1d_blocks.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. import math -from typing import Optional, Tuple, Union import torch import torch.nn.functional as F @@ -26,13 +25,13 @@ class DownResnetBlock1D(nn.Module): def __init__( self, in_channels: int, - out_channels: Optional[int] = None, + out_channels: int | None = None, num_layers: int = 1, conv_shortcut: bool = False, temb_channels: int = 32, groups: int = 32, - groups_out: Optional[int] = None, - non_linearity: Optional[str] = None, + groups_out: int | None = None, + non_linearity: str | None = None, time_embedding_norm: str = "default", output_scale_factor: float = 1.0, add_downsample: bool = True, @@ -66,7 +65,7 @@ def __init__( if add_downsample: self.downsample = Downsample1D(out_channels, use_conv=True, padding=1) - def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor: output_states = () hidden_states = self.resnets[0](hidden_states, temb) @@ -88,12 +87,12 @@ class UpResnetBlock1D(nn.Module): def __init__( self, in_channels: int, - out_channels: Optional[int] = None, + out_channels: int | None = None, num_layers: int = 1, temb_channels: int = 32, groups: int = 32, - groups_out: Optional[int] = None, - non_linearity: Optional[str] = None, + groups_out: int | None = None, + non_linearity: str | None = None, time_embedding_norm: str = "default", output_scale_factor: float = 1.0, add_upsample: bool = True, @@ -129,8 +128,8 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Optional[Tuple[torch.Tensor, ...]] = None, - temb: Optional[torch.Tensor] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...] | None = None, + temb: torch.Tensor | None = None, ) -> torch.Tensor: if res_hidden_states_tuple is not None: res_hidden_states = res_hidden_states_tuple[-1] @@ -161,7 +160,7 @@ def __init__(self, in_channels: int, out_channels: int, embed_dim: int): self.res2 = ResidualTemporalBlock1D(in_channels // 2, in_channels // 4, embed_dim=embed_dim) self.down2 = Downsample1D(out_channels // 4, use_conv=True) - def forward(self, x: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: + def forward(self, x: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor: x = self.res1(x, temb) x = self.down1(x) x = self.res2(x, temb) @@ -178,7 +177,7 @@ def __init__( num_layers: int = 1, add_downsample: bool = False, add_upsample: bool = False, - non_linearity: Optional[str] = None, + non_linearity: str | None = None, ): super().__init__() self.in_channels = in_channels @@ -230,7 +229,7 @@ def __init__(self, num_groups_out: int, out_channels: int, embed_dim: int, act_f self.final_conv1d_act = get_activation(act_fn) self.final_conv1d_2 = nn.Conv1d(embed_dim, out_channels, 1) - def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor: hidden_states = self.final_conv1d_1(hidden_states) hidden_states = rearrange_dims(hidden_states) hidden_states = self.final_conv1d_gn(hidden_states) @@ -305,7 +304,7 @@ def __init__(self, kernel: str = "linear", pad_mode: str = "reflect"): self.pad = kernel_1d.shape[0] // 2 - 1 self.register_buffer("kernel", kernel_1d) - def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor: hidden_states = F.pad(hidden_states, ((self.pad + 1) // 2,) * 2, self.pad_mode) weight = hidden_states.new_zeros([hidden_states.shape[1], hidden_states.shape[1], self.kernel.shape[0]]) indices = torch.arange(hidden_states.shape[1], device=hidden_states.device) @@ -407,7 +406,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class UNetMidBlock1D(nn.Module): - def __init__(self, mid_channels: int, in_channels: int, out_channels: Optional[int] = None): + def __init__(self, mid_channels: int, in_channels: int, out_channels: int | None = None): super().__init__() out_channels = in_channels if out_channels is None else out_channels @@ -435,7 +434,7 @@ def __init__(self, mid_channels: int, in_channels: int, out_channels: Optional[i self.attentions = nn.ModuleList(attentions) self.resnets = nn.ModuleList(resnets) - def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor: hidden_states = self.down(hidden_states) for attn, resnet in zip(self.attentions, self.resnets): hidden_states = resnet(hidden_states) @@ -447,7 +446,7 @@ def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = No class AttnDownBlock1D(nn.Module): - def __init__(self, out_channels: int, in_channels: int, mid_channels: Optional[int] = None): + def __init__(self, out_channels: int, in_channels: int, mid_channels: int | None = None): super().__init__() mid_channels = out_channels if mid_channels is None else mid_channels @@ -466,7 +465,7 @@ def __init__(self, out_channels: int, in_channels: int, mid_channels: Optional[i self.attentions = nn.ModuleList(attentions) self.resnets = nn.ModuleList(resnets) - def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor: hidden_states = self.down(hidden_states) for resnet, attn in zip(self.resnets, self.attentions): @@ -477,7 +476,7 @@ def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = No class DownBlock1D(nn.Module): - def __init__(self, out_channels: int, in_channels: int, mid_channels: Optional[int] = None): + def __init__(self, out_channels: int, in_channels: int, mid_channels: int | None = None): super().__init__() mid_channels = out_channels if mid_channels is None else mid_channels @@ -490,7 +489,7 @@ def __init__(self, out_channels: int, in_channels: int, mid_channels: Optional[i self.resnets = nn.ModuleList(resnets) - def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor: hidden_states = self.down(hidden_states) for resnet in self.resnets: @@ -500,7 +499,7 @@ def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = No class DownBlock1DNoSkip(nn.Module): - def __init__(self, out_channels: int, in_channels: int, mid_channels: Optional[int] = None): + def __init__(self, out_channels: int, in_channels: int, mid_channels: int | None = None): super().__init__() mid_channels = out_channels if mid_channels is None else mid_channels @@ -512,7 +511,7 @@ def __init__(self, out_channels: int, in_channels: int, mid_channels: Optional[i self.resnets = nn.ModuleList(resnets) - def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor: hidden_states = torch.cat([hidden_states, temb], dim=1) for resnet in self.resnets: hidden_states = resnet(hidden_states) @@ -521,7 +520,7 @@ def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = No class AttnUpBlock1D(nn.Module): - def __init__(self, in_channels: int, out_channels: int, mid_channels: Optional[int] = None): + def __init__(self, in_channels: int, out_channels: int, mid_channels: int | None = None): super().__init__() mid_channels = out_channels if mid_channels is None else mid_channels @@ -543,8 +542,8 @@ def __init__(self, in_channels: int, out_channels: int, mid_channels: Optional[i def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, ) -> torch.Tensor: res_hidden_states = res_hidden_states_tuple[-1] hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) @@ -559,7 +558,7 @@ def forward( class UpBlock1D(nn.Module): - def __init__(self, in_channels: int, out_channels: int, mid_channels: Optional[int] = None): + def __init__(self, in_channels: int, out_channels: int, mid_channels: int | None = None): super().__init__() mid_channels = in_channels if mid_channels is None else mid_channels @@ -575,8 +574,8 @@ def __init__(self, in_channels: int, out_channels: int, mid_channels: Optional[i def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, ) -> torch.Tensor: res_hidden_states = res_hidden_states_tuple[-1] hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) @@ -590,7 +589,7 @@ def forward( class UpBlock1DNoSkip(nn.Module): - def __init__(self, in_channels: int, out_channels: int, mid_channels: Optional[int] = None): + def __init__(self, in_channels: int, out_channels: int, mid_channels: int | None = None): super().__init__() mid_channels = in_channels if mid_channels is None else mid_channels @@ -605,8 +604,8 @@ def __init__(self, in_channels: int, out_channels: int, mid_channels: Optional[i def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, ) -> torch.Tensor: res_hidden_states = res_hidden_states_tuple[-1] hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) @@ -617,10 +616,10 @@ def forward( return hidden_states -DownBlockType = Union[DownResnetBlock1D, DownBlock1D, AttnDownBlock1D, DownBlock1DNoSkip] -MidBlockType = Union[MidResTemporalBlock1D, ValueFunctionMidBlock1D, UNetMidBlock1D] -OutBlockType = Union[OutConv1DBlock, OutValueFunctionBlock] -UpBlockType = Union[UpResnetBlock1D, UpBlock1D, AttnUpBlock1D, UpBlock1DNoSkip] +DownBlockType = DownResnetBlock1D | DownBlock1D | AttnDownBlock1D | DownBlock1DNoSkip +MidBlockType = MidResTemporalBlock1D | ValueFunctionMidBlock1D | UNetMidBlock1D +OutBlockType = OutConv1DBlock | OutValueFunctionBlock +UpBlockType = UpResnetBlock1D | UpBlock1D | AttnUpBlock1D | UpBlock1DNoSkip def get_down_block( @@ -694,7 +693,7 @@ def get_mid_block( def get_out_block( *, out_block_type: str, num_groups_out: int, embed_dim: int, out_channels: int, act_fn: str, fc_dim: int -) -> Optional[OutBlockType]: +) -> OutBlockType | None: if out_block_type == "OutConv1DBlock": return OutConv1DBlock(num_groups_out, out_channels, embed_dim, act_fn) elif out_block_type == "ValueFunction": diff --git a/src/diffusers/models/unets/unet_2d.py b/src/diffusers/models/unets/unet_2d.py index 2588a9c518bd..4e54f757d120 100644 --- a/src/diffusers/models/unets/unet_2d.py +++ b/src/diffusers/models/unets/unet_2d.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. from dataclasses import dataclass -from typing import Optional, Tuple, Union import torch import torch.nn as nn @@ -45,7 +44,7 @@ class UNet2DModel(ModelMixin, ConfigMixin): for all models (such as downloading or saving). Parameters: - sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`): + sample_size (`int` or `tuple[int, int]`, *optional*, defaults to `None`): Height and width of input/output sample. Dimensions must be a multiple of `2 ** (len(block_out_channels) - 1)`. in_channels (`int`, *optional*, defaults to 3): Number of channels in the input sample. @@ -55,14 +54,14 @@ class UNet2DModel(ModelMixin, ConfigMixin): freq_shift (`int`, *optional*, defaults to 0): Frequency shift for Fourier time embedding. flip_sin_to_cos (`bool`, *optional*, defaults to `True`): Whether to flip sin to cos for Fourier time embedding. - down_block_types (`Tuple[str]`, *optional*, defaults to `("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")`): - Tuple of downsample block types. + down_block_types (`tuple[str]`, *optional*, defaults to `("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")`): + tuple of downsample block types. mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2D"`): Block type for middle of UNet, it can be either `UNetMidBlock2D` or `None`. - up_block_types (`Tuple[str]`, *optional*, defaults to `("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")`): - Tuple of upsample block types. - block_out_channels (`Tuple[int]`, *optional*, defaults to `(224, 448, 672, 896)`): - Tuple of block output channels. + up_block_types (`tuple[str]`, *optional*, defaults to `("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")`): + tuple of upsample block types. + block_out_channels (`tuple[int]`, *optional*, defaults to `(224, 448, 672, 896)`): + tuple of block output channels. layers_per_block (`int`, *optional*, defaults to `2`): The number of layers per block. mid_block_scale_factor (`float`, *optional*, defaults to `1`): The scale factor for the mid block. downsample_padding (`int`, *optional*, defaults to `1`): The padding for the downsample convolution. @@ -95,18 +94,18 @@ class UNet2DModel(ModelMixin, ConfigMixin): @register_to_config def __init__( self, - sample_size: Optional[Union[int, Tuple[int, int]]] = None, + sample_size: int | tuple[int, int] | None = None, in_channels: int = 3, out_channels: int = 3, center_input_sample: bool = False, time_embedding_type: str = "positional", - time_embedding_dim: Optional[int] = None, + time_embedding_dim: int | None = None, freq_shift: int = 0, flip_sin_to_cos: bool = True, - down_block_types: Tuple[str, ...] = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D"), - mid_block_type: Optional[str] = "UNetMidBlock2D", - up_block_types: Tuple[str, ...] = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D"), - block_out_channels: Tuple[int, ...] = (224, 448, 672, 896), + down_block_types: tuple[str, ...] = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D"), + mid_block_type: str | None = "UNetMidBlock2D", + up_block_types: tuple[str, ...] = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D"), + block_out_channels: tuple[int, ...] = (224, 448, 672, 896), layers_per_block: int = 2, mid_block_scale_factor: float = 1, downsample_padding: int = 1, @@ -114,15 +113,15 @@ def __init__( upsample_type: str = "conv", dropout: float = 0.0, act_fn: str = "silu", - attention_head_dim: Optional[int] = 8, + attention_head_dim: int | None = 8, norm_num_groups: int = 32, - attn_norm_num_groups: Optional[int] = None, + attn_norm_num_groups: int | None = None, norm_eps: float = 1e-5, resnet_time_scale_shift: str = "default", add_attention: bool = True, - class_embed_type: Optional[str] = None, - num_class_embeds: Optional[int] = None, - num_train_timesteps: Optional[int] = None, + class_embed_type: str | None = None, + num_class_embeds: int | None = None, + num_train_timesteps: int | None = None, ): super().__init__() @@ -250,10 +249,10 @@ def __init__( def forward( self, sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], - class_labels: Optional[torch.Tensor] = None, + timestep: torch.Tensor | float | int, + class_labels: torch.Tensor | None = None, return_dict: bool = True, - ) -> Union[UNet2DOutput, Tuple]: + ) -> UNet2DOutput | tuple: r""" The [`UNet2DModel`] forward method. diff --git a/src/diffusers/models/unets/unet_2d_blocks.py b/src/diffusers/models/unets/unet_2d_blocks.py index 94a9245e567c..18a0b4ec5659 100644 --- a/src/diffusers/models/unets/unet_2d_blocks.py +++ b/src/diffusers/models/unets/unet_2d_blocks.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any import numpy as np import torch @@ -50,10 +50,10 @@ def get_down_block( resnet_eps: float, resnet_act_fn: str, transformer_layers_per_block: int = 1, - num_attention_heads: Optional[int] = None, - resnet_groups: Optional[int] = None, - cross_attention_dim: Optional[int] = None, - downsample_padding: Optional[int] = None, + num_attention_heads: int | None = None, + resnet_groups: int | None = None, + cross_attention_dim: int | None = None, + downsample_padding: int | None = None, dual_cross_attention: bool = False, use_linear_projection: bool = False, only_cross_attention: bool = False, @@ -62,9 +62,9 @@ def get_down_block( attention_type: str = "default", resnet_skip_time_act: bool = False, resnet_out_scale_factor: float = 1.0, - cross_attention_norm: Optional[str] = None, - attention_head_dim: Optional[int] = None, - downsample_type: Optional[str] = None, + cross_attention_norm: str | None = None, + attention_head_dim: int | None = None, + downsample_type: str | None = None, dropout: float = 0.0, ): # If attn head dim is not defined, we default it to the number of heads @@ -258,8 +258,8 @@ def get_mid_block( resnet_groups: int, output_scale_factor: float = 1.0, transformer_layers_per_block: int = 1, - num_attention_heads: Optional[int] = None, - cross_attention_dim: Optional[int] = None, + num_attention_heads: int | None = None, + cross_attention_dim: int | None = None, dual_cross_attention: bool = False, use_linear_projection: bool = False, mid_block_only_cross_attention: bool = False, @@ -267,8 +267,8 @@ def get_mid_block( resnet_time_scale_shift: str = "default", attention_type: str = "default", resnet_skip_time_act: bool = False, - cross_attention_norm: Optional[str] = None, - attention_head_dim: Optional[int] = 1, + cross_attention_norm: str | None = None, + attention_head_dim: int | None = 1, dropout: float = 0.0, ): if mid_block_type == "UNetMidBlock2DCrossAttn": @@ -334,11 +334,11 @@ def get_up_block( add_upsample: bool, resnet_eps: float, resnet_act_fn: str, - resolution_idx: Optional[int] = None, + resolution_idx: int | None = None, transformer_layers_per_block: int = 1, - num_attention_heads: Optional[int] = None, - resnet_groups: Optional[int] = None, - cross_attention_dim: Optional[int] = None, + num_attention_heads: int | None = None, + resnet_groups: int | None = None, + cross_attention_dim: int | None = None, dual_cross_attention: bool = False, use_linear_projection: bool = False, only_cross_attention: bool = False, @@ -347,9 +347,9 @@ def get_up_block( attention_type: str = "default", resnet_skip_time_act: bool = False, resnet_out_scale_factor: float = 1.0, - cross_attention_norm: Optional[str] = None, - attention_head_dim: Optional[int] = None, - upsample_type: Optional[str] = None, + cross_attention_norm: str | None = None, + attention_head_dim: int | None = None, + upsample_type: str | None = None, dropout: float = 0.0, ) -> nn.Module: # If attn head dim is not defined, we default it to the number of heads @@ -602,7 +602,7 @@ class UNetMidBlock2D(nn.Module): resnet_act_fn (`str`, *optional*, defaults to `swish`): The activation function for the resnet blocks. resnet_groups (`int`, *optional*, defaults to 32): The number of groups to use in the group normalization layers of the resnet blocks. - attn_groups (`Optional[int]`, *optional*, defaults to None): The number of groups for the attention blocks. + attn_groups (`int | None`, *optional*, defaults to None): The number of groups for the attention blocks. resnet_pre_norm (`bool`, *optional*, defaults to `True`): Whether to use pre-normalization for the resnet blocks. add_attention (`bool`, *optional*, defaults to `True`): Whether to add attention blocks. @@ -627,7 +627,7 @@ def __init__( resnet_time_scale_shift: str = "default", # default, spatial resnet_act_fn: str = "swish", resnet_groups: int = 32, - attn_groups: Optional[int] = None, + attn_groups: int | None = None, resnet_pre_norm: bool = True, add_attention: bool = True, attention_head_dim: int = 1, @@ -733,7 +733,7 @@ def __init__( self.gradient_checkpointing = False - def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor: hidden_states = self.resnets[0](hidden_states, temb) for attn, resnet in zip(self.attentions, self.resnets[1:]): if torch.is_grad_enabled() and self.gradient_checkpointing: @@ -753,15 +753,15 @@ def __init__( self, in_channels: int, temb_channels: int, - out_channels: Optional[int] = None, + out_channels: int | None = None, dropout: float = 0.0, num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, + transformer_layers_per_block: int | tuple[int] = 1, resnet_eps: float = 1e-6, resnet_time_scale_shift: str = "default", resnet_act_fn: str = "swish", resnet_groups: int = 32, - resnet_groups_out: Optional[int] = None, + resnet_groups_out: int | None = None, resnet_pre_norm: bool = True, num_attention_heads: int = 1, output_scale_factor: float = 1.0, @@ -854,11 +854,11 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + encoder_attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: if cross_attention_kwargs is not None: if cross_attention_kwargs.get("scale", None) is not None: @@ -907,7 +907,7 @@ def __init__( cross_attention_dim: int = 1280, skip_time_act: bool = False, only_cross_attention: bool = False, - cross_attention_norm: Optional[str] = None, + cross_attention_norm: str | None = None, ): super().__init__() @@ -978,11 +978,11 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + encoder_attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} if cross_attention_kwargs.get("scale", None) is not None: @@ -1112,10 +1112,10 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - upsample_size: Optional[int] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]: + temb: torch.Tensor | None = None, + upsample_size: int | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]: cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} if cross_attention_kwargs.get("scale", None) is not None: logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.") @@ -1152,7 +1152,7 @@ def __init__( temb_channels: int, dropout: float = 0.0, num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, + transformer_layers_per_block: int | tuple[int] = 1, resnet_eps: float = 1e-6, resnet_time_scale_shift: str = "default", resnet_act_fn: str = "swish", @@ -1239,13 +1239,13 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - additional_residuals: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]: + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + encoder_attention_mask: torch.Tensor | None = None, + additional_residuals: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]: if cross_attention_kwargs is not None: if cross_attention_kwargs.get("scale", None) is not None: logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.") @@ -1344,8 +1344,8 @@ def __init__( self.gradient_checkpointing = False def forward( - self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None, *args, **kwargs - ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]: + self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None, *args, **kwargs + ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]: if len(args) > 0 or kwargs.get("scale", None) is not None: deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`." deprecate("scale", "1.0.0", deprecation_message) @@ -1635,11 +1635,11 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - skip_sample: Optional[torch.Tensor] = None, + temb: torch.Tensor | None = None, + skip_sample: torch.Tensor | None = None, *args, **kwargs, - ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...], torch.Tensor]: + ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...], torch.Tensor]: if len(args) > 0 or kwargs.get("scale", None) is not None: deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`." deprecate("scale", "1.0.0", deprecation_message) @@ -1726,11 +1726,11 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - skip_sample: Optional[torch.Tensor] = None, + temb: torch.Tensor | None = None, + skip_sample: torch.Tensor | None = None, *args, **kwargs, - ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...], torch.Tensor]: + ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...], torch.Tensor]: if len(args) > 0 or kwargs.get("scale", None) is not None: deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`." deprecate("scale", "1.0.0", deprecation_message) @@ -1818,8 +1818,8 @@ def __init__( self.gradient_checkpointing = False def forward( - self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None, *args, **kwargs - ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]: + self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None, *args, **kwargs + ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]: if len(args) > 0 or kwargs.get("scale", None) is not None: deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`." deprecate("scale", "1.0.0", deprecation_message) @@ -1862,7 +1862,7 @@ def __init__( add_downsample: bool = True, skip_time_act: bool = False, only_cross_attention: bool = False, - cross_attention_norm: Optional[str] = None, + cross_attention_norm: str | None = None, ): super().__init__() @@ -1941,12 +1941,12 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]: + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + encoder_attention_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]: cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} if cross_attention_kwargs.get("scale", None) is not None: logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.") @@ -2041,8 +2041,8 @@ def __init__( self.gradient_checkpointing = False def forward( - self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None, *args, **kwargs - ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]: + self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None, *args, **kwargs + ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]: if len(args) > 0 or kwargs.get("scale", None) is not None: deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`." deprecate("scale", "1.0.0", deprecation_message) @@ -2132,12 +2132,12 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]: + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + encoder_attention_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]: cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} if cross_attention_kwargs.get("scale", None) is not None: logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.") @@ -2278,9 +2278,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - upsample_size: Optional[int] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, + upsample_size: int | None = None, *args, **kwargs, ) -> torch.Tensor: @@ -2318,10 +2318,10 @@ def __init__( out_channels: int, prev_output_channel: int, temb_channels: int, - resolution_idx: Optional[int] = None, + resolution_idx: int | None = None, dropout: float = 0.0, num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, + transformer_layers_per_block: int | tuple[int] = 1, resnet_eps: float = 1e-6, resnet_time_scale_shift: str = "default", resnet_act_fn: str = "swish", @@ -2405,13 +2405,13 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - upsample_size: Optional[int] = None, - attention_mask: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + upsample_size: int | None = None, + attention_mask: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: if cross_attention_kwargs is not None: if cross_attention_kwargs.get("scale", None) is not None: @@ -2478,7 +2478,7 @@ def __init__( prev_output_channel: int, out_channels: int, temb_channels: int, - resolution_idx: Optional[int] = None, + resolution_idx: int | None = None, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, @@ -2524,9 +2524,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - upsample_size: Optional[int] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, + upsample_size: int | None = None, *args, **kwargs, ) -> torch.Tensor: @@ -2577,7 +2577,7 @@ def __init__( self, in_channels: int, out_channels: int, - resolution_idx: Optional[int] = None, + resolution_idx: int | None = None, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, @@ -2587,7 +2587,7 @@ def __init__( resnet_pre_norm: bool = True, output_scale_factor: float = 1.0, add_upsample: bool = True, - temb_channels: Optional[int] = None, + temb_channels: int | None = None, ): super().__init__() resnets = [] @@ -2634,7 +2634,7 @@ def __init__( self.resolution_idx = resolution_idx - def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor: for resnet in self.resnets: hidden_states = resnet(hidden_states, temb=temb) @@ -2650,7 +2650,7 @@ def __init__( self, in_channels: int, out_channels: int, - resolution_idx: Optional[int] = None, + resolution_idx: int | None = None, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, @@ -2661,7 +2661,7 @@ def __init__( attention_head_dim: int = 1, output_scale_factor: float = 1.0, add_upsample: bool = True, - temb_channels: Optional[int] = None, + temb_channels: int | None = None, ): super().__init__() resnets = [] @@ -2732,7 +2732,7 @@ def __init__( self.resolution_idx = resolution_idx - def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor: for resnet, attn in zip(self.resnets, self.attentions): hidden_states = resnet(hidden_states, temb=temb) hidden_states = attn(hidden_states, temb=temb) @@ -2751,7 +2751,7 @@ def __init__( prev_output_channel: int, out_channels: int, temb_channels: int, - resolution_idx: Optional[int] = None, + resolution_idx: int | None = None, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, @@ -2841,12 +2841,12 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, skip_sample=None, *args, **kwargs, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: if len(args) > 0 or kwargs.get("scale", None) is not None: deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`." deprecate("scale", "1.0.0", deprecation_message) @@ -2885,7 +2885,7 @@ def __init__( prev_output_channel: int, out_channels: int, temb_channels: int, - resolution_idx: Optional[int] = None, + resolution_idx: int | None = None, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, @@ -2953,12 +2953,12 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, skip_sample=None, *args, **kwargs, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: if len(args) > 0 or kwargs.get("scale", None) is not None: deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`." deprecate("scale", "1.0.0", deprecation_message) @@ -2995,7 +2995,7 @@ def __init__( prev_output_channel: int, out_channels: int, temb_channels: int, - resolution_idx: Optional[int] = None, + resolution_idx: int | None = None, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, @@ -3060,9 +3060,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - upsample_size: Optional[int] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, + upsample_size: int | None = None, *args, **kwargs, ) -> torch.Tensor: @@ -3095,7 +3095,7 @@ def __init__( out_channels: int, prev_output_channel: int, temb_channels: int, - resolution_idx: Optional[int] = None, + resolution_idx: int | None = None, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, @@ -3109,7 +3109,7 @@ def __init__( add_upsample: bool = True, skip_time_act: bool = False, only_cross_attention: bool = False, - cross_attention_norm: Optional[str] = None, + cross_attention_norm: str | None = None, ): super().__init__() resnets = [] @@ -3190,13 +3190,13 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - upsample_size: Optional[int] = None, - attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + upsample_size: int | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + encoder_attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} if cross_attention_kwargs.get("scale", None) is not None: @@ -3256,7 +3256,7 @@ def __init__( num_layers: int = 5, resnet_eps: float = 1e-5, resnet_act_fn: str = "gelu", - resnet_group_size: Optional[int] = 32, + resnet_group_size: int | None = 32, add_upsample: bool = True, ): super().__init__() @@ -3298,9 +3298,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - upsample_size: Optional[int] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, + upsample_size: int | None = None, *args, **kwargs, ) -> torch.Tensor: @@ -3414,13 +3414,13 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - upsample_size: Optional[int] = None, - attention_mask: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + upsample_size: int | None = None, + attention_mask: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: res_hidden_states_tuple = res_hidden_states_tuple[-1] if res_hidden_states_tuple is not None: @@ -3490,12 +3490,12 @@ def __init__( num_attention_heads: int, attention_head_dim: int, dropout: float = 0.0, - cross_attention_dim: Optional[int] = None, + cross_attention_dim: int | None = None, attention_bias: bool = False, upcast_attention: bool = False, temb_channels: int = 768, # for ada_group_norm add_self_attention: bool = False, - cross_attention_norm: Optional[str] = None, + cross_attention_norm: str | None = None, group_size: int = 32, ): super().__init__() @@ -3536,13 +3536,13 @@ def _to_4d(self, hidden_states: torch.Tensor, height: int, weight: int) -> torch def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_hidden_states: torch.Tensor | None = None, # TODO: mark emb as non-optional (self.norm2 requires it). # requires assessing impact of change to positional param interface. - emb: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, + emb: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + encoder_attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} if cross_attention_kwargs.get("scale", None) is not None: diff --git a/src/diffusers/models/unets/unet_2d_condition.py b/src/diffusers/models/unets/unet_2d_condition.py index e669aa51a54e..deae25899475 100644 --- a/src/diffusers/models/unets/unet_2d_condition.py +++ b/src/diffusers/models/unets/unet_2d_condition.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -20,7 +20,12 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin, UNet2DConditionLoadersMixin from ...loaders.single_file_model import FromOriginalModelMixin -from ...utils import USE_PEFT_BACKEND, BaseOutput, deprecate, logging, scale_lora_layers, unscale_lora_layers +from ...utils import ( + BaseOutput, + apply_lora_scale, + deprecate, + logging, +) from ..activations import get_activation from ..attention import AttentionMixin from ..attention_processor import ( @@ -78,7 +83,7 @@ class UNet2DConditionModel( for all models (such as downloading or saving). Parameters: - sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`): + sample_size (`int` or `tuple[int, int]`, *optional*, defaults to `None`): Height and width of input/output sample. in_channels (`int`, *optional*, defaults to 4): Number of channels in the input sample. out_channels (`int`, *optional*, defaults to 4): Number of channels in the output. @@ -86,17 +91,17 @@ class UNet2DConditionModel( flip_sin_to_cos (`bool`, *optional*, defaults to `True`): Whether to flip the sin to cos in the time embedding. freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding. - down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`): + down_block_types (`tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`): The tuple of downsample blocks to use. mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`): Block type for middle of UNet, it can be one of `UNetMidBlock2DCrossAttn`, `UNetMidBlock2D`, or `UNetMidBlock2DSimpleCrossAttn`. If `None`, the mid block layer is skipped. - up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`): + up_block_types (`tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`): The tuple of upsample blocks to use. - only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`): + only_cross_attention(`bool` or `tuple[bool]`, *optional*, default to `False`): Whether to include self-attention in the basic transformer blocks, see [`~models.attention.BasicTransformerBlock`]. - block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): + block_out_channels (`tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): The tuple of output channels for each block. layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block. downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution. @@ -106,15 +111,15 @@ class UNet2DConditionModel( norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization. If `None`, normalization and activation layers is skipped in post-processing. norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization. - cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280): + cross_attention_dim (`int` or `tuple[int]`, *optional*, defaults to 1280): The dimension of the cross attention features. - transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple]` , *optional*, defaults to 1): + transformer_layers_per_block (`int`, `tuple[int]`, or `tuple[tuple]` , *optional*, defaults to 1): The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for [`~models.unets.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unets.unet_2d_blocks.CrossAttnUpBlock2D`], [`~models.unets.unet_2d_blocks.UNetMidBlock2DCrossAttn`]. - reverse_transformer_layers_per_block : (`Tuple[Tuple]`, *optional*, defaults to None): + reverse_transformer_layers_per_block : (`tuple[tuple]`, *optional*, defaults to None): The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`], in the upsampling - blocks of the U-Net. Only relevant if `transformer_layers_per_block` is of type `Tuple[Tuple]` and for + blocks of the U-Net. Only relevant if `transformer_layers_per_block` is of type `tuple[tuple]` and for [`~models.unets.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unets.unet_2d_blocks.CrossAttnUpBlock2D`], [`~models.unets.unet_2d_blocks.UNetMidBlock2DCrossAttn`]. encoder_hid_dim (`int`, *optional*, defaults to None): @@ -171,63 +176,63 @@ class conditioning with `class_embed_type` equal to `None`. @register_to_config def __init__( self, - sample_size: Optional[Union[int, Tuple[int, int]]] = None, + sample_size: int | tuple[int, int] | None = None, in_channels: int = 4, out_channels: int = 4, center_input_sample: bool = False, flip_sin_to_cos: bool = True, freq_shift: int = 0, - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str, ...] = ( "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D", ), - mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn", - up_block_types: Tuple[str, ...] = ( + mid_block_type: str | None = "UNetMidBlock2DCrossAttn", + up_block_types: tuple[str, ...] = ( "UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", ), - only_cross_attention: Union[bool, Tuple[bool]] = False, - block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280), - layers_per_block: Union[int, Tuple[int]] = 2, + only_cross_attention: bool | tuple[bool] = False, + block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280), + layers_per_block: int | tuple[int] = 2, downsample_padding: int = 1, mid_block_scale_factor: float = 1, dropout: float = 0.0, act_fn: str = "silu", - norm_num_groups: Optional[int] = 32, + norm_num_groups: int | None = 32, norm_eps: float = 1e-5, - cross_attention_dim: Union[int, Tuple[int]] = 1280, - transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1, - reverse_transformer_layers_per_block: Optional[Tuple[Tuple[int]]] = None, - encoder_hid_dim: Optional[int] = None, - encoder_hid_dim_type: Optional[str] = None, - attention_head_dim: Union[int, Tuple[int]] = 8, - num_attention_heads: Optional[Union[int, Tuple[int]]] = None, + cross_attention_dim: int | tuple[int] = 1280, + transformer_layers_per_block: int | tuple[int] | tuple[tuple] = 1, + reverse_transformer_layers_per_block: tuple[tuple[int]] | None = None, + encoder_hid_dim: int | None = None, + encoder_hid_dim_type: str | None = None, + attention_head_dim: int | tuple[int] = 8, + num_attention_heads: int | tuple[int] | None = None, dual_cross_attention: bool = False, use_linear_projection: bool = False, - class_embed_type: Optional[str] = None, - addition_embed_type: Optional[str] = None, - addition_time_embed_dim: Optional[int] = None, - num_class_embeds: Optional[int] = None, + class_embed_type: str | None = None, + addition_embed_type: str | None = None, + addition_time_embed_dim: int | None = None, + num_class_embeds: int | None = None, upcast_attention: bool = False, resnet_time_scale_shift: str = "default", resnet_skip_time_act: bool = False, resnet_out_scale_factor: float = 1.0, time_embedding_type: str = "positional", - time_embedding_dim: Optional[int] = None, - time_embedding_act_fn: Optional[str] = None, - timestep_post_act: Optional[str] = None, - time_cond_proj_dim: Optional[int] = None, + time_embedding_dim: int | None = None, + time_embedding_act_fn: str | None = None, + timestep_post_act: str | None = None, + time_cond_proj_dim: int | None = None, conv_in_kernel: int = 3, conv_out_kernel: int = 3, - projection_class_embeddings_input_dim: Optional[int] = None, + projection_class_embeddings_input_dim: int | None = None, attention_type: str = "default", class_embeddings_concat: bool = False, - mid_block_only_cross_attention: Optional[bool] = None, - cross_attention_norm: Optional[str] = None, + mid_block_only_cross_attention: bool | None = None, + cross_attention_norm: str | None = None, addition_embed_type_num_heads: int = 64, ): super().__init__() @@ -491,16 +496,16 @@ def __init__( def _check_config( self, - down_block_types: Tuple[str, ...], - up_block_types: Tuple[str, ...], - only_cross_attention: Union[bool, Tuple[bool]], - block_out_channels: Tuple[int, ...], - layers_per_block: Union[int, Tuple[int]], - cross_attention_dim: Union[int, Tuple[int]], - transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple[int]]], + down_block_types: tuple[str, ...], + up_block_types: tuple[str, ...], + only_cross_attention: bool | tuple[bool], + block_out_channels: tuple[int, ...], + layers_per_block: int | tuple[int], + cross_attention_dim: int | tuple[int], + transformer_layers_per_block: int | tuple[int, tuple[tuple[int]]], reverse_transformer_layers_per_block: bool, attention_head_dim: int, - num_attention_heads: Optional[Union[int, Tuple[int]]], + num_attention_heads: int | tuple[int] | None, ): if len(down_block_types) != len(up_block_types): raise ValueError( @@ -548,7 +553,7 @@ def _set_time_proj( flip_sin_to_cos: bool, freq_shift: float, time_embedding_dim: int, - ) -> Tuple[int, int]: + ) -> tuple[int, int]: if time_embedding_type == "fourier": time_embed_dim = time_embedding_dim or block_out_channels[0] * 2 if time_embed_dim % 2 != 0: @@ -571,9 +576,9 @@ def _set_time_proj( def _set_encoder_hid_proj( self, - encoder_hid_dim_type: Optional[str], - cross_attention_dim: Union[int, Tuple[int]], - encoder_hid_dim: Optional[int], + encoder_hid_dim_type: str | None, + cross_attention_dim: int | tuple[int], + encoder_hid_dim: int | None, ): if encoder_hid_dim_type is None and encoder_hid_dim is not None: encoder_hid_dim_type = "text_proj" @@ -611,10 +616,10 @@ def _set_encoder_hid_proj( def _set_class_embedding( self, - class_embed_type: Optional[str], + class_embed_type: str | None, act_fn: str, - num_class_embeds: Optional[int], - projection_class_embeddings_input_dim: Optional[int], + num_class_embeds: int | None, + projection_class_embeddings_input_dim: int | None, time_embed_dim: int, timestep_input_dim: int, ): @@ -650,12 +655,12 @@ def _set_add_embedding( self, addition_embed_type: str, addition_embed_type_num_heads: int, - addition_time_embed_dim: Optional[int], + addition_time_embed_dim: int | None, flip_sin_to_cos: bool, freq_shift: float, - cross_attention_dim: Optional[int], - encoder_hid_dim: Optional[int], - projection_class_embeddings_input_dim: Optional[int], + cross_attention_dim: int | None, + encoder_hid_dim: int | None, + projection_class_embeddings_input_dim: int | None, time_embed_dim: int, ): if addition_embed_type == "text": @@ -716,7 +721,7 @@ def set_default_attn_processor(self): self.set_attn_processor(processor) - def set_attention_slice(self, slice_size: Union[str, int, List[int]] = "auto"): + def set_attention_slice(self, slice_size: str | int | list[int] = "auto"): r""" Enable sliced attention computation. @@ -770,7 +775,7 @@ def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): # Recursively walk through all the children. # Any children which exposes the set_attention_slice method # gets the message - def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]): + def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: list[int]): if hasattr(module, "set_attention_slice"): module.set_attention_slice(slice_size.pop()) @@ -843,9 +848,7 @@ def unfuse_qkv_projections(self): if self.original_attn_processors is not None: self.set_attn_processor(self.original_attn_processors) - def get_time_embed( - self, sample: torch.Tensor, timestep: Union[torch.Tensor, float, int] - ) -> Optional[torch.Tensor]: + def get_time_embed(self, sample: torch.Tensor, timestep: torch.Tensor | float | int) -> torch.Tensor | None: timesteps = timestep if not torch.is_tensor(timesteps): # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can @@ -870,7 +873,7 @@ def get_time_embed( t_emb = t_emb.to(dtype=sample.dtype) return t_emb - def get_class_embed(self, sample: torch.Tensor, class_labels: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + def get_class_embed(self, sample: torch.Tensor, class_labels: torch.Tensor | None) -> torch.Tensor | None: class_emb = None if self.class_embedding is not None: if class_labels is None: @@ -887,8 +890,8 @@ def get_class_embed(self, sample: torch.Tensor, class_labels: Optional[torch.Ten return class_emb def get_aug_embed( - self, emb: torch.Tensor, encoder_hidden_states: torch.Tensor, added_cond_kwargs: Dict[str, Any] - ) -> Optional[torch.Tensor]: + self, emb: torch.Tensor, encoder_hidden_states: torch.Tensor, added_cond_kwargs: dict[str, Any] + ) -> torch.Tensor | None: aug_emb = None if self.config.addition_embed_type == "text": aug_emb = self.add_embedding(encoder_hidden_states) @@ -939,7 +942,7 @@ def get_aug_embed( return aug_emb def process_encoder_hidden_states( - self, encoder_hidden_states: torch.Tensor, added_cond_kwargs: Dict[str, Any] + self, encoder_hidden_states: torch.Tensor, added_cond_kwargs: dict[str, Any] ) -> torch.Tensor: if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_proj": encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states) @@ -974,22 +977,23 @@ def process_encoder_hidden_states( encoder_hidden_states = (encoder_hidden_states, image_embeds) return encoder_hidden_states + @apply_lora_scale("cross_attention_kwargs") def forward( self, sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], + timestep: torch.Tensor | float | int, encoder_hidden_states: torch.Tensor, - class_labels: Optional[torch.Tensor] = None, - timestep_cond: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, - down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None, - mid_block_additional_residual: Optional[torch.Tensor] = None, - down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, + class_labels: torch.Tensor | None = None, + timestep_cond: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + added_cond_kwargs: dict[str, torch.Tensor] | None = None, + down_block_additional_residuals: tuple[torch.Tensor] | None = None, + mid_block_additional_residual: torch.Tensor | None = None, + down_intrablock_additional_residuals: tuple[torch.Tensor] | None = None, + encoder_attention_mask: torch.Tensor | None = None, return_dict: bool = True, - ) -> Union[UNet2DConditionOutput, Tuple]: + ) -> UNet2DConditionOutput | tuple: r""" The [`UNet2DConditionModel`] forward method. @@ -1112,18 +1116,6 @@ def forward( cross_attention_kwargs["gligen"] = {"objs": self.position_net(**gligen_args)} # 3. down - # we're popping the `scale` instead of getting it because otherwise `scale` will be propagated - # to the internal blocks and will raise deprecation warnings. this will be confusing for our users. - if cross_attention_kwargs is not None: - cross_attention_kwargs = cross_attention_kwargs.copy() - lora_scale = cross_attention_kwargs.pop("scale", 1.0) - else: - lora_scale = 1.0 - - if USE_PEFT_BACKEND: - # weight the lora layers by setting `lora_scale` for each PEFT layer - scale_lora_layers(self, lora_scale) - is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None # using new arg down_intrablock_additional_residuals for T2I-Adapters, to distinguish from controlnets is_adapter = down_intrablock_additional_residuals is not None @@ -1239,10 +1231,6 @@ def forward( sample = self.conv_act(sample) sample = self.conv_out(sample) - if USE_PEFT_BACKEND: - # remove `lora_scale` from each PEFT layer - unscale_lora_layers(self, lora_scale) - if not return_dict: return (sample,) diff --git a/src/diffusers/models/unets/unet_2d_condition_flax.py b/src/diffusers/models/unets/unet_2d_condition_flax.py index 8d9a309afbcc..a361026fc0ea 100644 --- a/src/diffusers/models/unets/unet_2d_condition_flax.py +++ b/src/diffusers/models/unets/unet_2d_condition_flax.py @@ -11,7 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Dict, Optional, Tuple, Union import flax import flax.linen as nn @@ -74,20 +73,20 @@ class FlaxUNet2DConditionModel(nn.Module, FlaxModelMixin, ConfigMixin): The number of channels in the input sample. out_channels (`int`, *optional*, defaults to 4): The number of channels in the output. - down_block_types (`Tuple[str]`, *optional*, defaults to `("FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxDownBlock2D")`): + down_block_types (`tuple[str]`, *optional*, defaults to `("FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxDownBlock2D")`): The tuple of downsample blocks to use. - up_block_types (`Tuple[str]`, *optional*, defaults to `("FlaxUpBlock2D", "FlaxCrossAttnUpBlock2D", "FlaxCrossAttnUpBlock2D", "FlaxCrossAttnUpBlock2D")`): + up_block_types (`tuple[str]`, *optional*, defaults to `("FlaxUpBlock2D", "FlaxCrossAttnUpBlock2D", "FlaxCrossAttnUpBlock2D", "FlaxCrossAttnUpBlock2D")`): The tuple of upsample blocks to use. mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`): Block type for middle of UNet, it can be one of `UNetMidBlock2DCrossAttn`. If `None`, the mid block layer is skipped. - block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): + block_out_channels (`tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): The tuple of output channels for each block. layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block. - attention_head_dim (`int` or `Tuple[int]`, *optional*, defaults to 8): + attention_head_dim (`int` or `tuple[int]`, *optional*, defaults to 8): The dimension of the attention heads. - num_attention_heads (`int` or `Tuple[int]`, *optional*): + num_attention_heads (`int` or `tuple[int]`, *optional*): The number of attention heads. cross_attention_dim (`int`, *optional*, defaults to 768): The dimension of the cross attention features. @@ -106,19 +105,19 @@ class FlaxUNet2DConditionModel(nn.Module, FlaxModelMixin, ConfigMixin): sample_size: int = 32 in_channels: int = 4 out_channels: int = 4 - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str, ...] = ( "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D", ) - up_block_types: Tuple[str, ...] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D") - mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn" - only_cross_attention: Union[bool, Tuple[bool]] = False - block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280) + up_block_types: tuple[str, ...] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D") + mid_block_type: str | None = "UNetMidBlock2DCrossAttn" + only_cross_attention: bool | tuple[bool] = False + block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280) layers_per_block: int = 2 - attention_head_dim: Union[int, Tuple[int, ...]] = 8 - num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None + attention_head_dim: int | tuple[int, ...] = 8 + num_attention_heads: int | tuple[int, ...] | None = None cross_attention_dim: int = 1280 dropout: float = 0.0 use_linear_projection: bool = False @@ -127,11 +126,11 @@ class FlaxUNet2DConditionModel(nn.Module, FlaxModelMixin, ConfigMixin): freq_shift: int = 0 use_memory_efficient_attention: bool = False split_head_dim: bool = False - transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1 - addition_embed_type: Optional[str] = None - addition_time_embed_dim: Optional[int] = None + transformer_layers_per_block: int | tuple[int, ...] = 1 + addition_embed_type: str | None = None + addition_time_embed_dim: int | None = None addition_embed_type_num_heads: int = 64 - projection_class_embeddings_input_dim: Optional[int] = None + projection_class_embeddings_input_dim: int | None = None def init_weights(self, rng: jax.Array) -> FrozenDict: # init input tensors @@ -338,14 +337,14 @@ def setup(self) -> None: def __call__( self, sample: jnp.ndarray, - timesteps: Union[jnp.ndarray, float, int], + timesteps: jnp.ndarray | float | int, encoder_hidden_states: jnp.ndarray, - added_cond_kwargs: Optional[Union[Dict, FrozenDict]] = None, - down_block_additional_residuals: Optional[Tuple[jnp.ndarray, ...]] = None, - mid_block_additional_residual: Optional[jnp.ndarray] = None, + added_cond_kwargs: dict | FrozenDict | None = None, + down_block_additional_residuals: tuple[jnp.ndarray, ...] | None = None, + mid_block_additional_residual: jnp.ndarray | None = None, return_dict: bool = True, train: bool = False, - ) -> Union[FlaxUNet2DConditionOutput, Tuple[jnp.ndarray]]: + ) -> FlaxUNet2DConditionOutput | tuple[jnp.ndarray]: r""" Args: sample (`jnp.ndarray`): (batch, channel, height, width) noisy inputs tensor diff --git a/src/diffusers/models/unets/unet_3d_blocks.py b/src/diffusers/models/unets/unet_3d_blocks.py index 53c0f4bae38b..b5f5aae342b4 100644 --- a/src/diffusers/models/unets/unet_3d_blocks.py +++ b/src/diffusers/models/unets/unet_3d_blocks.py @@ -12,7 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, Optional, Tuple, Union +from __future__ import annotations + +from typing import Any import torch from torch import nn @@ -89,9 +91,9 @@ def get_down_block( resnet_eps: float, resnet_act_fn: str, num_attention_heads: int, - resnet_groups: Optional[int] = None, - cross_attention_dim: Optional[int] = None, - downsample_padding: Optional[int] = None, + resnet_groups: int | None = None, + cross_attention_dim: int | None = None, + downsample_padding: int | None = None, dual_cross_attention: bool = False, use_linear_projection: bool = True, only_cross_attention: bool = False, @@ -99,15 +101,10 @@ def get_down_block( resnet_time_scale_shift: str = "default", temporal_num_attention_heads: int = 8, temporal_max_seq_length: int = 32, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, - temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1, + transformer_layers_per_block: int | tuple[int] = 1, + temporal_transformer_layers_per_block: int | tuple[int] = 1, dropout: float = 0.0, -) -> Union[ - "DownBlock3D", - "CrossAttnDownBlock3D", - "DownBlockSpatioTemporal", - "CrossAttnDownBlockSpatioTemporal", -]: +) -> "DownBlock3D" | "CrossAttnDownBlock3D" | "DownBlockSpatioTemporal" | "CrossAttnDownBlockSpatioTemporal": if down_block_type == "DownBlock3D": return DownBlock3D( num_layers=num_layers, @@ -182,26 +179,21 @@ def get_up_block( resnet_eps: float, resnet_act_fn: str, num_attention_heads: int, - resolution_idx: Optional[int] = None, - resnet_groups: Optional[int] = None, - cross_attention_dim: Optional[int] = None, + resolution_idx: int | None = None, + resnet_groups: int | None = None, + cross_attention_dim: int | None = None, dual_cross_attention: bool = False, use_linear_projection: bool = True, only_cross_attention: bool = False, upcast_attention: bool = False, resnet_time_scale_shift: str = "default", temporal_num_attention_heads: int = 8, - temporal_cross_attention_dim: Optional[int] = None, + temporal_cross_attention_dim: int | None = None, temporal_max_seq_length: int = 32, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, - temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1, + transformer_layers_per_block: int | tuple[int] = 1, + temporal_transformer_layers_per_block: int | tuple[int] = 1, dropout: float = 0.0, -) -> Union[ - "UpBlock3D", - "CrossAttnUpBlock3D", - "UpBlockSpatioTemporal", - "CrossAttnUpBlockSpatioTemporal", -]: +) -> "UpBlock3D" | "CrossAttnUpBlock3D" | "UpBlockSpatioTemporal" | "CrossAttnUpBlockSpatioTemporal": if up_block_type == "UpBlock3D": return UpBlock3D( num_layers=num_layers, @@ -376,11 +368,11 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, num_frames: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, ) -> torch.Tensor: hidden_states = self.resnets[0](hidden_states, temb) hidden_states = self.temp_convs[0](hidden_states, num_frames=num_frames) @@ -509,12 +501,12 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, num_frames: int = 1, - cross_attention_kwargs: Dict[str, Any] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: + cross_attention_kwargs: dict[str, Any] = None, + ) -> torch.Tensor | tuple[torch.Tensor, ...]: # TODO(Patrick, William) - attention mask is not used output_states = () @@ -616,9 +608,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, + temb: torch.Tensor | None = None, num_frames: int = 1, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: + ) -> torch.Tensor | tuple[torch.Tensor, ...]: output_states = () for resnet, temp_conv in zip(self.resnets, self.temp_convs): @@ -658,7 +650,7 @@ def __init__( use_linear_projection: bool = False, only_cross_attention: bool = False, upcast_attention: bool = False, - resolution_idx: Optional[int] = None, + resolution_idx: int | None = None, ): super().__init__() resnets = [] @@ -734,13 +726,13 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - upsample_size: Optional[int] = None, - attention_mask: Optional[torch.Tensor] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + upsample_size: int | None = None, + attention_mask: torch.Tensor | None = None, num_frames: int = 1, - cross_attention_kwargs: Dict[str, Any] = None, + cross_attention_kwargs: dict[str, Any] = None, ) -> torch.Tensor: is_freeu_enabled = ( getattr(self, "s1", None) @@ -809,7 +801,7 @@ def __init__( resnet_pre_norm: bool = True, output_scale_factor: float = 1.0, add_upsample: bool = True, - resolution_idx: Optional[int] = None, + resolution_idx: int | None = None, ): super().__init__() resnets = [] @@ -856,9 +848,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - upsample_size: Optional[int] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, + upsample_size: int | None = None, num_frames: int = 1, ) -> torch.Tensor: is_freeu_enabled = ( @@ -1015,7 +1007,7 @@ def __init__( in_channels: int, temb_channels: int, num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, + transformer_layers_per_block: int | tuple[int] = 1, num_attention_heads: int = 1, cross_attention_dim: int = 1280, ): @@ -1067,9 +1059,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - image_only_indicator: Optional[torch.Tensor] = None, + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + image_only_indicator: torch.Tensor | None = None, ) -> torch.Tensor: hidden_states = self.resnets[0]( hidden_states, @@ -1142,9 +1134,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - image_only_indicator: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]: + temb: torch.Tensor | None = None, + image_only_indicator: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]: output_states = () for resnet in self.resnets: if torch.is_grad_enabled() and self.gradient_checkpointing: @@ -1170,7 +1162,7 @@ def __init__( out_channels: int, temb_channels: int, num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, + transformer_layers_per_block: int | tuple[int] = 1, num_attention_heads: int = 1, cross_attention_dim: int = 1280, add_downsample: bool = True, @@ -1227,10 +1219,10 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - image_only_indicator: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]: + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + image_only_indicator: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]: output_states = () blocks = list(zip(self.resnets, self.attentions)) @@ -1271,7 +1263,7 @@ def __init__( prev_output_channel: int, out_channels: int, temb_channels: int, - resolution_idx: Optional[int] = None, + resolution_idx: int | None = None, num_layers: int = 1, resnet_eps: float = 1e-6, add_upsample: bool = True, @@ -1305,10 +1297,10 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - image_only_indicator: Optional[torch.Tensor] = None, - upsample_size: Optional[int] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, + image_only_indicator: torch.Tensor | None = None, + upsample_size: int | None = None, ) -> torch.Tensor: for resnet in self.resnets: # pop res hidden states @@ -1336,9 +1328,9 @@ def __init__( out_channels: int, prev_output_channel: int, temb_channels: int, - resolution_idx: Optional[int] = None, + resolution_idx: int | None = None, num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, + transformer_layers_per_block: int | tuple[int] = 1, resnet_eps: float = 1e-6, num_attention_heads: int = 1, cross_attention_dim: int = 1280, @@ -1390,11 +1382,11 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - image_only_indicator: Optional[torch.Tensor] = None, - upsample_size: Optional[int] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + image_only_indicator: torch.Tensor | None = None, + upsample_size: int | None = None, ) -> torch.Tensor: for resnet, attn in zip(self.resnets, self.attentions): # pop res hidden states diff --git a/src/diffusers/models/unets/unet_3d_condition.py b/src/diffusers/models/unets/unet_3d_condition.py index 26dc50f84acd..5006e48feb46 100644 --- a/src/diffusers/models/unets/unet_3d_condition.py +++ b/src/diffusers/models/unets/unet_3d_condition.py @@ -14,7 +14,7 @@ # limitations under the License. from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -67,15 +67,15 @@ class UNet3DConditionModel(ModelMixin, AttentionMixin, ConfigMixin, UNet2DCondit for all models (such as downloading or saving). Parameters: - sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`): + sample_size (`int` or `tuple[int, int]`, *optional*, defaults to `None`): Height and width of input/output sample. in_channels (`int`, *optional*, defaults to 4): The number of channels in the input sample. out_channels (`int`, *optional*, defaults to 4): The number of channels in the output. - down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "DownBlock3D")`): + down_block_types (`tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "DownBlock3D")`): The tuple of downsample blocks to use. - up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D")`): + up_block_types (`tuple[str]`, *optional*, defaults to `("UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D")`): The tuple of upsample blocks to use. - block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): + block_out_channels (`tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): The tuple of output channels for each block. layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block. downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution. @@ -97,32 +97,32 @@ class UNet3DConditionModel(ModelMixin, AttentionMixin, ConfigMixin, UNet2DCondit @register_to_config def __init__( self, - sample_size: Optional[int] = None, + sample_size: int | None = None, in_channels: int = 4, out_channels: int = 4, - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str, ...] = ( "CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "DownBlock3D", ), - up_block_types: Tuple[str, ...] = ( + up_block_types: tuple[str, ...] = ( "UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", ), - block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280), + block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280), layers_per_block: int = 2, downsample_padding: int = 1, mid_block_scale_factor: float = 1, act_fn: str = "silu", - norm_num_groups: Optional[int] = 32, + norm_num_groups: int | None = 32, norm_eps: float = 1e-5, cross_attention_dim: int = 1024, - attention_head_dim: Union[int, Tuple[int]] = 64, - num_attention_heads: Optional[Union[int, Tuple[int]]] = None, - time_cond_proj_dim: Optional[int] = None, + attention_head_dim: int | tuple[int] = 64, + num_attention_heads: int | tuple[int] | None = None, + time_cond_proj_dim: int | None = None, ): super().__init__() @@ -286,7 +286,7 @@ def __init__( ) # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attention_slice - def set_attention_slice(self, slice_size: Union[str, int, List[int]]) -> None: + def set_attention_slice(self, slice_size: str | int | list[int]) -> None: r""" Enable sliced attention computation. @@ -340,7 +340,7 @@ def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): # Recursively walk through all the children. # Any children which exposes the set_attention_slice method # gets the message - def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]): + def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: list[int]): if hasattr(module, "set_attention_slice"): module.set_attention_slice(slice_size.pop()) @@ -351,7 +351,7 @@ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[i for module in self.children(): fn_recursive_set_attention_slice(module, reversed_slice_size) - def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None: + def enable_forward_chunking(self, chunk_size: int | None = None, dim: int = 0) -> None: """ Sets the attention processor to use [feed forward chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers). @@ -476,16 +476,16 @@ def unfuse_qkv_projections(self): def forward( self, sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], + timestep: torch.Tensor | float | int, encoder_hidden_states: torch.Tensor, - class_labels: Optional[torch.Tensor] = None, - timestep_cond: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None, - mid_block_additional_residual: Optional[torch.Tensor] = None, + class_labels: torch.Tensor | None = None, + timestep_cond: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + down_block_additional_residuals: tuple[torch.Tensor] | None = None, + mid_block_additional_residual: torch.Tensor | None = None, return_dict: bool = True, - ) -> Union[UNet3DConditionOutput, Tuple[torch.Tensor]]: + ) -> UNet3DConditionOutput | tuple[torch.Tensor]: r""" The [`UNet3DConditionModel`] forward method. diff --git a/src/diffusers/models/unets/unet_i2vgen_xl.py b/src/diffusers/models/unets/unet_i2vgen_xl.py index 0ada264417dd..5c3cfe91d5bd 100644 --- a/src/diffusers/models/unets/unet_i2vgen_xl.py +++ b/src/diffusers/models/unets/unet_i2vgen_xl.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -51,7 +51,7 @@ def __init__( attention_head_dim: int, activation_fn: str = "geglu", upcast_attention: bool = False, - ff_inner_dim: Optional[int] = None, + ff_inner_dim: int | None = None, dropout: int = 0.0, ): super().__init__() @@ -101,15 +101,15 @@ class I2VGenXLUNet(ModelMixin, AttentionMixin, ConfigMixin, UNet2DConditionLoade for all models (such as downloading or saving). Parameters: - sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`): + sample_size (`int` or `tuple[int, int]`, *optional*, defaults to `None`): Height and width of input/output sample. in_channels (`int`, *optional*, defaults to 4): The number of channels in the input sample. out_channels (`int`, *optional*, defaults to 4): The number of channels in the output. - down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`): + down_block_types (`tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`): The tuple of downsample blocks to use. - up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`): + up_block_types (`tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`): The tuple of upsample blocks to use. - block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): + block_out_channels (`tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): The tuple of output channels for each block. layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block. norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization. @@ -124,27 +124,27 @@ class I2VGenXLUNet(ModelMixin, AttentionMixin, ConfigMixin, UNet2DConditionLoade @register_to_config def __init__( self, - sample_size: Optional[int] = None, + sample_size: int | None = None, in_channels: int = 4, out_channels: int = 4, - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str, ...] = ( "CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "DownBlock3D", ), - up_block_types: Tuple[str, ...] = ( + up_block_types: tuple[str, ...] = ( "UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", ), - block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280), + block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280), layers_per_block: int = 2, - norm_num_groups: Optional[int] = 32, + norm_num_groups: int | None = 32, cross_attention_dim: int = 1024, - attention_head_dim: Union[int, Tuple[int]] = 64, - num_attention_heads: Optional[Union[int, Tuple[int]]] = None, + attention_head_dim: int | tuple[int] = 64, + num_attention_heads: int | tuple[int] | None = None, ): super().__init__() @@ -313,7 +313,7 @@ def __init__( self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, kernel_size=3, padding=1) # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking - def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None: + def enable_forward_chunking(self, chunk_size: int | None = None, dim: int = 0) -> None: """ Sets the attention processor to use [feed forward chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers). @@ -439,15 +439,15 @@ def unfuse_qkv_projections(self): def forward( self, sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], + timestep: torch.Tensor | float | int, fps: torch.Tensor, image_latents: torch.Tensor, - image_embeddings: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - timestep_cond: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + image_embeddings: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + timestep_cond: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - ) -> Union[UNet3DConditionOutput, Tuple[torch.Tensor]]: + ) -> UNet3DConditionOutput | tuple[torch.Tensor]: r""" The [`I2VGenXLUNet`] forward method. diff --git a/src/diffusers/models/unets/unet_kandinsky3.py b/src/diffusers/models/unets/unet_kandinsky3.py index 13f4641a4c50..6fa68b42ee30 100644 --- a/src/diffusers/models/unets/unet_kandinsky3.py +++ b/src/diffusers/models/unets/unet_kandinsky3.py @@ -13,7 +13,6 @@ # limitations under the License. from dataclasses import dataclass -from typing import Tuple, Union import torch from torch import nn @@ -54,9 +53,9 @@ def __init__( time_embedding_dim: int = 1536, groups: int = 32, attention_head_dim: int = 64, - layers_per_block: Union[int, Tuple[int]] = 3, - block_out_channels: Tuple[int, ...] = (384, 768, 1536, 3072), - cross_attention_dim: Union[int, Tuple[int]] = 4096, + layers_per_block: int | tuple[int] = 3, + block_out_channels: tuple[int, ...] = (384, 768, 1536, 3072), + cross_attention_dim: int | tuple[int] = 4096, encoder_hid_dim: int = 4096, ): super().__init__() diff --git a/src/diffusers/models/unets/unet_motion_model.py b/src/diffusers/models/unets/unet_motion_model.py index 5a93541501d3..97452eff05aa 100644 --- a/src/diffusers/models/unets/unet_motion_model.py +++ b/src/diffusers/models/unets/unet_motion_model.py @@ -13,7 +13,7 @@ # limitations under the License. from dataclasses import dataclass -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -21,7 +21,7 @@ from ...configuration_utils import ConfigMixin, FrozenDict, register_to_config from ...loaders import FromOriginalModelMixin, PeftAdapterMixin, UNet2DConditionLoadersMixin -from ...utils import BaseOutput, deprecate, logging +from ...utils import BaseOutput, apply_lora_scale, deprecate, logging from ...utils.torch_utils import apply_freeu from ..attention import AttentionMixin, BasicTransformerBlock from ..attention_processor import ( @@ -93,19 +93,19 @@ def __init__( self, num_attention_heads: int = 16, attention_head_dim: int = 88, - in_channels: Optional[int] = None, - out_channels: Optional[int] = None, + in_channels: int | None = None, + out_channels: int | None = None, num_layers: int = 1, dropout: float = 0.0, norm_num_groups: int = 32, - cross_attention_dim: Optional[int] = None, + cross_attention_dim: int | None = None, attention_bias: bool = False, - sample_size: Optional[int] = None, + sample_size: int | None = None, activation_fn: str = "geglu", norm_elementwise_affine: bool = True, double_self_attention: bool = True, - positional_embeddings: Optional[str] = None, - num_positional_embeddings: Optional[int] = None, + positional_embeddings: str | None = None, + num_positional_embeddings: int | None = None, ): super().__init__() self.num_attention_heads = num_attention_heads @@ -142,11 +142,11 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.LongTensor] = None, - timestep: Optional[torch.LongTensor] = None, - class_labels: Optional[torch.LongTensor] = None, + encoder_hidden_states: torch.LongTensor | None = None, + timestep: torch.LongTensor | None = None, + class_labels: torch.LongTensor | None = None, num_frames: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, ) -> torch.Tensor: """ The [`AnimateDiffTransformer3D`] forward method. @@ -227,10 +227,10 @@ def __init__( output_scale_factor: float = 1.0, add_downsample: bool = True, downsample_padding: int = 1, - temporal_num_attention_heads: Union[int, Tuple[int]] = 1, - temporal_cross_attention_dim: Optional[int] = None, + temporal_num_attention_heads: int | tuple[int] = 1, + temporal_cross_attention_dim: int | None = None, temporal_max_seq_length: int = 32, - temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1, + temporal_transformer_layers_per_block: int | tuple[int] = 1, temporal_double_self_attention: bool = True, ): super().__init__() @@ -308,11 +308,11 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, + temb: torch.Tensor | None = None, num_frames: int = 1, *args, **kwargs, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: + ) -> torch.Tensor | tuple[torch.Tensor, ...]: if len(args) > 0 or kwargs.get("scale", None) is not None: deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`." deprecate("scale", "1.0.0", deprecation_message) @@ -347,7 +347,7 @@ def __init__( temb_channels: int, dropout: float = 0.0, num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, + transformer_layers_per_block: int | tuple[int] = 1, resnet_eps: float = 1e-6, resnet_time_scale_shift: str = "default", resnet_act_fn: str = "swish", @@ -363,10 +363,10 @@ def __init__( only_cross_attention: bool = False, upcast_attention: bool = False, attention_type: str = "default", - temporal_cross_attention_dim: Optional[int] = None, + temporal_cross_attention_dim: int | None = None, temporal_num_attention_heads: int = 8, temporal_max_seq_length: int = 32, - temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1, + temporal_transformer_layers_per_block: int | tuple[int] = 1, temporal_double_self_attention: bool = True, ): super().__init__() @@ -477,13 +477,13 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, num_frames: int = 1, - encoder_attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - additional_residuals: Optional[torch.Tensor] = None, + encoder_attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + additional_residuals: torch.Tensor | None = None, ): if cross_attention_kwargs is not None: if cross_attention_kwargs.get("scale", None) is not None: @@ -531,10 +531,10 @@ def __init__( out_channels: int, prev_output_channel: int, temb_channels: int, - resolution_idx: Optional[int] = None, + resolution_idx: int | None = None, dropout: float = 0.0, num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, + transformer_layers_per_block: int | tuple[int] = 1, resnet_eps: float = 1e-6, resnet_time_scale_shift: str = "default", resnet_act_fn: str = "swish", @@ -549,10 +549,10 @@ def __init__( only_cross_attention: bool = False, upcast_attention: bool = False, attention_type: str = "default", - temporal_cross_attention_dim: Optional[int] = None, + temporal_cross_attention_dim: int | None = None, temporal_num_attention_heads: int = 8, temporal_max_seq_length: int = 32, - temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1, + temporal_transformer_layers_per_block: int | tuple[int] = 1, ): super().__init__() resnets = [] @@ -653,13 +653,13 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - upsample_size: Optional[int] = None, - attention_mask: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + upsample_size: int | None = None, + attention_mask: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, num_frames: int = 1, ) -> torch.Tensor: if cross_attention_kwargs is not None: @@ -723,7 +723,7 @@ def __init__( prev_output_channel: int, out_channels: int, temb_channels: int, - resolution_idx: Optional[int] = None, + resolution_idx: int | None = None, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, @@ -733,10 +733,10 @@ def __init__( resnet_pre_norm: bool = True, output_scale_factor: float = 1.0, add_upsample: bool = True, - temporal_cross_attention_dim: Optional[int] = None, + temporal_cross_attention_dim: int | None = None, temporal_num_attention_heads: int = 8, temporal_max_seq_length: int = 32, - temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1, + temporal_transformer_layers_per_block: int | tuple[int] = 1, ): super().__init__() resnets = [] @@ -798,8 +798,8 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, upsample_size=None, num_frames: int = 1, *args, @@ -858,7 +858,7 @@ def __init__( temb_channels: int, dropout: float = 0.0, num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, + transformer_layers_per_block: int | tuple[int] = 1, resnet_eps: float = 1e-6, resnet_time_scale_shift: str = "default", resnet_act_fn: str = "swish", @@ -872,9 +872,9 @@ def __init__( upcast_attention: bool = False, attention_type: str = "default", temporal_num_attention_heads: int = 1, - temporal_cross_attention_dim: Optional[int] = None, + temporal_cross_attention_dim: int | None = None, temporal_max_seq_length: int = 32, - temporal_transformer_layers_per_block: Union[int, Tuple[int]] = 1, + temporal_transformer_layers_per_block: int | tuple[int] = 1, ): super().__init__() @@ -980,11 +980,11 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + encoder_attention_mask: torch.Tensor | None = None, num_frames: int = 1, ) -> torch.Tensor: if cross_attention_kwargs is not None: @@ -1021,10 +1021,10 @@ def __init__( self, in_channels: int, layers_per_block: int = 2, - transformer_layers_per_block: Union[int, Tuple[int]] = 8, - num_attention_heads: Union[int, Tuple[int]] = 8, + transformer_layers_per_block: int | tuple[int] = 8, + num_attention_heads: int | tuple[int] = 8, attention_bias: bool = False, - cross_attention_dim: Optional[int] = None, + cross_attention_dim: int | None = None, activation_fn: str = "geglu", norm_num_groups: int = 32, max_seq_length: int = 32, @@ -1061,31 +1061,31 @@ class MotionAdapter(ModelMixin, ConfigMixin, FromOriginalModelMixin): @register_to_config def __init__( self, - block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280), - motion_layers_per_block: Union[int, Tuple[int]] = 2, - motion_transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple[int]]] = 1, + block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280), + motion_layers_per_block: int | tuple[int] = 2, + motion_transformer_layers_per_block: int | tuple[int] | tuple[tuple[int]] = 1, motion_mid_block_layers_per_block: int = 1, - motion_transformer_layers_per_mid_block: Union[int, Tuple[int]] = 1, - motion_num_attention_heads: Union[int, Tuple[int]] = 8, + motion_transformer_layers_per_mid_block: int | tuple[int] = 1, + motion_num_attention_heads: int | tuple[int] = 8, motion_norm_num_groups: int = 32, motion_max_seq_length: int = 32, use_motion_mid_block: bool = True, - conv_in_channels: Optional[int] = None, + conv_in_channels: int | None = None, ): """Container to store AnimateDiff Motion Modules Args: - block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): + block_out_channels (`tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): The tuple of output channels for each UNet block. - motion_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 2): + motion_layers_per_block (`int` or `tuple[int]`, *optional*, defaults to 2): The number of motion layers per UNet block. - motion_transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple[int]]`, *optional*, defaults to 1): + motion_transformer_layers_per_block (`int`, `tuple[int]`, or `tuple[tuple[int]]`, *optional*, defaults to 1): The number of transformer layers to use in each motion layer in each block. motion_mid_block_layers_per_block (`int`, *optional*, defaults to 1): The number of motion layers in the middle UNet block. - motion_transformer_layers_per_mid_block (`int` or `Tuple[int]`, *optional*, defaults to 1): + motion_transformer_layers_per_mid_block (`int` or `tuple[int]`, *optional*, defaults to 1): The number of transformer layers to use in each motion layer in the middle block. - motion_num_attention_heads (`int` or `Tuple[int]`, *optional*, defaults to 8): + motion_num_attention_heads (`int` or `tuple[int]`, *optional*, defaults to 8): The number of heads to use in each attention layer of the motion module. motion_norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use in each group normalization layer of the motion module. @@ -1209,48 +1209,48 @@ class UNetMotionModel(ModelMixin, AttentionMixin, ConfigMixin, UNet2DConditionLo @register_to_config def __init__( self, - sample_size: Optional[int] = None, + sample_size: int | None = None, in_channels: int = 4, out_channels: int = 4, - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str, ...] = ( "CrossAttnDownBlockMotion", "CrossAttnDownBlockMotion", "CrossAttnDownBlockMotion", "DownBlockMotion", ), - up_block_types: Tuple[str, ...] = ( + up_block_types: tuple[str, ...] = ( "UpBlockMotion", "CrossAttnUpBlockMotion", "CrossAttnUpBlockMotion", "CrossAttnUpBlockMotion", ), - block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280), - layers_per_block: Union[int, Tuple[int]] = 2, + block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280), + layers_per_block: int | tuple[int] = 2, downsample_padding: int = 1, mid_block_scale_factor: float = 1, act_fn: str = "silu", norm_num_groups: int = 32, norm_eps: float = 1e-5, cross_attention_dim: int = 1280, - transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1, - reverse_transformer_layers_per_block: Optional[Union[int, Tuple[int], Tuple[Tuple]]] = None, - temporal_transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1, - reverse_temporal_transformer_layers_per_block: Optional[Union[int, Tuple[int], Tuple[Tuple]]] = None, - transformer_layers_per_mid_block: Optional[Union[int, Tuple[int]]] = None, - temporal_transformer_layers_per_mid_block: Optional[Union[int, Tuple[int]]] = 1, + transformer_layers_per_block: int | tuple[int] | tuple[tuple] = 1, + reverse_transformer_layers_per_block: int | tuple[int] | tuple[tuple] | None = None, + temporal_transformer_layers_per_block: int | tuple[int] | tuple[tuple] = 1, + reverse_temporal_transformer_layers_per_block: int | tuple[int] | tuple[tuple] | None = None, + transformer_layers_per_mid_block: int | tuple[int] | None = None, + temporal_transformer_layers_per_mid_block: int | tuple[int] | None = 1, use_linear_projection: bool = False, - num_attention_heads: Union[int, Tuple[int, ...]] = 8, + num_attention_heads: int | tuple[int, ...] = 8, motion_max_seq_length: int = 32, - motion_num_attention_heads: Union[int, Tuple[int, ...]] = 8, - reverse_motion_num_attention_heads: Optional[Union[int, Tuple[int, ...], Tuple[Tuple[int, ...], ...]]] = None, + motion_num_attention_heads: int | tuple[int, ...] = 8, + reverse_motion_num_attention_heads: int | tuple[int, ...] | tuple[tuple[int, ...], ...] | None = None, use_motion_mid_block: bool = True, mid_block_layers: int = 1, - encoder_hid_dim: Optional[int] = None, - encoder_hid_dim_type: Optional[str] = None, - addition_embed_type: Optional[str] = None, - addition_time_embed_dim: Optional[int] = None, - projection_class_embeddings_input_dim: Optional[int] = None, - time_cond_proj_dim: Optional[int] = None, + encoder_hid_dim: int | None = None, + encoder_hid_dim_type: str | None = None, + addition_embed_type: str | None = None, + addition_time_embed_dim: int | None = None, + projection_class_embeddings_input_dim: int | None = None, + time_cond_proj_dim: int | None = None, ): super().__init__() @@ -1534,7 +1534,7 @@ def __init__( def from_unet2d( cls, unet: UNet2DConditionModel, - motion_adapter: Optional[MotionAdapter] = None, + motion_adapter: MotionAdapter | None = None, load_weights: bool = True, ): has_motion_adapter = motion_adapter is not None @@ -1708,7 +1708,7 @@ def freeze_unet2d_params(self) -> None: for param in motion_modules.parameters(): param.requires_grad = True - def load_motion_modules(self, motion_adapter: Optional[MotionAdapter]) -> None: + def load_motion_modules(self, motion_adapter: MotionAdapter | None) -> None: for i, down_block in enumerate(motion_adapter.down_blocks): self.down_blocks[i].motion_modules.load_state_dict(down_block.motion_modules.state_dict()) for i, up_block in enumerate(motion_adapter.up_blocks): @@ -1723,7 +1723,7 @@ def save_motion_modules( save_directory: str, is_main_process: bool = True, safe_serialization: bool = True, - variant: Optional[str] = None, + variant: str | None = None, push_to_hub: bool = False, **kwargs, ) -> None: @@ -1753,7 +1753,7 @@ def save_motion_modules( **kwargs, ) - def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None: + def enable_forward_chunking(self, chunk_size: int | None = None, dim: int = 0) -> None: """ Sets the attention processor to use [feed forward chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers). @@ -1875,19 +1875,20 @@ def unfuse_qkv_projections(self): if self.original_attn_processors is not None: self.set_attn_processor(self.original_attn_processors) + @apply_lora_scale("cross_attention_kwargs") def forward( self, sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], + timestep: torch.Tensor | float | int, encoder_hidden_states: torch.Tensor, - timestep_cond: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, - down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None, - mid_block_additional_residual: Optional[torch.Tensor] = None, + timestep_cond: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + added_cond_kwargs: dict[str, torch.Tensor] | None = None, + down_block_additional_residuals: tuple[torch.Tensor] | None = None, + mid_block_additional_residual: torch.Tensor | None = None, return_dict: bool = True, - ) -> Union[UNetMotionOutput, Tuple[torch.Tensor]]: + ) -> UNetMotionOutput | tuple[torch.Tensor]: r""" The [`UNetMotionModel`] forward method. diff --git a/src/diffusers/models/unets/unet_spatio_temporal_condition.py b/src/diffusers/models/unets/unet_spatio_temporal_condition.py index c0cd5fbdd489..eddeb9826b0c 100644 --- a/src/diffusers/models/unets/unet_spatio_temporal_condition.py +++ b/src/diffusers/models/unets/unet_spatio_temporal_condition.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import Optional, Tuple, Union import torch import torch.nn as nn @@ -39,29 +38,29 @@ class UNetSpatioTemporalConditionModel(ModelMixin, AttentionMixin, ConfigMixin, for all models (such as downloading or saving). Parameters: - sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`): + sample_size (`int` or `tuple[int, int]`, *optional*, defaults to `None`): Height and width of input/output sample. in_channels (`int`, *optional*, defaults to 8): Number of channels in the input sample. out_channels (`int`, *optional*, defaults to 4): Number of channels in the output. - down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "DownBlockSpatioTemporal")`): + down_block_types (`tuple[str]`, *optional*, defaults to `("CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "DownBlockSpatioTemporal")`): The tuple of downsample blocks to use. - up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal")`): + up_block_types (`tuple[str]`, *optional*, defaults to `("UpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal")`): The tuple of upsample blocks to use. - block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): + block_out_channels (`tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): The tuple of output channels for each block. addition_time_embed_dim: (`int`, defaults to 256): Dimension to to encode the additional time ids. projection_class_embeddings_input_dim (`int`, defaults to 768): The dimension of the projection of encoded `added_time_ids`. layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block. - cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280): + cross_attention_dim (`int` or `tuple[int]`, *optional*, defaults to 1280): The dimension of the cross attention features. - transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple]` , *optional*, defaults to 1): + transformer_layers_per_block (`int`, `tuple[int]`, or `tuple[tuple]` , *optional*, defaults to 1): The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for [`~models.unets.unet_3d_blocks.CrossAttnDownBlockSpatioTemporal`], [`~models.unets.unet_3d_blocks.CrossAttnUpBlockSpatioTemporal`], [`~models.unets.unet_3d_blocks.UNetMidBlockSpatioTemporal`]. - num_attention_heads (`int`, `Tuple[int]`, defaults to `(5, 10, 10, 20)`): + num_attention_heads (`int`, `tuple[int]`, defaults to `(5, 10, 10, 20)`): The number of attention heads. dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. """ @@ -71,28 +70,28 @@ class UNetSpatioTemporalConditionModel(ModelMixin, AttentionMixin, ConfigMixin, @register_to_config def __init__( self, - sample_size: Optional[int] = None, + sample_size: int | None = None, in_channels: int = 8, out_channels: int = 4, - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str, ...] = ( "CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "DownBlockSpatioTemporal", ), - up_block_types: Tuple[str, ...] = ( + up_block_types: tuple[str, ...] = ( "UpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", ), - block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280), + block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280), addition_time_embed_dim: int = 256, projection_class_embeddings_input_dim: int = 768, - layers_per_block: Union[int, Tuple[int]] = 2, - cross_attention_dim: Union[int, Tuple[int]] = 1024, - transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1, - num_attention_heads: Union[int, Tuple[int, ...]] = (5, 10, 20, 20), + layers_per_block: int | tuple[int] = 2, + cross_attention_dim: int | tuple[int] = 1024, + transformer_layers_per_block: int | tuple[int, tuple[tuple]] = 1, + num_attention_heads: int | tuple[int, ...] = (5, 10, 20, 20), num_frames: int = 25, ): super().__init__() @@ -260,7 +259,7 @@ def set_default_attn_processor(self): self.set_attn_processor(processor) # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking - def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None: + def enable_forward_chunking(self, chunk_size: int | None = None, dim: int = 0) -> None: """ Sets the attention processor to use [feed forward chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers). @@ -292,11 +291,11 @@ def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int def forward( self, sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], + timestep: torch.Tensor | float | int, encoder_hidden_states: torch.Tensor, added_time_ids: torch.Tensor, return_dict: bool = True, - ) -> Union[UNetSpatioTemporalConditionOutput, Tuple]: + ) -> UNetSpatioTemporalConditionOutput | tuple: r""" The [`UNetSpatioTemporalConditionModel`] forward method. diff --git a/src/diffusers/models/unets/unet_stable_cascade.py b/src/diffusers/models/unets/unet_stable_cascade.py index 23d358c1bf51..5a6f24ab794b 100644 --- a/src/diffusers/models/unets/unet_stable_cascade.py +++ b/src/diffusers/models/unets/unet_stable_cascade.py @@ -14,7 +14,6 @@ import math from dataclasses import dataclass -from typing import Optional, Tuple, Union import numpy as np import torch @@ -145,30 +144,30 @@ def __init__( timestep_ratio_embedding_dim: int = 64, patch_size: int = 1, conditioning_dim: int = 2048, - block_out_channels: Tuple[int, ...] = (2048, 2048), - num_attention_heads: Tuple[int, ...] = (32, 32), - down_num_layers_per_block: Tuple[int, ...] = (8, 24), - up_num_layers_per_block: Tuple[int, ...] = (24, 8), - down_blocks_repeat_mappers: Optional[Tuple[int]] = ( + block_out_channels: tuple[int, ...] = (2048, 2048), + num_attention_heads: tuple[int, ...] = (32, 32), + down_num_layers_per_block: tuple[int, ...] = (8, 24), + up_num_layers_per_block: tuple[int, ...] = (24, 8), + down_blocks_repeat_mappers: tuple[int] | None = ( 1, 1, ), - up_blocks_repeat_mappers: Optional[Tuple[int]] = (1, 1), - block_types_per_layer: Tuple[Tuple[str]] = ( + up_blocks_repeat_mappers: tuple[int] | None = (1, 1), + block_types_per_layer: tuple[tuple[str]] = ( ("SDCascadeResBlock", "SDCascadeTimestepBlock", "SDCascadeAttnBlock"), ("SDCascadeResBlock", "SDCascadeTimestepBlock", "SDCascadeAttnBlock"), ), - clip_text_in_channels: Optional[int] = None, + clip_text_in_channels: int | None = None, clip_text_pooled_in_channels=1280, - clip_image_in_channels: Optional[int] = None, + clip_image_in_channels: int | None = None, clip_seq=4, - effnet_in_channels: Optional[int] = None, - pixel_mapper_in_channels: Optional[int] = None, + effnet_in_channels: int | None = None, + pixel_mapper_in_channels: int | None = None, kernel_size=3, - dropout: Union[float, Tuple[float]] = (0.1, 0.1), - self_attn: Union[bool, Tuple[bool]] = True, - timestep_conditioning_type: Tuple[str, ...] = ("sca", "crp"), - switch_level: Optional[Tuple[bool]] = None, + dropout: float | tuple[float] = (0.1, 0.1), + self_attn: bool | tuple[bool] = True, + timestep_conditioning_type: tuple[str, ...] = ("sca", "crp"), + switch_level: tuple[bool] | None = None, ): """ @@ -183,20 +182,20 @@ def __init__( Patch size to use for pixel unshuffling layer conditioning_dim (`int`, defaults to 2048): Dimension of the image and text conditional embedding. - block_out_channels (Tuple[int], defaults to (2048, 2048)): - Tuple of output channels for each block. - num_attention_heads (Tuple[int], defaults to (32, 32)): + block_out_channels (tuple[int], defaults to (2048, 2048)): + tuple of output channels for each block. + num_attention_heads (tuple[int], defaults to (32, 32)): Number of attention heads in each attention block. Set to -1 to if block types in a layer do not have attention. - down_num_layers_per_block (Tuple[int], defaults to [8, 24]): + down_num_layers_per_block (tuple[int], defaults to [8, 24]): Number of layers in each down block. - up_num_layers_per_block (Tuple[int], defaults to [24, 8]): + up_num_layers_per_block (tuple[int], defaults to [24, 8]): Number of layers in each up block. - down_blocks_repeat_mappers (Tuple[int], optional, defaults to [1, 1]): + down_blocks_repeat_mappers (tuple[int], optional, defaults to [1, 1]): Number of 1x1 Convolutional layers to repeat in each down block. - up_blocks_repeat_mappers (Tuple[int], optional, defaults to [1, 1]): + up_blocks_repeat_mappers (tuple[int], optional, defaults to [1, 1]): Number of 1x1 Convolutional layers to repeat in each up block. - block_types_per_layer (Tuple[Tuple[str]], optional, + block_types_per_layer (tuple[tuple[str]], optional, defaults to ( ("SDCascadeResBlock", "SDCascadeTimestepBlock", "SDCascadeAttnBlock"), ("SDCascadeResBlock", "SDCascadeTimestepBlock", "SDCascadeAttnBlock") @@ -214,14 +213,14 @@ def __init__( Number of input channels for pixel mapper conditioning. kernel_size (`int`, *optional*, defaults to 3): Kernel size to use in the block convolutional layers. - dropout (Tuple[float], *optional*, defaults to (0.1, 0.1)): + dropout (tuple[float], *optional*, defaults to (0.1, 0.1)): Dropout to use per block. - self_attn (Union[bool, Tuple[bool]]): - Tuple of booleans that determine whether to use self attention in a block or not. - timestep_conditioning_type (Tuple[str], defaults to ("sca", "crp")): + self_attn (bool | tuple[bool]): + tuple of booleans that determine whether to use self attention in a block or not. + timestep_conditioning_type (tuple[str], defaults to ("sca", "crp")): Timestep conditioning type. - switch_level (Optional[Tuple[bool]], *optional*, defaults to `None`): - Tuple that indicates whether upsampling or downsampling should be applied in a block + switch_level (tuple[bool] | None, *optional*, defaults to `None`): + tuple that indicates whether upsampling or downsampling should be applied in a block """ super().__init__() diff --git a/src/diffusers/models/unets/uvit_2d.py b/src/diffusers/models/unets/uvit_2d.py index 4c99ef88ca19..836d41a7f946 100644 --- a/src/diffusers/models/unets/uvit_2d.py +++ b/src/diffusers/models/unets/uvit_2d.py @@ -21,6 +21,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin +from ...utils import apply_lora_scale from ..attention import AttentionMixin, BasicTransformerBlock, SkipFFTransformerBlock from ..attention_processor import ( ADDED_KV_ATTENTION_PROCESSORS, @@ -146,6 +147,7 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("cross_attention_kwargs") def forward(self, input_ids, encoder_hidden_states, pooled_text_emb, micro_conds, cross_attention_kwargs=None): encoder_hidden_states = self.encoder_proj(encoder_hidden_states) encoder_hidden_states = self.encoder_proj_layer_norm(encoder_hidden_states) diff --git a/src/diffusers/models/upsampling.py b/src/diffusers/models/upsampling.py index 8a47c69f1264..cd3986287303 100644 --- a/src/diffusers/models/upsampling.py +++ b/src/diffusers/models/upsampling.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Tuple - import torch import torch.nn as nn import torch.nn.functional as F @@ -44,7 +42,7 @@ def __init__( channels: int, use_conv: bool = False, use_conv_transpose: bool = False, - out_channels: Optional[int] = None, + out_channels: int | None = None, name: str = "conv", ): super().__init__() @@ -94,9 +92,9 @@ def __init__( channels: int, use_conv: bool = False, use_conv_transpose: bool = False, - out_channels: Optional[int] = None, + out_channels: int | None = None, name: str = "conv", - kernel_size: Optional[int] = None, + kernel_size: int | None = None, padding=1, norm_type=None, eps=None, @@ -139,7 +137,7 @@ def __init__( else: self.Conv2d_0 = conv - def forward(self, hidden_states: torch.Tensor, output_size: Optional[int] = None, *args, **kwargs) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, output_size: int | None = None, *args, **kwargs) -> torch.Tensor: if len(args) > 0 or kwargs.get("scale", None) is not None: deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`." deprecate("scale", "1.0.0", deprecation_message) @@ -208,10 +206,10 @@ class FirUpsample2D(nn.Module): def __init__( self, - channels: Optional[int] = None, - out_channels: Optional[int] = None, + channels: int | None = None, + out_channels: int | None = None, use_conv: bool = False, - fir_kernel: Tuple[int, int, int, int] = (1, 3, 3, 1), + fir_kernel: tuple[int, int, int, int] = (1, 3, 3, 1), ): super().__init__() out_channels = out_channels if out_channels else channels @@ -224,8 +222,8 @@ def __init__( def _upsample_2d( self, hidden_states: torch.Tensor, - weight: Optional[torch.Tensor] = None, - kernel: Optional[torch.Tensor] = None, + weight: torch.Tensor | None = None, + kernel: torch.Tensor | None = None, factor: int = 2, gain: float = 1, ) -> torch.Tensor: @@ -425,7 +423,7 @@ def upfirdn2d_native( kernel: torch.Tensor, up: int = 1, down: int = 1, - pad: Tuple[int, int] = (0, 0), + pad: tuple[int, int] = (0, 0), ) -> torch.Tensor: up_x = up_y = up down_x = down_y = down @@ -472,7 +470,7 @@ def upfirdn2d_native( def upsample_2d( hidden_states: torch.Tensor, - kernel: Optional[torch.Tensor] = None, + kernel: torch.Tensor | None = None, factor: int = 2, gain: float = 1, ) -> torch.Tensor: diff --git a/src/diffusers/models/vae_flax.py b/src/diffusers/models/vae_flax.py index 5aad386a89e8..c7042840e4e0 100644 --- a/src/diffusers/models/vae_flax.py +++ b/src/diffusers/models/vae_flax.py @@ -16,7 +16,6 @@ import math from functools import partial -from typing import Tuple import flax import flax.linen as nn @@ -514,10 +513,10 @@ class FlaxEncoder(nn.Module): Input channels out_channels (:obj:`int`, *optional*, defaults to 3): Output channels - down_block_types (:obj:`Tuple[str]`, *optional*, defaults to `(DownEncoderBlock2D)`): + down_block_types (:obj:`tuple[str]`, *optional*, defaults to `(DownEncoderBlock2D)`): DownEncoder block type - block_out_channels (:obj:`Tuple[str]`, *optional*, defaults to `(64,)`): - Tuple containing the number of output channels for each block + block_out_channels (:obj:`tuple[str]`, *optional*, defaults to `(64,)`): + tuple[ containing the number of output channels for each block layers_per_block (:obj:`int`, *optional*, defaults to `2`): Number of Resnet layer for each block norm_num_groups (:obj:`int`, *optional*, defaults to `32`): @@ -532,8 +531,8 @@ class FlaxEncoder(nn.Module): in_channels: int = 3 out_channels: int = 3 - down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",) - block_out_channels: Tuple[int, ...] = (64,) + down_block_types: tuple[str, ...] = ("DownEncoderBlock2D",) + block_out_channels: tuple[int, ...] = (64,) layers_per_block: int = 2 norm_num_groups: int = 32 act_fn: str = "silu" @@ -632,10 +631,10 @@ class FlaxDecoder(nn.Module): Input channels out_channels (:obj:`int`, *optional*, defaults to 3): Output channels - up_block_types (:obj:`Tuple[str]`, *optional*, defaults to `(UpDecoderBlock2D)`): + up_block_types (:obj:`tuple[str]`, *optional*, defaults to `(UpDecoderBlock2D)`): UpDecoder block type - block_out_channels (:obj:`Tuple[str]`, *optional*, defaults to `(64,)`): - Tuple containing the number of output channels for each block + block_out_channels (:obj:`tuple[str]`, *optional*, defaults to `(64,)`): + tuple[ containing the number of output channels for each block layers_per_block (:obj:`int`, *optional*, defaults to `2`): Number of Resnet layer for each block norm_num_groups (:obj:`int`, *optional*, defaults to `32`): @@ -650,8 +649,8 @@ class FlaxDecoder(nn.Module): in_channels: int = 3 out_channels: int = 3 - up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",) - block_out_channels: Tuple[int, ...] = (64,) + up_block_types: tuple[str, ...] = ("UpDecoderBlock2D",) + block_out_channels: tuple[int, ...] = (64,) layers_per_block: int = 2 norm_num_groups: int = 32 act_fn: str = "silu" @@ -794,12 +793,12 @@ class FlaxAutoencoderKL(nn.Module, FlaxModelMixin, ConfigMixin): Number of channels in the input image. out_channels (`int`, *optional*, defaults to 3): Number of channels in the output. - down_block_types (`Tuple[str]`, *optional*, defaults to `(DownEncoderBlock2D)`): - Tuple of downsample block types. - up_block_types (`Tuple[str]`, *optional*, defaults to `(UpDecoderBlock2D)`): - Tuple of upsample block types. - block_out_channels (`Tuple[str]`, *optional*, defaults to `(64,)`): - Tuple of block output channels. + down_block_types (`tuple[str]`, *optional*, defaults to `(DownEncoderBlock2D)`): + tuple[ of downsample block types. + up_block_types (`tuple[str]`, *optional*, defaults to `(UpDecoderBlock2D)`): + tuple[ of upsample block types. + block_out_channels (`tuple[str]`, *optional*, defaults to `(64,)`): + tuple[ of block output channels. layers_per_block (`int`, *optional*, defaults to `2`): Number of ResNet layer for each block. act_fn (`str`, *optional*, defaults to `silu`): @@ -823,9 +822,9 @@ class FlaxAutoencoderKL(nn.Module, FlaxModelMixin, ConfigMixin): in_channels: int = 3 out_channels: int = 3 - down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",) - up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",) - block_out_channels: Tuple[int, ...] = (64,) + down_block_types: tuple[str, ...] = ("DownEncoderBlock2D",) + up_block_types: tuple[str, ...] = ("UpDecoderBlock2D",) + block_out_channels: tuple[int, ...] = (64,) layers_per_block: int = 1 act_fn: str = "silu" latent_channels: int = 4 diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 5fcc1a176d1b..c9bebd8644f7 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -33,6 +33,7 @@ "ModularPipeline", "AutoPipelineBlocks", "SequentialPipelineBlocks", + "ConditionalPipelineBlocks", "LoopSequentialPipelineBlocks", "PipelineState", "BlockState", @@ -45,7 +46,16 @@ "InsertableDict", ] _import_structure["stable_diffusion_xl"] = ["StableDiffusionXLAutoBlocks", "StableDiffusionXLModularPipeline"] - _import_structure["wan"] = ["WanAutoBlocks", "Wan22AutoBlocks", "WanModularPipeline"] + _import_structure["wan"] = [ + "WanBlocks", + "Wan22Blocks", + "WanImage2VideoAutoBlocks", + "Wan22Image2VideoBlocks", + "WanModularPipeline", + "Wan22ModularPipeline", + "WanImage2VideoModularPipeline", + "Wan22Image2VideoModularPipeline", + ] _import_structure["flux"] = [ "FluxAutoBlocks", "FluxModularPipeline", @@ -54,7 +64,11 @@ ] _import_structure["flux2"] = [ "Flux2AutoBlocks", + "Flux2KleinAutoBlocks", + "Flux2KleinBaseAutoBlocks", "Flux2ModularPipeline", + "Flux2KleinModularPipeline", + "Flux2KleinBaseModularPipeline", ] _import_structure["qwenimage"] = [ "QwenImageAutoBlocks", @@ -63,6 +77,8 @@ "QwenImageEditAutoBlocks", "QwenImageEditPlusModularPipeline", "QwenImageEditPlusAutoBlocks", + "QwenImageLayeredModularPipeline", + "QwenImageLayeredAutoBlocks", ] _import_structure["z_image"] = [ "ZImageAutoBlocks", @@ -79,10 +95,18 @@ else: from .components_manager import ComponentsManager from .flux import FluxAutoBlocks, FluxKontextAutoBlocks, FluxKontextModularPipeline, FluxModularPipeline - from .flux2 import Flux2AutoBlocks, Flux2ModularPipeline + from .flux2 import ( + Flux2AutoBlocks, + Flux2KleinAutoBlocks, + Flux2KleinBaseAutoBlocks, + Flux2KleinBaseModularPipeline, + Flux2KleinModularPipeline, + Flux2ModularPipeline, + ) from .modular_pipeline import ( AutoPipelineBlocks, BlockState, + ConditionalPipelineBlocks, LoopSequentialPipelineBlocks, ModularPipeline, ModularPipelineBlocks, @@ -96,10 +120,21 @@ QwenImageEditModularPipeline, QwenImageEditPlusAutoBlocks, QwenImageEditPlusModularPipeline, + QwenImageLayeredAutoBlocks, + QwenImageLayeredModularPipeline, QwenImageModularPipeline, ) from .stable_diffusion_xl import StableDiffusionXLAutoBlocks, StableDiffusionXLModularPipeline - from .wan import Wan22AutoBlocks, WanAutoBlocks, WanModularPipeline + from .wan import ( + Wan22Blocks, + Wan22Image2VideoBlocks, + Wan22Image2VideoModularPipeline, + Wan22ModularPipeline, + WanBlocks, + WanImage2VideoAutoBlocks, + WanImage2VideoModularPipeline, + WanModularPipeline, + ) from .z_image import ZImageAutoBlocks, ZImageModularPipeline else: import sys diff --git a/src/diffusers/modular_pipelines/components_manager.py b/src/diffusers/modular_pipelines/components_manager.py index cb7e8fb73697..e018381ba859 100644 --- a/src/diffusers/modular_pipelines/components_manager.py +++ b/src/diffusers/modular_pipelines/components_manager.py @@ -12,11 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import copy import time from collections import OrderedDict from itertools import combinations -from typing import Any, Dict, List, Optional, Union +from typing import Any import torch @@ -53,9 +55,9 @@ class CustomOffloadHook(ModelHook): def __init__( self, - execution_device: Optional[Union[str, int, torch.device]] = None, - other_hooks: Optional[List["UserCustomOffloadHook"]] = None, - offload_strategy: Optional["AutoOffloadStrategy"] = None, + execution_device: str | int | torch.device | None = None, + other_hooks: list["UserCustomOffloadHook"] | None = None, + offload_strategy: "AutoOffloadStrategy" | None = None, ): self.execution_device = execution_device if execution_device is not None else PartialState().default_device self.other_hooks = other_hooks @@ -135,8 +137,8 @@ def add_other_hook(self, hook: "UserCustomOffloadHook"): def custom_offload_with_hook( model_id: str, model: torch.nn.Module, - execution_device: Union[str, int, torch.device] = None, - offload_strategy: Optional["AutoOffloadStrategy"] = None, + execution_device: str | int | torch.device = None, + offload_strategy: "AutoOffloadStrategy" | None = None, ): hook = CustomOffloadHook(execution_device=execution_device, offload_strategy=offload_strategy) user_hook = UserCustomOffloadHook(model_id=model_id, model=model, hook=hook) @@ -160,7 +162,10 @@ def __call__(self, hooks, model_id, model, execution_device): if len(hooks) == 0: return [] - current_module_size = model.get_memory_footprint() + try: + current_module_size = model.get_memory_footprint() + except AttributeError: + raise AttributeError(f"Do not know how to compute memory footprint of `{model.__class__.__name__}.") device_type = execution_device.type device_module = getattr(torch, device_type, torch.cuda) @@ -223,7 +228,7 @@ def search_best_candidate(module_sizes, min_memory_offload): # utils for display component info in a readable format # TODO: move to a different file -def summarize_dict_by_value_and_parts(d: Dict[str, Any]) -> Dict[str, Any]: +def summarize_dict_by_value_and_parts(d: dict[str, Any]) -> dict[str, Any]: """Summarizes a dictionary by finding common prefixes that share the same value. For a dictionary with dot-separated keys like: { @@ -244,7 +249,7 @@ def summarize_dict_by_value_and_parts(d: Dict[str, Any]) -> Dict[str, Any]: value_to_keys[value_tuple] = [] value_to_keys[value_tuple].append(key) - def find_common_prefix(keys: List[str]) -> str: + def find_common_prefix(keys: list[str]) -> str: """Find the shortest common prefix among a list of dot-separated keys.""" if not keys: return "" @@ -321,6 +326,7 @@ class ComponentsManager: "has_hook", "execution_device", "ip_adapter", + "quantization", ] def __init__(self): @@ -333,10 +339,10 @@ def __init__(self): def _lookup_ids( self, - name: Optional[str] = None, - collection: Optional[str] = None, - load_id: Optional[str] = None, - components: Optional[OrderedDict] = None, + name: str | None = None, + collection: str | None = None, + load_id: str | None = None, + components: OrderedDict | None = None, ): """ Lookup component_ids by name, collection, or load_id. Does not support pattern matching. Returns a set of @@ -353,7 +359,9 @@ def _lookup_ids( ids_by_name.add(component_id) else: ids_by_name = set(components.keys()) - if collection: + if collection and collection not in self.collections: + return set() + elif collection and collection in self.collections: ids_by_collection = set() for component_id, component in components.items(): if component_id in self.collections[collection]: @@ -375,14 +383,14 @@ def _lookup_ids( def _id_to_name(component_id: str): return "_".join(component_id.split("_")[:-1]) - def add(self, name: str, component: Any, collection: Optional[str] = None): + def add(self, name: str, component: Any, collection: str | None = None): """ Add a component to the ComponentsManager. Args: name (str): The name of the component component (Any): The component to add - collection (Optional[str]): The collection to add the component to + collection (str | None): The collection to add the component to Returns: str: The unique component ID, which is generated as "{name}_{id(component)}" where @@ -420,7 +428,8 @@ def add(self, name: str, component: Any, collection: Optional[str] = None): # add component to components manager self.components[component_id] = component - self.added_time[component_id] = time.time() + if is_new_component: + self.added_time[component_id] = time.time() if collection: if collection not in self.collections: @@ -499,9 +508,9 @@ def remove(self, component_id: str = None): # YiYi TODO: rename to search_components for now, may remove this method def search_components( self, - names: Optional[str] = None, - collection: Optional[str] = None, - load_id: Optional[str] = None, + names: str | None = None, + collection: str | None = None, + load_id: str | None = None, return_dict_with_names: bool = True, ): """ @@ -683,7 +692,7 @@ def matches_pattern(component_id, pattern, exact_match=False): return get_return_dict(matches, return_dict_with_names) - def enable_auto_cpu_offload(self, device: Union[str, int, torch.device] = None, memory_reserve_margin="3GB"): + def enable_auto_cpu_offload(self, device: str | int | torch.device = None, memory_reserve_margin="3GB"): """ Enable automatic CPU offloading for all components. @@ -695,7 +704,7 @@ def enable_auto_cpu_offload(self, device: Union[str, int, torch.device] = None, 5. Models stay on the execution device until another model needs memory and forces them off Args: - device (Union[str, int, torch.device]): The execution device where models are moved for forward passes + device (str | int | torch.device): The execution device where models are moved for forward passes memory_reserve_margin (str): The memory reserve margin to use, default is 3GB. This is the amount of memory to keep free on the device to avoid running out of memory during model execution (e.g., for intermediate activations, gradients, etc.) @@ -703,7 +712,20 @@ def enable_auto_cpu_offload(self, device: Union[str, int, torch.device] = None, if not is_accelerate_available(): raise ImportError("Make sure to install accelerate to use auto_cpu_offload") - # TODO: add a warning if mem_get_info isn't available on `device`. + if device is None: + device = get_device() + if not isinstance(device, torch.device): + device = torch.device(device) + + device_type = device.type + device_module = getattr(torch, device_type, torch.cuda) + if not hasattr(device_module, "mem_get_info"): + raise NotImplementedError( + f"`enable_auto_cpu_offload() relies on the `mem_get_info()` method. It's not implemented for {str(device.type)}." + ) + + if device.index is None: + device = torch.device(f"{device.type}:{0}") for name, component in self.components.items(): if isinstance(component, torch.nn.Module) and hasattr(component, "_hf_hook"): @@ -711,11 +733,7 @@ def enable_auto_cpu_offload(self, device: Union[str, int, torch.device] = None, self.disable_auto_cpu_offload() offload_strategy = AutoOffloadStrategy(memory_reserve_margin=memory_reserve_margin) - if device is None: - device = get_device() - device = torch.device(device) - if device.index is None: - device = torch.device(f"{device.type}:{0}") + all_hooks = [] for name, component in self.components.items(): if isinstance(component, torch.nn.Module): @@ -748,17 +766,16 @@ def disable_auto_cpu_offload(self): self.model_hooks = None self._auto_offload_enabled = False - # YiYi TODO: (1) add quantization info def get_model_info( self, component_id: str, - fields: Optional[Union[str, List[str]]] = None, - ) -> Optional[Dict[str, Any]]: + fields: str | list[str] | None = None, + ) -> dict[str, Any] | None: """Get comprehensive information about a component. Args: component_id (str): Name of the component to get info for - fields (Optional[Union[str, List[str]]]): + fields (str | list[str] | None): Field(s) to return. Can be a string for single field or list of fields. If None, uses the available_info_fields setting. @@ -824,6 +841,17 @@ def get_model_info( if scales: info["ip_adapter"] = summarize_dict_by_value_and_parts(scales) + # Check for quantization + hf_quantizer = getattr(component, "hf_quantizer", None) + if hf_quantizer is not None: + quant_config = hf_quantizer.quantization_config + if hasattr(quant_config, "to_diff_dict"): + info["quantization"] = quant_config.to_diff_dict() + else: + info["quantization"] = quant_config.to_dict() + else: + info["quantization"] = None + # If fields specified, filter info if fields is not None: return {k: v for k, v in info.items() if k in fields} @@ -954,21 +982,25 @@ def format_device(component, info): output += "\nAdditional Component Info:\n" + "=" * 50 + "\n" for name in self.components: info = self.get_model_info(name) - if info is not None and (info.get("adapters") is not None or info.get("ip_adapter")): + if info is not None and ( + info.get("adapters") is not None or info.get("ip_adapter") or info.get("quantization") + ): output += f"\n{name}:\n" if info.get("adapters") is not None: output += f" Adapters: {info['adapters']}\n" if info.get("ip_adapter"): output += " IP-Adapter: Enabled\n" + if info.get("quantization"): + output += f" Quantization: {info['quantization']}\n" return output def get_one( self, - component_id: Optional[str] = None, - name: Optional[str] = None, - collection: Optional[str] = None, - load_id: Optional[str] = None, + component_id: str | None = None, + name: str | None = None, + collection: str | None = None, + load_id: str | None = None, ) -> Any: """ Get a single component by either: @@ -977,10 +1009,10 @@ def get_one( Raises an error if multiple components match or none are found. Args: - component_id (Optional[str]): Optional component ID to get - name (Optional[str]): Component name or pattern - collection (Optional[str]): Optional collection to filter by - load_id (Optional[str]): Optional load_id to filter by + component_id (str | None): Optional component ID to get + name (str | None): Component name or pattern + collection (str | None): Optional collection to filter by + load_id (str | None): Optional load_id to filter by Returns: A single component @@ -1008,16 +1040,16 @@ def get_one( return next(iter(results.values())) - def get_ids(self, names: Union[str, List[str]] = None, collection: Optional[str] = None): + def get_ids(self, names: str | list[str] = None, collection: str | None = None): """ Get component IDs by a list of names, optionally filtered by collection. Args: - names (Union[str, List[str]]): List of component names - collection (Optional[str]): Optional collection to filter by + names (str | list[str]): list of component names + collection (str | None): Optional collection to filter by Returns: - List[str]: List of component IDs + list[str]: list of component IDs """ ids = set() if not isinstance(names, list): @@ -1026,18 +1058,18 @@ def get_ids(self, names: Union[str, List[str]] = None, collection: Optional[str] ids.update(self._lookup_ids(name=name, collection=collection)) return list(ids) - def get_components_by_ids(self, ids: List[str], return_dict_with_names: Optional[bool] = True): + def get_components_by_ids(self, ids: list[str], return_dict_with_names: bool | None = True): """ Get components by a list of IDs. Args: - ids (List[str]): - List of component IDs - return_dict_with_names (Optional[bool]): + ids (list[str]): + list of component IDs + return_dict_with_names (bool | None): Whether to return a dictionary with component names as keys: Returns: - Dict[str, Any]: Dictionary of components. + dict[str, Any]: Dictionary of components. - If return_dict_with_names=True, keys are component names. - If return_dict_with_names=False, keys are component IDs. @@ -1059,16 +1091,16 @@ def get_components_by_ids(self, ids: List[str], return_dict_with_names: Optional else: return components - def get_components_by_names(self, names: List[str], collection: Optional[str] = None): + def get_components_by_names(self, names: list[str], collection: str | None = None): """ Get components by a list of names, optionally filtered by collection. Args: - names (List[str]): List of component names - collection (Optional[str]): Optional collection to filter by + names (list[str]): list of component names + collection (str | None): Optional collection to filter by Returns: - Dict[str, Any]: Dictionary of components with component names as keys + dict[str, Any]: Dictionary of components with component names as keys Raises: ValueError: If duplicate component names are found in the search results diff --git a/src/diffusers/modular_pipelines/flux/__init__.py b/src/diffusers/modular_pipelines/flux/__init__.py index ec00986611c8..4754ed01ce6a 100644 --- a/src/diffusers/modular_pipelines/flux/__init__.py +++ b/src/diffusers/modular_pipelines/flux/__init__.py @@ -21,21 +21,8 @@ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) else: - _import_structure["encoders"] = ["FluxTextEncoderStep"] - _import_structure["modular_blocks"] = [ - "ALL_BLOCKS", - "AUTO_BLOCKS", - "AUTO_BLOCKS_KONTEXT", - "FLUX_KONTEXT_BLOCKS", - "TEXT2IMAGE_BLOCKS", - "FluxAutoBeforeDenoiseStep", - "FluxAutoBlocks", - "FluxAutoDecodeStep", - "FluxAutoDenoiseStep", - "FluxKontextAutoBlocks", - "FluxKontextAutoDenoiseStep", - "FluxKontextBeforeDenoiseStep", - ] + _import_structure["modular_blocks_flux"] = ["FluxAutoBlocks"] + _import_structure["modular_blocks_flux_kontext"] = ["FluxKontextAutoBlocks"] _import_structure["modular_pipeline"] = ["FluxKontextModularPipeline", "FluxModularPipeline"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: @@ -45,21 +32,8 @@ except OptionalDependencyNotAvailable: from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 else: - from .encoders import FluxTextEncoderStep - from .modular_blocks import ( - ALL_BLOCKS, - AUTO_BLOCKS, - AUTO_BLOCKS_KONTEXT, - FLUX_KONTEXT_BLOCKS, - TEXT2IMAGE_BLOCKS, - FluxAutoBeforeDenoiseStep, - FluxAutoBlocks, - FluxAutoDecodeStep, - FluxAutoDenoiseStep, - FluxKontextAutoBlocks, - FluxKontextAutoDenoiseStep, - FluxKontextBeforeDenoiseStep, - ) + from .modular_blocks_flux import FluxAutoBlocks + from .modular_blocks_flux_kontext import FluxKontextAutoBlocks from .modular_pipeline import FluxKontextModularPipeline, FluxModularPipeline else: import sys diff --git a/src/diffusers/modular_pipelines/flux/before_denoise.py b/src/diffusers/modular_pipelines/flux/before_denoise.py index daffec986535..c28154775f5a 100644 --- a/src/diffusers/modular_pipelines/flux/before_denoise.py +++ b/src/diffusers/modular_pipelines/flux/before_denoise.py @@ -13,7 +13,6 @@ # limitations under the License. import inspect -from typing import List, Optional, Union import numpy as np import torch @@ -33,10 +32,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -51,15 +50,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -106,7 +105,7 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -156,7 +155,7 @@ class FluxSetTimestepsStep(ModularPipelineBlocks): model_name = "flux" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] @property @@ -164,7 +163,7 @@ def description(self) -> str: return "Step that sets the scheduler's timesteps for inference" @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("num_inference_steps", default=50), InputParam("timesteps"), @@ -183,7 +182,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference"), OutputParam( @@ -232,7 +231,7 @@ class FluxImg2ImgSetTimestepsStep(ModularPipelineBlocks): model_name = "flux" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] @property @@ -240,7 +239,7 @@ def description(self) -> str: return "Step that sets the scheduler's timesteps for inference" @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("num_inference_steps", default=50), InputParam("timesteps"), @@ -259,7 +258,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference"), OutputParam( @@ -322,7 +321,7 @@ class FluxPrepareLatentsStep(ModularPipelineBlocks): model_name = "flux" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [] @property @@ -330,11 +329,11 @@ def description(self) -> str: return "Prepare latents step that prepares the latents for the text-to-image generation process" @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("height", type_hint=int), InputParam("width", type_hint=int), - InputParam("latents", type_hint=Optional[torch.Tensor]), + InputParam("latents", type_hint=torch.Tensor | None), InputParam("num_images_per_prompt", type_hint=int, default=1), InputParam("generator"), InputParam( @@ -347,7 +346,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process" @@ -431,11 +430,11 @@ def description(self) -> str: " `prepare_latents`. Both noise and image latents should already be patchified." @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam( name="latents", @@ -458,7 +457,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( name="initial_noise", @@ -507,7 +506,7 @@ def description(self) -> str: return "Step that prepares the RoPE inputs for the denoising process. Should be placed after text encoder and latent preparation steps." @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam(name="height", required=True), InputParam(name="width", required=True), @@ -515,18 +514,18 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( name="txt_ids", kwargs_type="denoiser_input_fields", - type_hint=List[int], + type_hint=list[int], description="The sequence lengths of the prompt embeds, used for RoPE calculation.", ), OutputParam( name="img_ids", kwargs_type="denoiser_input_fields", - type_hint=List[int], + type_hint=list[int], description="The sequence lengths of the image latents, used for RoPE calculation.", ), ] @@ -557,7 +556,7 @@ def description(self) -> str: return "Step that prepares the RoPE inputs for the denoising process of Flux Kontext. Should be placed after text encoder and latent preparation steps." @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam(name="image_height"), InputParam(name="image_width"), @@ -567,18 +566,18 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( name="txt_ids", kwargs_type="denoiser_input_fields", - type_hint=List[int], + type_hint=list[int], description="The sequence lengths of the prompt embeds, used for RoPE calculation.", ), OutputParam( name="img_ids", kwargs_type="denoiser_input_fields", - type_hint=List[int], + type_hint=list[int], description="The sequence lengths of the image latents, used for RoPE calculation.", ), ] diff --git a/src/diffusers/modular_pipelines/flux/decoders.py b/src/diffusers/modular_pipelines/flux/decoders.py index 846549b1a376..5da861e78fcb 100644 --- a/src/diffusers/modular_pipelines/flux/decoders.py +++ b/src/diffusers/modular_pipelines/flux/decoders.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, List, Tuple, Union +from typing import Any import numpy as np import PIL @@ -49,7 +49,7 @@ class FluxDecodeStep(ModularPipelineBlocks): model_name = "flux" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("vae", AutoencoderKL), ComponentSpec( @@ -65,7 +65,7 @@ def description(self) -> str: return "Step that decodes the denoised latents into images" @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [ InputParam("output_type", default="pil"), InputParam("height", default=1024), @@ -79,11 +79,11 @@ def inputs(self) -> List[Tuple[str, Any]]: ] @property - def intermediate_outputs(self) -> List[str]: + def intermediate_outputs(self) -> list[str]: return [ OutputParam( "images", - type_hint=Union[List[PIL.Image.Image], torch.Tensor, np.ndarray], + type_hint=list[PIL.Image.Image] | torch.Tensor | np.ndarray, description="The generated images, can be a list of PIL.Image.Image, torch.Tensor or a numpy array", ) ] diff --git a/src/diffusers/modular_pipelines/flux/denoise.py b/src/diffusers/modular_pipelines/flux/denoise.py index 5a769df1036d..babb4a867e59 100644 --- a/src/diffusers/modular_pipelines/flux/denoise.py +++ b/src/diffusers/modular_pipelines/flux/denoise.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, List, Tuple +from typing import Any import torch @@ -36,7 +36,7 @@ class FluxLoopDenoiser(ModularPipelineBlocks): model_name = "flux" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("transformer", FluxTransformer2DModel)] @property @@ -48,7 +48,7 @@ def description(self) -> str: ) @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [ InputParam("joint_attention_kwargs"), InputParam( @@ -113,7 +113,7 @@ class FluxKontextLoopDenoiser(ModularPipelineBlocks): model_name = "flux-kontext" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("transformer", FluxTransformer2DModel)] @property @@ -125,7 +125,7 @@ def description(self) -> str: ) @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [ InputParam("joint_attention_kwargs"), InputParam( @@ -203,7 +203,7 @@ class FluxLoopAfterDenoiser(ModularPipelineBlocks): model_name = "flux" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] @property @@ -215,15 +215,15 @@ def description(self) -> str: ) @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [] @property - def intermediate_inputs(self) -> List[str]: + def intermediate_inputs(self) -> list[str]: return [InputParam("generator")] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")] @torch.no_grad() @@ -254,14 +254,14 @@ def description(self) -> str: ) @property - def loop_expected_components(self) -> List[ComponentSpec]: + def loop_expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), ComponentSpec("transformer", FluxTransformer2DModel), ] @property - def loop_inputs(self) -> List[InputParam]: + def loop_inputs(self) -> list[InputParam]: return [ InputParam( "timesteps", diff --git a/src/diffusers/modular_pipelines/flux/encoders.py b/src/diffusers/modular_pipelines/flux/encoders.py index f0314d4771b0..583c139ff22e 100644 --- a/src/diffusers/modular_pipelines/flux/encoders.py +++ b/src/diffusers/modular_pipelines/flux/encoders.py @@ -13,7 +13,6 @@ # limitations under the License. import html -from typing import List, Optional, Union import regex as re import torch @@ -55,7 +54,7 @@ def prompt_clean(text): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -90,7 +89,7 @@ def description(self) -> str: return "Image Preprocess step." @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec( "image_processor", @@ -101,11 +100,11 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [InputParam("resized_image"), InputParam("image"), InputParam("height"), InputParam("width")] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [OutputParam(name="processed_image")] @staticmethod @@ -151,7 +150,7 @@ def description(self) -> str: ) @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec( "image_processor", @@ -162,11 +161,11 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [InputParam("image"), InputParam("_auto_resize", type_hint=bool, default=True)] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [OutputParam(name="processed_image")] @torch.no_grad() @@ -206,7 +205,7 @@ def __call__(self, components: FluxModularPipeline, state: PipelineState): return components, state -class FluxVaeEncoderDynamicStep(ModularPipelineBlocks): +class FluxVaeEncoderStep(ModularPipelineBlocks): model_name = "flux" def __init__( @@ -241,17 +240,17 @@ def description(self) -> str: return f"Dynamic VAE Encoder step that converts {self._image_input_name} into latent representations {self._image_latents_output_name}.\n" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: components = [ComponentSpec("vae", AutoencoderKL)] return components @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: inputs = [InputParam(self._image_input_name), InputParam("generator")] return inputs @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( self._image_latents_output_name, @@ -291,7 +290,7 @@ def description(self) -> str: return "Text Encoder step that generate text_embeddings to guide the image generation" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("text_encoder", CLIPTextModel), ComponentSpec("tokenizer", CLIPTokenizer), @@ -300,7 +299,7 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("prompt"), InputParam("prompt_2"), @@ -309,7 +308,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "prompt_embeds", @@ -332,9 +331,7 @@ def check_inputs(block_state): raise ValueError(f"`prompt` or `prompt_2` has to be of type `str` or `list` but is {type(prompt)}") @staticmethod - def _get_t5_prompt_embeds( - components, prompt: Union[str, List[str]], max_sequence_length: int, device: torch.device - ): + def _get_t5_prompt_embeds(components, prompt: str | list[str], max_sequence_length: int, device: torch.device): dtype = components.text_encoder_2.dtype prompt = [prompt] if isinstance(prompt, str) else prompt @@ -365,7 +362,7 @@ def _get_t5_prompt_embeds( return prompt_embeds @staticmethod - def _get_clip_prompt_embeds(components, prompt: Union[str, List[str]], device: torch.device): + def _get_clip_prompt_embeds(components, prompt: str | list[str], device: torch.device): prompt = [prompt] if isinstance(prompt, str) else prompt if isinstance(components, TextualInversionLoaderMixin): @@ -401,13 +398,13 @@ def _get_clip_prompt_embeds(components, prompt: Union[str, List[str]], device: t @staticmethod def encode_prompt( components, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], - device: Optional[torch.device] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + prompt: str | list[str], + prompt_2: str | list[str], + device: torch.device | None = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 512, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): device = device or components._execution_device diff --git a/src/diffusers/modular_pipelines/flux/inputs.py b/src/diffusers/modular_pipelines/flux/inputs.py index 8309eebfeb37..9d2f69dbe26f 100644 --- a/src/diffusers/modular_pipelines/flux/inputs.py +++ b/src/diffusers/modular_pipelines/flux/inputs.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List import torch @@ -42,7 +41,7 @@ def description(self) -> str: ) @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("num_images_per_prompt", default=1), InputParam( @@ -62,7 +61,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[str]: + def intermediate_outputs(self) -> list[str]: return [ OutputParam( "batch_size", @@ -121,14 +120,14 @@ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> Pip return components, state -# Adapted from `QwenImageInputsDynamicStep` -class FluxInputsDynamicStep(ModularPipelineBlocks): +# Adapted from `QwenImageAdditionalInputsStep` +class FluxAdditionalInputsStep(ModularPipelineBlocks): model_name = "flux" def __init__( self, - image_latent_inputs: List[str] = ["image_latents"], - additional_batch_inputs: List[str] = [], + image_latent_inputs: list[str] = ["image_latents"], + additional_batch_inputs: list[str] = [], ): if not isinstance(image_latent_inputs, list): image_latent_inputs = [image_latent_inputs] @@ -163,7 +162,7 @@ def description(self) -> str: return summary_section + inputs_info + placement_section @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: inputs = [ InputParam(name="num_images_per_prompt", default=1), InputParam(name="batch_size", required=True), @@ -182,7 +181,7 @@ def inputs(self) -> List[InputParam]: return inputs @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam(name="image_height", type_hint=int, description="The height of the image latents"), OutputParam(name="image_width", type_hint=int, description="The width of the image latents"), @@ -244,7 +243,7 @@ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> Pip return components, state -class FluxKontextInputsDynamicStep(FluxInputsDynamicStep): +class FluxKontextAdditionalInputsStep(FluxAdditionalInputsStep): model_name = "flux-kontext" def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState: @@ -257,7 +256,7 @@ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> Pip continue # 1. Calculate height/width from latents - # Unlike the `FluxInputsDynamicStep`, we don't overwrite the `block.height` and `block.width` + # Unlike the `FluxAdditionalInputsStep`, we don't overwrite the `block.height` and `block.width` height, width = calculate_dimension_from_latents(image_latent_tensor, components.vae_scale_factor) if not hasattr(block_state, "image_height"): block_state.image_height = height @@ -304,6 +303,7 @@ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> Pip class FluxKontextSetResolutionStep(ModularPipelineBlocks): model_name = "flux-kontext" + @property def description(self): return ( "Determines the height and width to be used during the subsequent computations.\n" @@ -311,7 +311,7 @@ def description(self): ) @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: inputs = [ InputParam(name="height"), InputParam(name="width"), @@ -320,7 +320,7 @@ def inputs(self) -> List[InputParam]: return inputs @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam(name="height", type_hint=int, description="The height of the initial noisy latents"), OutputParam(name="width", type_hint=int, description="The width of the initial noisy latents"), diff --git a/src/diffusers/modular_pipelines/flux/modular_blocks.py b/src/diffusers/modular_pipelines/flux/modular_blocks.py deleted file mode 100644 index bd9b2d1b40c9..000000000000 --- a/src/diffusers/modular_pipelines/flux/modular_blocks.py +++ /dev/null @@ -1,446 +0,0 @@ -# Copyright 2025 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from ...utils import logging -from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks -from ..modular_pipeline_utils import InsertableDict -from .before_denoise import ( - FluxImg2ImgPrepareLatentsStep, - FluxImg2ImgSetTimestepsStep, - FluxKontextRoPEInputsStep, - FluxPrepareLatentsStep, - FluxRoPEInputsStep, - FluxSetTimestepsStep, -) -from .decoders import FluxDecodeStep -from .denoise import FluxDenoiseStep, FluxKontextDenoiseStep -from .encoders import ( - FluxKontextProcessImagesInputStep, - FluxProcessImagesInputStep, - FluxTextEncoderStep, - FluxVaeEncoderDynamicStep, -) -from .inputs import ( - FluxInputsDynamicStep, - FluxKontextInputsDynamicStep, - FluxKontextSetResolutionStep, - FluxTextInputStep, -) - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -# vae encoder (run before before_denoise) -FluxImg2ImgVaeEncoderBlocks = InsertableDict( - [("preprocess", FluxProcessImagesInputStep()), ("encode", FluxVaeEncoderDynamicStep())] -) - - -class FluxImg2ImgVaeEncoderStep(SequentialPipelineBlocks): - model_name = "flux" - - block_classes = FluxImg2ImgVaeEncoderBlocks.values() - block_names = FluxImg2ImgVaeEncoderBlocks.keys() - - @property - def description(self) -> str: - return "Vae encoder step that preprocess andencode the image inputs into their latent representations." - - -class FluxAutoVaeEncoderStep(AutoPipelineBlocks): - block_classes = [FluxImg2ImgVaeEncoderStep] - block_names = ["img2img"] - block_trigger_inputs = ["image"] - - @property - def description(self): - return ( - "Vae encoder step that encode the image inputs into their latent representations.\n" - + "This is an auto pipeline block that works for img2img tasks.\n" - + " - `FluxImg2ImgVaeEncoderStep` (img2img) is used when only `image` is provided." - + " - if `image` is not provided, step will be skipped." - ) - - -# Flux Kontext vae encoder (run before before_denoise) - -FluxKontextVaeEncoderBlocks = InsertableDict( - [("preprocess", FluxKontextProcessImagesInputStep()), ("encode", FluxVaeEncoderDynamicStep(sample_mode="argmax"))] -) - - -class FluxKontextVaeEncoderStep(SequentialPipelineBlocks): - model_name = "flux-kontext" - - block_classes = FluxKontextVaeEncoderBlocks.values() - block_names = FluxKontextVaeEncoderBlocks.keys() - - @property - def description(self) -> str: - return "Vae encoder step that preprocess andencode the image inputs into their latent representations." - - -class FluxKontextAutoVaeEncoderStep(AutoPipelineBlocks): - block_classes = [FluxKontextVaeEncoderStep] - block_names = ["img2img"] - block_trigger_inputs = ["image"] - - @property - def description(self): - return ( - "Vae encoder step that encode the image inputs into their latent representations.\n" - + "This is an auto pipeline block that works for img2img tasks.\n" - + " - `FluxKontextVaeEncoderStep` (img2img) is used when only `image` is provided." - + " - if `image` is not provided, step will be skipped." - ) - - -# before_denoise: text2img -FluxBeforeDenoiseBlocks = InsertableDict( - [ - ("prepare_latents", FluxPrepareLatentsStep()), - ("set_timesteps", FluxSetTimestepsStep()), - ("prepare_rope_inputs", FluxRoPEInputsStep()), - ] -) - - -class FluxBeforeDenoiseStep(SequentialPipelineBlocks): - block_classes = FluxBeforeDenoiseBlocks.values() - block_names = FluxBeforeDenoiseBlocks.keys() - - @property - def description(self): - return "Before denoise step that prepares the inputs for the denoise step in text-to-image generation." - - -# before_denoise: img2img -FluxImg2ImgBeforeDenoiseBlocks = InsertableDict( - [ - ("prepare_latents", FluxPrepareLatentsStep()), - ("set_timesteps", FluxImg2ImgSetTimestepsStep()), - ("prepare_img2img_latents", FluxImg2ImgPrepareLatentsStep()), - ("prepare_rope_inputs", FluxRoPEInputsStep()), - ] -) - - -class FluxImg2ImgBeforeDenoiseStep(SequentialPipelineBlocks): - block_classes = FluxImg2ImgBeforeDenoiseBlocks.values() - block_names = FluxImg2ImgBeforeDenoiseBlocks.keys() - - @property - def description(self): - return "Before denoise step that prepare the inputs for the denoise step for img2img task." - - -# before_denoise: all task (text2img, img2img) -class FluxAutoBeforeDenoiseStep(AutoPipelineBlocks): - model_name = "flux-kontext" - block_classes = [FluxImg2ImgBeforeDenoiseStep, FluxBeforeDenoiseStep] - block_names = ["img2img", "text2image"] - block_trigger_inputs = ["image_latents", None] - - @property - def description(self): - return ( - "Before denoise step that prepare the inputs for the denoise step.\n" - + "This is an auto pipeline block that works for text2image.\n" - + " - `FluxBeforeDenoiseStep` (text2image) is used.\n" - + " - `FluxImg2ImgBeforeDenoiseStep` (img2img) is used when only `image_latents` is provided.\n" - ) - - -# before_denoise: FluxKontext - -FluxKontextBeforeDenoiseBlocks = InsertableDict( - [ - ("prepare_latents", FluxPrepareLatentsStep()), - ("set_timesteps", FluxSetTimestepsStep()), - ("prepare_rope_inputs", FluxKontextRoPEInputsStep()), - ] -) - - -class FluxKontextBeforeDenoiseStep(SequentialPipelineBlocks): - block_classes = FluxKontextBeforeDenoiseBlocks.values() - block_names = FluxKontextBeforeDenoiseBlocks.keys() - - @property - def description(self): - return ( - "Before denoise step that prepare the inputs for the denoise step\n" - "for img2img/text2img task for Flux Kontext." - ) - - -class FluxKontextAutoBeforeDenoiseStep(AutoPipelineBlocks): - block_classes = [FluxKontextBeforeDenoiseStep, FluxBeforeDenoiseStep] - block_names = ["img2img", "text2image"] - block_trigger_inputs = ["image_latents", None] - - @property - def description(self): - return ( - "Before denoise step that prepare the inputs for the denoise step.\n" - + "This is an auto pipeline block that works for text2image.\n" - + " - `FluxBeforeDenoiseStep` (text2image) is used.\n" - + " - `FluxKontextBeforeDenoiseStep` (img2img) is used when only `image_latents` is provided.\n" - ) - - -# denoise: text2image -class FluxAutoDenoiseStep(AutoPipelineBlocks): - block_classes = [FluxDenoiseStep] - block_names = ["denoise"] - block_trigger_inputs = [None] - - @property - def description(self) -> str: - return ( - "Denoise step that iteratively denoise the latents. " - "This is a auto pipeline block that works for text2image and img2img tasks." - " - `FluxDenoiseStep` (denoise) for text2image and img2img tasks." - ) - - -# denoise: Flux Kontext - - -class FluxKontextAutoDenoiseStep(AutoPipelineBlocks): - block_classes = [FluxKontextDenoiseStep] - block_names = ["denoise"] - block_trigger_inputs = [None] - - @property - def description(self) -> str: - return ( - "Denoise step that iteratively denoise the latents for Flux Kontext. " - "This is a auto pipeline block that works for text2image and img2img tasks." - " - `FluxDenoiseStep` (denoise) for text2image and img2img tasks." - ) - - -# decode: all task (text2img, img2img) -class FluxAutoDecodeStep(AutoPipelineBlocks): - block_classes = [FluxDecodeStep] - block_names = ["non-inpaint"] - block_trigger_inputs = [None] - - @property - def description(self): - return "Decode step that decode the denoised latents into image outputs.\n - `FluxDecodeStep`" - - -# inputs: text2image/img2img -FluxImg2ImgBlocks = InsertableDict( - [("text_inputs", FluxTextInputStep()), ("additional_inputs", FluxInputsDynamicStep())] -) - - -class FluxImg2ImgInputStep(SequentialPipelineBlocks): - model_name = "flux" - block_classes = FluxImg2ImgBlocks.values() - block_names = FluxImg2ImgBlocks.keys() - - @property - def description(self): - return "Input step that prepares the inputs for the img2img denoising step. It:\n" - " - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents`).\n" - " - update height/width based `image_latents`, patchify `image_latents`." - - -class FluxAutoInputStep(AutoPipelineBlocks): - block_classes = [FluxImg2ImgInputStep, FluxTextInputStep] - block_names = ["img2img", "text2image"] - block_trigger_inputs = ["image_latents", None] - - @property - def description(self): - return ( - "Input step that standardize the inputs for the denoising step, e.g. make sure inputs have consistent batch size, and patchified. \n" - " This is an auto pipeline block that works for text2image/img2img tasks.\n" - + " - `FluxImg2ImgInputStep` (img2img) is used when `image_latents` is provided.\n" - + " - `FluxTextInputStep` (text2image) is used when `image_latents` are not provided.\n" - ) - - -# inputs: Flux Kontext - -FluxKontextBlocks = InsertableDict( - [ - ("set_resolution", FluxKontextSetResolutionStep()), - ("text_inputs", FluxTextInputStep()), - ("additional_inputs", FluxKontextInputsDynamicStep()), - ] -) - - -class FluxKontextInputStep(SequentialPipelineBlocks): - model_name = "flux-kontext" - block_classes = FluxKontextBlocks.values() - block_names = FluxKontextBlocks.keys() - - @property - def description(self): - return ( - "Input step that prepares the inputs for the both text2img and img2img denoising step. It:\n" - " - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents`).\n" - " - update height/width based `image_latents`, patchify `image_latents`." - ) - - -class FluxKontextAutoInputStep(AutoPipelineBlocks): - block_classes = [FluxKontextInputStep, FluxTextInputStep] - # block_classes = [FluxKontextInputStep] - block_names = ["img2img", "text2img"] - # block_names = ["img2img"] - block_trigger_inputs = ["image_latents", None] - # block_trigger_inputs = ["image_latents"] - - @property - def description(self): - return ( - "Input step that standardize the inputs for the denoising step, e.g. make sure inputs have consistent batch size, and patchified. \n" - " This is an auto pipeline block that works for text2image/img2img tasks.\n" - + " - `FluxKontextInputStep` (img2img) is used when `image_latents` is provided.\n" - + " - `FluxKontextInputStep` is also capable of handling text2image task when `image_latent` isn't present." - ) - - -class FluxCoreDenoiseStep(SequentialPipelineBlocks): - model_name = "flux" - block_classes = [FluxAutoInputStep, FluxAutoBeforeDenoiseStep, FluxAutoDenoiseStep] - block_names = ["input", "before_denoise", "denoise"] - - @property - def description(self): - return ( - "Core step that performs the denoising process. \n" - + " - `FluxAutoInputStep` (input) standardizes the inputs for the denoising step.\n" - + " - `FluxAutoBeforeDenoiseStep` (before_denoise) prepares the inputs for the denoising step.\n" - + " - `FluxAutoDenoiseStep` (denoise) iteratively denoises the latents.\n" - + "This step supports text-to-image and image-to-image tasks for Flux:\n" - + " - for image-to-image generation, you need to provide `image_latents`\n" - + " - for text-to-image generation, all you need to provide is prompt embeddings." - ) - - -class FluxKontextCoreDenoiseStep(SequentialPipelineBlocks): - model_name = "flux-kontext" - block_classes = [FluxKontextAutoInputStep, FluxKontextAutoBeforeDenoiseStep, FluxKontextAutoDenoiseStep] - block_names = ["input", "before_denoise", "denoise"] - - @property - def description(self): - return ( - "Core step that performs the denoising process. \n" - + " - `FluxKontextAutoInputStep` (input) standardizes the inputs for the denoising step.\n" - + " - `FluxKontextAutoBeforeDenoiseStep` (before_denoise) prepares the inputs for the denoising step.\n" - + " - `FluxKontextAutoDenoiseStep` (denoise) iteratively denoises the latents.\n" - + "This step supports text-to-image and image-to-image tasks for Flux:\n" - + " - for image-to-image generation, you need to provide `image_latents`\n" - + " - for text-to-image generation, all you need to provide is prompt embeddings." - ) - - -# Auto blocks (text2image and img2img) -AUTO_BLOCKS = InsertableDict( - [ - ("text_encoder", FluxTextEncoderStep()), - ("vae_encoder", FluxAutoVaeEncoderStep()), - ("denoise", FluxCoreDenoiseStep()), - ("decode", FluxDecodeStep()), - ] -) - -AUTO_BLOCKS_KONTEXT = InsertableDict( - [ - ("text_encoder", FluxTextEncoderStep()), - ("vae_encoder", FluxKontextAutoVaeEncoderStep()), - ("denoise", FluxKontextCoreDenoiseStep()), - ("decode", FluxDecodeStep()), - ] -) - - -class FluxAutoBlocks(SequentialPipelineBlocks): - model_name = "flux" - - block_classes = AUTO_BLOCKS.values() - block_names = AUTO_BLOCKS.keys() - - @property - def description(self): - return ( - "Auto Modular pipeline for text-to-image and image-to-image using Flux.\n" - + "- for text-to-image generation, all you need to provide is `prompt`\n" - + "- for image-to-image generation, you need to provide either `image` or `image_latents`" - ) - - -class FluxKontextAutoBlocks(FluxAutoBlocks): - model_name = "flux-kontext" - - block_classes = AUTO_BLOCKS_KONTEXT.values() - block_names = AUTO_BLOCKS_KONTEXT.keys() - - -TEXT2IMAGE_BLOCKS = InsertableDict( - [ - ("text_encoder", FluxTextEncoderStep()), - ("input", FluxTextInputStep()), - ("prepare_latents", FluxPrepareLatentsStep()), - ("set_timesteps", FluxSetTimestepsStep()), - ("prepare_rope_inputs", FluxRoPEInputsStep()), - ("denoise", FluxDenoiseStep()), - ("decode", FluxDecodeStep()), - ] -) - -IMAGE2IMAGE_BLOCKS = InsertableDict( - [ - ("text_encoder", FluxTextEncoderStep()), - ("vae_encoder", FluxVaeEncoderDynamicStep()), - ("input", FluxImg2ImgInputStep()), - ("prepare_latents", FluxPrepareLatentsStep()), - ("set_timesteps", FluxImg2ImgSetTimestepsStep()), - ("prepare_img2img_latents", FluxImg2ImgPrepareLatentsStep()), - ("prepare_rope_inputs", FluxRoPEInputsStep()), - ("denoise", FluxDenoiseStep()), - ("decode", FluxDecodeStep()), - ] -) - -FLUX_KONTEXT_BLOCKS = InsertableDict( - [ - ("text_encoder", FluxTextEncoderStep()), - ("vae_encoder", FluxVaeEncoderDynamicStep(sample_mode="argmax")), - ("input", FluxKontextInputStep()), - ("prepare_latents", FluxPrepareLatentsStep()), - ("set_timesteps", FluxSetTimestepsStep()), - ("prepare_rope_inputs", FluxKontextRoPEInputsStep()), - ("denoise", FluxKontextDenoiseStep()), - ("decode", FluxDecodeStep()), - ] -) - -ALL_BLOCKS = { - "text2image": TEXT2IMAGE_BLOCKS, - "img2img": IMAGE2IMAGE_BLOCKS, - "auto": AUTO_BLOCKS, - "auto_kontext": AUTO_BLOCKS_KONTEXT, - "kontext": FLUX_KONTEXT_BLOCKS, -} diff --git a/src/diffusers/modular_pipelines/flux/modular_blocks_flux.py b/src/diffusers/modular_pipelines/flux/modular_blocks_flux.py new file mode 100644 index 000000000000..f2e78e933448 --- /dev/null +++ b/src/diffusers/modular_pipelines/flux/modular_blocks_flux.py @@ -0,0 +1,586 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...utils import logging +from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import ( + FluxImg2ImgPrepareLatentsStep, + FluxImg2ImgSetTimestepsStep, + FluxPrepareLatentsStep, + FluxRoPEInputsStep, + FluxSetTimestepsStep, +) +from .decoders import FluxDecodeStep +from .denoise import FluxDenoiseStep +from .encoders import ( + FluxProcessImagesInputStep, + FluxTextEncoderStep, + FluxVaeEncoderStep, +) +from .inputs import ( + FluxAdditionalInputsStep, + FluxTextInputStep, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# vae encoder (run before before_denoise) + + +# auto_docstring +class FluxImg2ImgVaeEncoderStep(SequentialPipelineBlocks): + """ + Vae encoder step that preprocess andencode the image inputs into their latent representations. + + Components: + image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`) + + Inputs: + resized_image (`None`, *optional*): + TODO: Add description. + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + processed_image (`None`): + TODO: Add description. + image_latents (`Tensor`): + The latents representing the reference image + """ + + model_name = "flux" + + block_classes = [FluxProcessImagesInputStep(), FluxVaeEncoderStep()] + block_names = ["preprocess", "encode"] + + @property + def description(self) -> str: + return "Vae encoder step that preprocess andencode the image inputs into their latent representations." + + +# auto_docstring +class FluxAutoVaeEncoderStep(AutoPipelineBlocks): + """ + Vae encoder step that encode the image inputs into their latent representations. + This is an auto pipeline block that works for img2img tasks. + - `FluxImg2ImgVaeEncoderStep` (img2img) is used when only `image` is provided. - if `image` is not provided, + step will be skipped. + + Components: + image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`) + + Inputs: + resized_image (`None`, *optional*): + TODO: Add description. + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + processed_image (`None`): + TODO: Add description. + image_latents (`Tensor`): + The latents representing the reference image + """ + + model_name = "flux" + block_classes = [FluxImg2ImgVaeEncoderStep] + block_names = ["img2img"] + block_trigger_inputs = ["image"] + + @property + def description(self): + return ( + "Vae encoder step that encode the image inputs into their latent representations.\n" + + "This is an auto pipeline block that works for img2img tasks.\n" + + " - `FluxImg2ImgVaeEncoderStep` (img2img) is used when only `image` is provided." + + " - if `image` is not provided, step will be skipped." + ) + + +# before_denoise: text2img +# auto_docstring +class FluxBeforeDenoiseStep(SequentialPipelineBlocks): + """ + Before denoise step that prepares the inputs for the denoise step in text-to-image generation. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + num_images_per_prompt (`int`, *optional*, defaults to 1): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + batch_size (`int`): + Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`. + Can be generated in input step. + dtype (`dtype`, *optional*): + The dtype of the model inputs + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + guidance_scale (`None`, *optional*, defaults to 3.5): + TODO: Add description. + prompt_embeds (`None`, *optional*): + TODO: Add description. + + Outputs: + latents (`Tensor`): + The initial latents to use for the denoising process + timesteps (`Tensor`): + The timesteps to use for inference + num_inference_steps (`int`): + The number of denoising steps to perform at inference time + guidance (`Tensor`): + Optional guidance to be used. + txt_ids (`list`): + The sequence lengths of the prompt embeds, used for RoPE calculation. + img_ids (`list`): + The sequence lengths of the image latents, used for RoPE calculation. + """ + + model_name = "flux" + block_classes = [FluxPrepareLatentsStep(), FluxSetTimestepsStep(), FluxRoPEInputsStep()] + block_names = ["prepare_latents", "set_timesteps", "prepare_rope_inputs"] + + @property + def description(self): + return "Before denoise step that prepares the inputs for the denoise step in text-to-image generation." + + +# before_denoise: img2img +# auto_docstring +class FluxImg2ImgBeforeDenoiseStep(SequentialPipelineBlocks): + """ + Before denoise step that prepare the inputs for the denoise step for img2img task. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + num_images_per_prompt (`int`, *optional*, defaults to 1): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + batch_size (`int`): + Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`. + Can be generated in input step. + dtype (`dtype`, *optional*): + The dtype of the model inputs + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + strength (`None`, *optional*, defaults to 0.6): + TODO: Add description. + guidance_scale (`None`, *optional*, defaults to 3.5): + TODO: Add description. + image_latents (`Tensor`): + The image latents to use for the denoising process. Can be generated in vae encoder and packed in input + step. + prompt_embeds (`None`, *optional*): + TODO: Add description. + + Outputs: + latents (`Tensor`): + The initial latents to use for the denoising process + timesteps (`Tensor`): + The timesteps to use for inference + num_inference_steps (`int`): + The number of denoising steps to perform at inference time + guidance (`Tensor`): + Optional guidance to be used. + initial_noise (`Tensor`): + The initial random noised used for inpainting denoising. + txt_ids (`list`): + The sequence lengths of the prompt embeds, used for RoPE calculation. + img_ids (`list`): + The sequence lengths of the image latents, used for RoPE calculation. + """ + + model_name = "flux" + block_classes = [ + FluxPrepareLatentsStep(), + FluxImg2ImgSetTimestepsStep(), + FluxImg2ImgPrepareLatentsStep(), + FluxRoPEInputsStep(), + ] + block_names = ["prepare_latents", "set_timesteps", "prepare_img2img_latents", "prepare_rope_inputs"] + + @property + def description(self): + return "Before denoise step that prepare the inputs for the denoise step for img2img task." + + +# before_denoise: all task (text2img, img2img) +# auto_docstring +class FluxAutoBeforeDenoiseStep(AutoPipelineBlocks): + """ + Before denoise step that prepare the inputs for the denoise step. + This is an auto pipeline block that works for text2image. + - `FluxBeforeDenoiseStep` (text2image) is used. + - `FluxImg2ImgBeforeDenoiseStep` (img2img) is used when only `image_latents` is provided. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + height (`int`): + TODO: Add description. + width (`int`): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + num_images_per_prompt (`int`, *optional*, defaults to 1): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + batch_size (`int`): + Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`. + Can be generated in input step. + dtype (`dtype`, *optional*): + The dtype of the model inputs + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + strength (`None`, *optional*, defaults to 0.6): + TODO: Add description. + guidance_scale (`None`, *optional*, defaults to 3.5): + TODO: Add description. + image_latents (`Tensor`, *optional*): + The image latents to use for the denoising process. Can be generated in vae encoder and packed in input + step. + prompt_embeds (`None`, *optional*): + TODO: Add description. + + Outputs: + latents (`Tensor`): + The initial latents to use for the denoising process + timesteps (`Tensor`): + The timesteps to use for inference + num_inference_steps (`int`): + The number of denoising steps to perform at inference time + guidance (`Tensor`): + Optional guidance to be used. + initial_noise (`Tensor`): + The initial random noised used for inpainting denoising. + txt_ids (`list`): + The sequence lengths of the prompt embeds, used for RoPE calculation. + img_ids (`list`): + The sequence lengths of the image latents, used for RoPE calculation. + """ + + model_name = "flux" + block_classes = [FluxImg2ImgBeforeDenoiseStep, FluxBeforeDenoiseStep] + block_names = ["img2img", "text2image"] + block_trigger_inputs = ["image_latents", None] + + @property + def description(self): + return ( + "Before denoise step that prepare the inputs for the denoise step.\n" + + "This is an auto pipeline block that works for text2image.\n" + + " - `FluxBeforeDenoiseStep` (text2image) is used.\n" + + " - `FluxImg2ImgBeforeDenoiseStep` (img2img) is used when only `image_latents` is provided.\n" + ) + + +# inputs: text2image/img2img + + +# auto_docstring +class FluxImg2ImgInputStep(SequentialPipelineBlocks): + """ + Input step that prepares the inputs for the img2img denoising step. It: + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + pooled_prompt_embeds (`Tensor`, *optional*): + Pre-generated pooled text embeddings. Can be generated from text_encoder step. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + image_latents (`None`, *optional*): + TODO: Add description. + + Outputs: + batch_size (`int`): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt + dtype (`dtype`): + Data type of model tensor inputs (determined by `prompt_embeds`) + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation + pooled_prompt_embeds (`Tensor`): + pooled text embeddings used to guide the image generation + image_height (`int`): + The height of the image latents + image_width (`int`): + The width of the image latents + """ + + model_name = "flux" + block_classes = [FluxTextInputStep(), FluxAdditionalInputsStep()] + block_names = ["text_inputs", "additional_inputs"] + + @property + def description(self): + return "Input step that prepares the inputs for the img2img denoising step. It:\n" + " - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents`).\n" + " - update height/width based `image_latents`, patchify `image_latents`." + + +# auto_docstring +class FluxAutoInputStep(AutoPipelineBlocks): + """ + Input step that standardize the inputs for the denoising step, e.g. make sure inputs have consistent batch size, + and patchified. + This is an auto pipeline block that works for text2image/img2img tasks. + - `FluxImg2ImgInputStep` (img2img) is used when `image_latents` is provided. + - `FluxTextInputStep` (text2image) is used when `image_latents` are not provided. + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + pooled_prompt_embeds (`Tensor`, *optional*): + Pre-generated pooled text embeddings. Can be generated from text_encoder step. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + image_latents (`None`, *optional*): + TODO: Add description. + + Outputs: + batch_size (`int`): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt + dtype (`dtype`): + Data type of model tensor inputs (determined by `prompt_embeds`) + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation + pooled_prompt_embeds (`Tensor`): + pooled text embeddings used to guide the image generation + image_height (`int`): + The height of the image latents + image_width (`int`): + The width of the image latents + """ + + model_name = "flux" + + block_classes = [FluxImg2ImgInputStep, FluxTextInputStep] + block_names = ["img2img", "text2image"] + block_trigger_inputs = ["image_latents", None] + + @property + def description(self): + return ( + "Input step that standardize the inputs for the denoising step, e.g. make sure inputs have consistent batch size, and patchified. \n" + " This is an auto pipeline block that works for text2image/img2img tasks.\n" + + " - `FluxImg2ImgInputStep` (img2img) is used when `image_latents` is provided.\n" + + " - `FluxTextInputStep` (text2image) is used when `image_latents` are not provided.\n" + ) + + +# auto_docstring +class FluxCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core step that performs the denoising process for Flux. + This step supports text-to-image and image-to-image tasks for Flux: + - for image-to-image generation, you need to provide `image_latents` + - for text-to-image generation, all you need to provide is prompt embeddings. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`FluxTransformer2DModel`) + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + pooled_prompt_embeds (`Tensor`, *optional*): + Pre-generated pooled text embeddings. Can be generated from text_encoder step. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + image_latents (`None`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + strength (`None`, *optional*, defaults to 0.6): + TODO: Add description. + guidance_scale (`None`, *optional*, defaults to 3.5): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "flux" + block_classes = [FluxAutoInputStep, FluxAutoBeforeDenoiseStep, FluxDenoiseStep] + block_names = ["input", "before_denoise", "denoise"] + + @property + def description(self): + return ( + "Core step that performs the denoising process for Flux.\n" + + "This step supports text-to-image and image-to-image tasks for Flux:\n" + + " - for image-to-image generation, you need to provide `image_latents`\n" + + " - for text-to-image generation, all you need to provide is prompt embeddings." + ) + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +# Auto blocks (text2image and img2img) +AUTO_BLOCKS = InsertableDict( + [ + ("text_encoder", FluxTextEncoderStep()), + ("vae_encoder", FluxAutoVaeEncoderStep()), + ("denoise", FluxCoreDenoiseStep()), + ("decode", FluxDecodeStep()), + ] +) + + +# auto_docstring +class FluxAutoBlocks(SequentialPipelineBlocks): + """ + Auto Modular pipeline for text-to-image and image-to-image using Flux. + + Supported workflows: + - `text2image`: requires `prompt` + - `image2image`: requires `image`, `prompt` + + Components: + text_encoder (`CLIPTextModel`) tokenizer (`CLIPTokenizer`) text_encoder_2 (`T5EncoderModel`) tokenizer_2 + (`T5TokenizerFast`) image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`) scheduler + (`FlowMatchEulerDiscreteScheduler`) transformer (`FluxTransformer2DModel`) + + Inputs: + prompt (`None`, *optional*): + TODO: Add description. + prompt_2 (`None`, *optional*): + TODO: Add description. + max_sequence_length (`int`, *optional*, defaults to 512): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + resized_image (`None`, *optional*): + TODO: Add description. + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + image_latents (`None`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + strength (`None`, *optional*, defaults to 0.6): + TODO: Add description. + guidance_scale (`None`, *optional*, defaults to 3.5): + TODO: Add description. + output_type (`None`, *optional*, defaults to pil): + TODO: Add description. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "flux" + + block_classes = AUTO_BLOCKS.values() + block_names = AUTO_BLOCKS.keys() + + _workflow_map = { + "text2image": {"prompt": True}, + "image2image": {"image": True, "prompt": True}, + } + + @property + def description(self): + return "Auto Modular pipeline for text-to-image and image-to-image using Flux." + + @property + def outputs(self): + return [OutputParam.template("images")] diff --git a/src/diffusers/modular_pipelines/flux/modular_blocks_flux_kontext.py b/src/diffusers/modular_pipelines/flux/modular_blocks_flux_kontext.py new file mode 100644 index 000000000000..b5a5dbf78c0e --- /dev/null +++ b/src/diffusers/modular_pipelines/flux/modular_blocks_flux_kontext.py @@ -0,0 +1,585 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...utils import logging +from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import ( + FluxKontextRoPEInputsStep, + FluxPrepareLatentsStep, + FluxRoPEInputsStep, + FluxSetTimestepsStep, +) +from .decoders import FluxDecodeStep +from .denoise import FluxKontextDenoiseStep +from .encoders import ( + FluxKontextProcessImagesInputStep, + FluxTextEncoderStep, + FluxVaeEncoderStep, +) +from .inputs import ( + FluxKontextAdditionalInputsStep, + FluxKontextSetResolutionStep, + FluxTextInputStep, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# Flux Kontext vae encoder (run before before_denoise) +# auto_docstring +class FluxKontextVaeEncoderStep(SequentialPipelineBlocks): + """ + Vae encoder step that preprocess andencode the image inputs into their latent representations. + + Components: + image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`) + + Inputs: + image (`None`, *optional*): + TODO: Add description. + _auto_resize (`bool`, *optional*, defaults to True): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + processed_image (`None`): + TODO: Add description. + image_latents (`Tensor`): + The latents representing the reference image + """ + + model_name = "flux-kontext" + + block_classes = [FluxKontextProcessImagesInputStep(), FluxVaeEncoderStep(sample_mode="argmax")] + block_names = ["preprocess", "encode"] + + @property + def description(self) -> str: + return "Vae encoder step that preprocess andencode the image inputs into their latent representations." + + +# auto_docstring +class FluxKontextAutoVaeEncoderStep(AutoPipelineBlocks): + """ + Vae encoder step that encode the image inputs into their latent representations. + This is an auto pipeline block that works for image-conditioned tasks. + - `FluxKontextVaeEncoderStep` (image_conditioned) is used when only `image` is provided. - if `image` is not + provided, step will be skipped. + + Components: + image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`) + + Inputs: + image (`None`, *optional*): + TODO: Add description. + _auto_resize (`bool`, *optional*, defaults to True): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + processed_image (`None`): + TODO: Add description. + image_latents (`Tensor`): + The latents representing the reference image + """ + + model_name = "flux-kontext" + + block_classes = [FluxKontextVaeEncoderStep] + block_names = ["image_conditioned"] + block_trigger_inputs = ["image"] + + @property + def description(self): + return ( + "Vae encoder step that encode the image inputs into their latent representations.\n" + + "This is an auto pipeline block that works for image-conditioned tasks.\n" + + " - `FluxKontextVaeEncoderStep` (image_conditioned) is used when only `image` is provided." + + " - if `image` is not provided, step will be skipped." + ) + + +# before_denoise: text2img +# auto_docstring +class FluxKontextBeforeDenoiseStep(SequentialPipelineBlocks): + """ + Before denoise step that prepares the inputs for the denoise step for Flux Kontext + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + num_images_per_prompt (`int`, *optional*, defaults to 1): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + batch_size (`int`): + Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`. + Can be generated in input step. + dtype (`dtype`, *optional*): + The dtype of the model inputs + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + guidance_scale (`None`, *optional*, defaults to 3.5): + TODO: Add description. + prompt_embeds (`None`, *optional*): + TODO: Add description. + + Outputs: + latents (`Tensor`): + The initial latents to use for the denoising process + timesteps (`Tensor`): + The timesteps to use for inference + num_inference_steps (`int`): + The number of denoising steps to perform at inference time + guidance (`Tensor`): + Optional guidance to be used. + txt_ids (`list`): + The sequence lengths of the prompt embeds, used for RoPE calculation. + img_ids (`list`): + The sequence lengths of the image latents, used for RoPE calculation. + """ + + model_name = "flux-kontext" + + block_classes = [FluxPrepareLatentsStep(), FluxSetTimestepsStep(), FluxRoPEInputsStep()] + block_names = ["prepare_latents", "set_timesteps", "prepare_rope_inputs"] + + @property + def description(self): + return "Before denoise step that prepares the inputs for the denoise step for Flux Kontext\n" + "for text-to-image tasks." + + +# before_denoise: image-conditioned +# auto_docstring +class FluxKontextImageConditionedBeforeDenoiseStep(SequentialPipelineBlocks): + """ + Before denoise step that prepare the inputs for the denoise step for Flux Kontext + for image-conditioned tasks. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + num_images_per_prompt (`int`, *optional*, defaults to 1): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + batch_size (`int`): + Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`. + Can be generated in input step. + dtype (`dtype`, *optional*): + The dtype of the model inputs + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + guidance_scale (`None`, *optional*, defaults to 3.5): + TODO: Add description. + image_height (`None`, *optional*): + TODO: Add description. + image_width (`None`, *optional*): + TODO: Add description. + prompt_embeds (`None`, *optional*): + TODO: Add description. + + Outputs: + latents (`Tensor`): + The initial latents to use for the denoising process + timesteps (`Tensor`): + The timesteps to use for inference + num_inference_steps (`int`): + The number of denoising steps to perform at inference time + guidance (`Tensor`): + Optional guidance to be used. + txt_ids (`list`): + The sequence lengths of the prompt embeds, used for RoPE calculation. + img_ids (`list`): + The sequence lengths of the image latents, used for RoPE calculation. + """ + + model_name = "flux-kontext" + + block_classes = [FluxPrepareLatentsStep(), FluxSetTimestepsStep(), FluxKontextRoPEInputsStep()] + block_names = ["prepare_latents", "set_timesteps", "prepare_rope_inputs"] + + @property + def description(self): + return ( + "Before denoise step that prepare the inputs for the denoise step for Flux Kontext\n" + "for image-conditioned tasks." + ) + + +# auto_docstring +class FluxKontextAutoBeforeDenoiseStep(AutoPipelineBlocks): + """ + Before denoise step that prepare the inputs for the denoise step. + This is an auto pipeline block that works for text2image. + - `FluxKontextBeforeDenoiseStep` (text2image) is used. + - `FluxKontextImageConditionedBeforeDenoiseStep` (image_conditioned) is used when only `image_latents` is + provided. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + num_images_per_prompt (`int`, *optional*, defaults to 1): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + batch_size (`int`): + Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`. + Can be generated in input step. + dtype (`dtype`, *optional*): + The dtype of the model inputs + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + guidance_scale (`None`, *optional*, defaults to 3.5): + TODO: Add description. + image_height (`None`, *optional*): + TODO: Add description. + image_width (`None`, *optional*): + TODO: Add description. + prompt_embeds (`None`, *optional*): + TODO: Add description. + + Outputs: + latents (`Tensor`): + The initial latents to use for the denoising process + timesteps (`Tensor`): + The timesteps to use for inference + num_inference_steps (`int`): + The number of denoising steps to perform at inference time + guidance (`Tensor`): + Optional guidance to be used. + txt_ids (`list`): + The sequence lengths of the prompt embeds, used for RoPE calculation. + img_ids (`list`): + The sequence lengths of the image latents, used for RoPE calculation. + """ + + model_name = "flux-kontext" + + block_classes = [FluxKontextImageConditionedBeforeDenoiseStep, FluxKontextBeforeDenoiseStep] + block_names = ["image_conditioned", "text2image"] + block_trigger_inputs = ["image_latents", None] + + @property + def description(self): + return ( + "Before denoise step that prepare the inputs for the denoise step.\n" + + "This is an auto pipeline block that works for text2image.\n" + + " - `FluxKontextBeforeDenoiseStep` (text2image) is used.\n" + + " - `FluxKontextImageConditionedBeforeDenoiseStep` (image_conditioned) is used when only `image_latents` is provided.\n" + ) + + +# inputs: Flux Kontext +# auto_docstring +class FluxKontextInputStep(SequentialPipelineBlocks): + """ + Input step that prepares the inputs for the both text2img and img2img denoising step. It: + - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents`). + - update height/width based `image_latents`, patchify `image_latents`. + + Inputs: + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + max_area (`int`, *optional*, defaults to 1048576): + TODO: Add description. + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + pooled_prompt_embeds (`Tensor`, *optional*): + Pre-generated pooled text embeddings. Can be generated from text_encoder step. + image_latents (`None`, *optional*): + TODO: Add description. + + Outputs: + height (`int`): + The height of the initial noisy latents + width (`int`): + The width of the initial noisy latents + batch_size (`int`): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt + dtype (`dtype`): + Data type of model tensor inputs (determined by `prompt_embeds`) + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation + pooled_prompt_embeds (`Tensor`): + pooled text embeddings used to guide the image generation + image_height (`int`): + The height of the image latents + image_width (`int`): + The width of the image latents + """ + + model_name = "flux-kontext" + block_classes = [FluxKontextSetResolutionStep(), FluxTextInputStep(), FluxKontextAdditionalInputsStep()] + block_names = ["set_resolution", "text_inputs", "additional_inputs"] + + @property + def description(self): + return ( + "Input step that prepares the inputs for the both text2img and img2img denoising step. It:\n" + " - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents`).\n" + " - update height/width based `image_latents`, patchify `image_latents`." + ) + + +# auto_docstring +class FluxKontextAutoInputStep(AutoPipelineBlocks): + """ + Input step that standardize the inputs for the denoising step, e.g. make sure inputs have consistent batch size, + and patchified. + This is an auto pipeline block that works for text2image/img2img tasks. + - `FluxKontextInputStep` (image_conditioned) is used when `image_latents` is provided. + - `FluxKontextInputStep` is also capable of handling text2image task when `image_latent` isn't present. + + Inputs: + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + max_area (`int`, *optional*, defaults to 1048576): + TODO: Add description. + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + pooled_prompt_embeds (`Tensor`, *optional*): + Pre-generated pooled text embeddings. Can be generated from text_encoder step. + image_latents (`None`, *optional*): + TODO: Add description. + + Outputs: + height (`int`): + The height of the initial noisy latents + width (`int`): + The width of the initial noisy latents + batch_size (`int`): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt + dtype (`dtype`): + Data type of model tensor inputs (determined by `prompt_embeds`) + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation + pooled_prompt_embeds (`Tensor`): + pooled text embeddings used to guide the image generation + image_height (`int`): + The height of the image latents + image_width (`int`): + The width of the image latents + """ + + model_name = "flux-kontext" + block_classes = [FluxKontextInputStep, FluxTextInputStep] + block_names = ["image_conditioned", "text2image"] + block_trigger_inputs = ["image_latents", None] + + @property + def description(self): + return ( + "Input step that standardize the inputs for the denoising step, e.g. make sure inputs have consistent batch size, and patchified. \n" + " This is an auto pipeline block that works for text2image/img2img tasks.\n" + + " - `FluxKontextInputStep` (image_conditioned) is used when `image_latents` is provided.\n" + + " - `FluxKontextInputStep` is also capable of handling text2image task when `image_latent` isn't present." + ) + + +# auto_docstring +class FluxKontextCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core step that performs the denoising process for Flux Kontext. + This step supports text-to-image and image-conditioned tasks for Flux Kontext: + - for image-conditioned generation, you need to provide `image_latents` + - for text-to-image generation, all you need to provide is prompt embeddings. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`FluxTransformer2DModel`) + + Inputs: + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + max_area (`int`, *optional*, defaults to 1048576): + TODO: Add description. + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + pooled_prompt_embeds (`Tensor`, *optional*): + Pre-generated pooled text embeddings. Can be generated from text_encoder step. + image_latents (`None`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + guidance_scale (`None`, *optional*, defaults to 3.5): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "flux-kontext" + block_classes = [FluxKontextAutoInputStep, FluxKontextAutoBeforeDenoiseStep, FluxKontextDenoiseStep] + block_names = ["input", "before_denoise", "denoise"] + + @property + def description(self): + return ( + "Core step that performs the denoising process for Flux Kontext.\n" + + "This step supports text-to-image and image-conditioned tasks for Flux Kontext:\n" + + " - for image-conditioned generation, you need to provide `image_latents`\n" + + " - for text-to-image generation, all you need to provide is prompt embeddings." + ) + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +AUTO_BLOCKS_KONTEXT = InsertableDict( + [ + ("text_encoder", FluxTextEncoderStep()), + ("vae_encoder", FluxKontextAutoVaeEncoderStep()), + ("denoise", FluxKontextCoreDenoiseStep()), + ("decode", FluxDecodeStep()), + ] +) + + +# auto_docstring +class FluxKontextAutoBlocks(SequentialPipelineBlocks): + """ + Modular pipeline for image-to-image using Flux Kontext. + + Supported workflows: + - `image_conditioned`: requires `image`, `prompt` + - `text2image`: requires `prompt` + + Components: + text_encoder (`CLIPTextModel`) tokenizer (`CLIPTokenizer`) text_encoder_2 (`T5EncoderModel`) tokenizer_2 + (`T5TokenizerFast`) image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`) scheduler + (`FlowMatchEulerDiscreteScheduler`) transformer (`FluxTransformer2DModel`) + + Inputs: + prompt (`None`, *optional*): + TODO: Add description. + prompt_2 (`None`, *optional*): + TODO: Add description. + max_sequence_length (`int`, *optional*, defaults to 512): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + image (`None`, *optional*): + TODO: Add description. + _auto_resize (`bool`, *optional*, defaults to True): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + max_area (`int`, *optional*, defaults to 1048576): + TODO: Add description. + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + image_latents (`None`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + guidance_scale (`None`, *optional*, defaults to 3.5): + TODO: Add description. + output_type (`None`, *optional*, defaults to pil): + TODO: Add description. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "flux-kontext" + + block_classes = AUTO_BLOCKS_KONTEXT.values() + block_names = AUTO_BLOCKS_KONTEXT.keys() + _workflow_map = { + "image_conditioned": {"image": True, "prompt": True}, + "text2image": {"prompt": True}, + } + + @property + def description(self): + return "Modular pipeline for image-to-image using Flux Kontext." + + @property + def outputs(self): + return [OutputParam.template("images")] diff --git a/src/diffusers/modular_pipelines/flux2/__init__.py b/src/diffusers/modular_pipelines/flux2/__init__.py index 21a41c1fe941..d7cc8badcaf7 100644 --- a/src/diffusers/modular_pipelines/flux2/__init__.py +++ b/src/diffusers/modular_pipelines/flux2/__init__.py @@ -21,40 +21,15 @@ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) else: - _import_structure["encoders"] = [ - "Flux2TextEncoderStep", - "Flux2RemoteTextEncoderStep", - "Flux2VaeEncoderStep", + _import_structure["encoders"] = ["Flux2RemoteTextEncoderStep"] + _import_structure["modular_blocks_flux2"] = ["Flux2AutoBlocks"] + _import_structure["modular_blocks_flux2_klein"] = ["Flux2KleinAutoBlocks"] + _import_structure["modular_blocks_flux2_klein_base"] = ["Flux2KleinBaseAutoBlocks"] + _import_structure["modular_pipeline"] = [ + "Flux2KleinBaseModularPipeline", + "Flux2KleinModularPipeline", + "Flux2ModularPipeline", ] - _import_structure["before_denoise"] = [ - "Flux2SetTimestepsStep", - "Flux2PrepareLatentsStep", - "Flux2RoPEInputsStep", - "Flux2PrepareImageLatentsStep", - ] - _import_structure["denoise"] = [ - "Flux2LoopDenoiser", - "Flux2LoopAfterDenoiser", - "Flux2DenoiseLoopWrapper", - "Flux2DenoiseStep", - ] - _import_structure["decoders"] = ["Flux2DecodeStep"] - _import_structure["inputs"] = [ - "Flux2ProcessImagesInputStep", - "Flux2TextInputStep", - ] - _import_structure["modular_blocks"] = [ - "ALL_BLOCKS", - "AUTO_BLOCKS", - "REMOTE_AUTO_BLOCKS", - "TEXT2IMAGE_BLOCKS", - "IMAGE_CONDITIONED_BLOCKS", - "Flux2AutoBlocks", - "Flux2AutoVaeEncoderStep", - "Flux2BeforeDenoiseStep", - "Flux2VaeEncoderSequentialStep", - ] - _import_structure["modular_pipeline"] = ["Flux2ModularPipeline"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: @@ -63,40 +38,11 @@ except OptionalDependencyNotAvailable: from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 else: - from .before_denoise import ( - Flux2PrepareImageLatentsStep, - Flux2PrepareLatentsStep, - Flux2RoPEInputsStep, - Flux2SetTimestepsStep, - ) - from .decoders import Flux2DecodeStep - from .denoise import ( - Flux2DenoiseLoopWrapper, - Flux2DenoiseStep, - Flux2LoopAfterDenoiser, - Flux2LoopDenoiser, - ) - from .encoders import ( - Flux2RemoteTextEncoderStep, - Flux2TextEncoderStep, - Flux2VaeEncoderStep, - ) - from .inputs import ( - Flux2ProcessImagesInputStep, - Flux2TextInputStep, - ) - from .modular_blocks import ( - ALL_BLOCKS, - AUTO_BLOCKS, - IMAGE_CONDITIONED_BLOCKS, - REMOTE_AUTO_BLOCKS, - TEXT2IMAGE_BLOCKS, - Flux2AutoBlocks, - Flux2AutoVaeEncoderStep, - Flux2BeforeDenoiseStep, - Flux2VaeEncoderSequentialStep, - ) - from .modular_pipeline import Flux2ModularPipeline + from .encoders import Flux2RemoteTextEncoderStep + from .modular_blocks_flux2 import Flux2AutoBlocks + from .modular_blocks_flux2_klein import Flux2KleinAutoBlocks + from .modular_blocks_flux2_klein_base import Flux2KleinBaseAutoBlocks + from .modular_pipeline import Flux2KleinBaseModularPipeline, Flux2KleinModularPipeline, Flux2ModularPipeline else: import sys diff --git a/src/diffusers/modular_pipelines/flux2/before_denoise.py b/src/diffusers/modular_pipelines/flux2/before_denoise.py index 42624688adfa..6b1b3bd96324 100644 --- a/src/diffusers/modular_pipelines/flux2/before_denoise.py +++ b/src/diffusers/modular_pipelines/flux2/before_denoise.py @@ -13,7 +13,6 @@ # limitations under the License. import inspect -from typing import List, Optional, Union import numpy as np import torch @@ -52,10 +51,10 @@ def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float: # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -70,15 +69,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -113,7 +112,7 @@ class Flux2SetTimestepsStep(ModularPipelineBlocks): model_name = "flux2" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), ComponentSpec("transformer", Flux2Transformer2DModel), @@ -124,26 +123,18 @@ def description(self) -> str: return "Step that sets the scheduler's timesteps for Flux2 inference using empirical mu calculation" @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("num_inference_steps", default=50), InputParam("timesteps"), InputParam("sigmas"), - InputParam("guidance_scale", default=4.0), InputParam("latents", type_hint=torch.Tensor), - InputParam("num_images_per_prompt", default=1), InputParam("height", type_hint=int), InputParam("width", type_hint=int), - InputParam( - "batch_size", - required=True, - type_hint=int, - description="Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`.", - ), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference"), OutputParam( @@ -151,13 +142,12 @@ def intermediate_outputs(self) -> List[OutputParam]: type_hint=int, description="The number of denoising steps to perform at inference time", ), - OutputParam("guidance", type_hint=torch.Tensor, description="Guidance scale tensor"), ] @torch.no_grad() def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - block_state.device = components._execution_device + device = components._execution_device scheduler = components.scheduler @@ -183,7 +173,7 @@ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> Pi timesteps, num_inference_steps = retrieve_timesteps( scheduler, num_inference_steps, - block_state.device, + device, timesteps=timesteps, sigmas=sigmas, mu=mu, @@ -191,11 +181,6 @@ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> Pi block_state.timesteps = timesteps block_state.num_inference_steps = num_inference_steps - batch_size = block_state.batch_size * block_state.num_images_per_prompt - guidance = torch.full([1], block_state.guidance_scale, device=block_state.device, dtype=torch.float32) - guidance = guidance.expand(batch_size) - block_state.guidance = guidance - components.scheduler.set_begin_index(0) self.set_block_state(state, block_state) @@ -206,7 +191,7 @@ class Flux2PrepareLatentsStep(ModularPipelineBlocks): model_name = "flux2" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [] @property @@ -214,11 +199,11 @@ def description(self) -> str: return "Prepare latents step that prepares the initial noise latents for Flux2 text-to-image generation" @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("height", type_hint=int), InputParam("width", type_hint=int), - InputParam("latents", type_hint=Optional[torch.Tensor]), + InputParam("latents", type_hint=torch.Tensor | None), InputParam("num_images_per_prompt", type_hint=int, default=1), InputParam("generator"), InputParam( @@ -231,7 +216,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process" @@ -350,14 +335,13 @@ def description(self) -> str: return "Step that prepares the 4D RoPE position IDs for Flux2 denoising. Should be placed after text encoder and latent preparation steps." @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam(name="prompt_embeds", required=True), - InputParam(name="latent_ids"), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( name="txt_ids", @@ -365,16 +349,71 @@ def intermediate_outputs(self) -> List[OutputParam]: type_hint=torch.Tensor, description="4D position IDs (T, H, W, L) for text tokens, used for RoPE calculation.", ), + ] + + @staticmethod + def _prepare_text_ids(x: torch.Tensor, t_coord: torch.Tensor | None = None): + """Prepare 4D position IDs for text tokens.""" + B, L, _ = x.shape + out_ids = [] + + for i in range(B): + t = torch.arange(1) if t_coord is None else t_coord[i] + h = torch.arange(1) + w = torch.arange(1) + seq_l = torch.arange(L) + + coords = torch.cartesian_prod(t, h, w, seq_l) + out_ids.append(coords) + + return torch.stack(out_ids) + + def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + prompt_embeds = block_state.prompt_embeds + device = prompt_embeds.device + + block_state.txt_ids = self._prepare_text_ids(prompt_embeds) + block_state.txt_ids = block_state.txt_ids.to(device) + + self.set_block_state(state, block_state) + return components, state + + +class Flux2KleinBaseRoPEInputsStep(ModularPipelineBlocks): + model_name = "flux2-klein" + + @property + def description(self) -> str: + return "Step that prepares the 4D RoPE position IDs for Flux2-Klein base model denoising. Should be placed after text encoder and latent preparation steps." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam(name="prompt_embeds", required=True), + InputParam(name="negative_prompt_embeds", required=False), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ OutputParam( - name="latent_ids", + name="txt_ids", kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, - description="4D position IDs (T, H, W, L) for image latents, used for RoPE calculation.", + description="4D position IDs (T, H, W, L) for text tokens, used for RoPE calculation.", + ), + OutputParam( + name="negative_txt_ids", + kwargs_type="denoiser_input_fields", + type_hint=torch.Tensor, + description="4D position IDs (T, H, W, L) for negative text tokens, used for RoPE calculation.", ), ] @staticmethod - def _prepare_text_ids(x: torch.Tensor, t_coord: Optional[torch.Tensor] = None): + def _prepare_text_ids(x: torch.Tensor, t_coord: torch.Tensor | None = None): """Prepare 4D position IDs for text tokens.""" B, L, _ = x.shape out_ids = [] @@ -399,6 +438,11 @@ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> Pi block_state.txt_ids = self._prepare_text_ids(prompt_embeds) block_state.txt_ids = block_state.txt_ids.to(device) + block_state.negative_txt_ids = None + if block_state.negative_prompt_embeds is not None: + block_state.negative_txt_ids = self._prepare_text_ids(block_state.negative_prompt_embeds) + block_state.negative_txt_ids = block_state.negative_txt_ids.to(device) + self.set_block_state(state, block_state) return components, state @@ -411,15 +455,15 @@ def description(self) -> str: return "Step that prepares image latents and their position IDs for Flux2 image conditioning." @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ - InputParam("image_latents", type_hint=List[torch.Tensor]), + InputParam("image_latents", type_hint=list[torch.Tensor]), InputParam("batch_size", required=True, type_hint=int), InputParam("num_images_per_prompt", default=1, type_hint=int), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "image_latents", @@ -434,7 +478,7 @@ def intermediate_outputs(self) -> List[OutputParam]: ] @staticmethod - def _prepare_image_ids(image_latents: List[torch.Tensor], scale: int = 10): + def _prepare_image_ids(image_latents: list[torch.Tensor], scale: int = 10): """ Generates 4D time-space coordinates (T, H, W, L) for a sequence of image latents. @@ -506,3 +550,42 @@ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> Pi self.set_block_state(state, block_state) return components, state + + +class Flux2PrepareGuidanceStep(ModularPipelineBlocks): + model_name = "flux2" + + @property + def description(self) -> str: + return "Step that prepares the guidance scale tensor for Flux2 inference" + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("guidance_scale", default=4.0), + InputParam("num_images_per_prompt", default=1), + InputParam( + "batch_size", + required=True, + type_hint=int, + description="Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("guidance", type_hint=torch.Tensor, description="Guidance scale tensor"), + ] + + @torch.no_grad() + def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + batch_size = block_state.batch_size * block_state.num_images_per_prompt + guidance = torch.full([1], block_state.guidance_scale, device=device, dtype=torch.float32) + guidance = guidance.expand(batch_size) + block_state.guidance = guidance + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/flux2/decoders.py b/src/diffusers/modular_pipelines/flux2/decoders.py index b769d9119891..c8ad9401efff 100644 --- a/src/diffusers/modular_pipelines/flux2/decoders.py +++ b/src/diffusers/modular_pipelines/flux2/decoders.py @@ -12,7 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, List, Tuple, Union +from __future__ import annotations + +from typing import Any, Union import numpy as np import PIL @@ -29,29 +31,16 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name -class Flux2DecodeStep(ModularPipelineBlocks): +class Flux2UnpackLatentsStep(ModularPipelineBlocks): model_name = "flux2" - @property - def expected_components(self) -> List[ComponentSpec]: - return [ - ComponentSpec("vae", AutoencoderKLFlux2), - ComponentSpec( - "image_processor", - Flux2ImageProcessor, - config=FrozenDict({"vae_scale_factor": 16, "vae_latent_channels": 32}), - default_creation_method="from_config", - ), - ] - @property def description(self) -> str: - return "Step that decodes the denoised latents into images using Flux2 VAE with batch norm denormalization" + return "Step that unpacks the latents from the denoising step" @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [ - InputParam("output_type", default="pil"), InputParam( "latents", required=True, @@ -67,12 +56,12 @@ def inputs(self) -> List[Tuple[str, Any]]: ] @property - def intermediate_outputs(self) -> List[str]: + def intermediate_outputs(self) -> list[str]: return [ OutputParam( - "images", - type_hint=Union[List[PIL.Image.Image], torch.Tensor, np.ndarray], - description="The generated images, can be a list of PIL.Image.Image, torch.Tensor or a numpy array", + "latents", + type_hint=torch.Tensor, + description="The denoise latents from denoising step, unpacked with position IDs.", ) ] @@ -107,6 +96,62 @@ def _unpack_latents_with_ids(x: torch.Tensor, x_ids: torch.Tensor) -> torch.Tens return torch.stack(x_list, dim=0) + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + latents = block_state.latents + latent_ids = block_state.latent_ids + + latents = self._unpack_latents_with_ids(latents, latent_ids) + + block_state.latents = latents + + self.set_block_state(state, block_state) + return components, state + + +class Flux2DecodeStep(ModularPipelineBlocks): + model_name = "flux2" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLFlux2), + ComponentSpec( + "image_processor", + Flux2ImageProcessor, + config=FrozenDict({"vae_scale_factor": 16, "vae_latent_channels": 32}), + default_creation_method="from_config", + ), + ] + + @property + def description(self) -> str: + return "Step that decodes the denoised latents into images using Flux2 VAE with batch norm denormalization" + + @property + def inputs(self) -> list[tuple[str, Any]]: + return [ + InputParam("output_type", default="pil"), + InputParam( + "latents", + required=True, + type_hint=torch.Tensor, + description="The denoised latents from the denoising step", + ), + ] + + @property + def intermediate_outputs(self) -> list[str]: + return [ + OutputParam( + "images", + type_hint=Union[list[PIL.Image.Image], torch.Tensor, np.ndarray], + description="The generated images, can be a list of PIL.Image.Image, torch.Tensor or a numpy array", + ) + ] + @staticmethod def _unpatchify_latents(latents): """Convert patchified latents back to regular format.""" @@ -121,26 +166,20 @@ def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) vae = components.vae - if block_state.output_type == "latent": - block_state.images = block_state.latents - else: - latents = block_state.latents - latent_ids = block_state.latent_ids - - latents = self._unpack_latents_with_ids(latents, latent_ids) + latents = block_state.latents - latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) - latents_bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to( - latents.device, latents.dtype - ) - latents = latents * latents_bn_std + latents_bn_mean + latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) + latents_bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to( + latents.device, latents.dtype + ) + latents = latents * latents_bn_std + latents_bn_mean - latents = self._unpatchify_latents(latents) + latents = self._unpatchify_latents(latents) - block_state.images = vae.decode(latents, return_dict=False)[0] - block_state.images = components.image_processor.postprocess( - block_state.images, output_type=block_state.output_type - ) + block_state.images = vae.decode(latents, return_dict=False)[0] + block_state.images = components.image_processor.postprocess( + block_state.images, output_type=block_state.output_type + ) self.set_block_state(state, block_state) return components, state diff --git a/src/diffusers/modular_pipelines/flux2/denoise.py b/src/diffusers/modular_pipelines/flux2/denoise.py index c12eca65c6a9..66783cc36953 100644 --- a/src/diffusers/modular_pipelines/flux2/denoise.py +++ b/src/diffusers/modular_pipelines/flux2/denoise.py @@ -12,10 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, List, Tuple +from typing import Any import torch +from ...configuration_utils import FrozenDict +from ...guiders import ClassifierFreeGuidance from ...models import Flux2Transformer2DModel from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import is_torch_xla_available, logging @@ -25,8 +27,8 @@ ModularPipelineBlocks, PipelineState, ) -from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam -from .modular_pipeline import Flux2ModularPipeline +from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam +from .modular_pipeline import Flux2KleinModularPipeline, Flux2ModularPipeline if is_torch_xla_available(): @@ -44,7 +46,7 @@ class Flux2LoopDenoiser(ModularPipelineBlocks): model_name = "flux2" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("transformer", Flux2Transformer2DModel)] @property @@ -56,7 +58,7 @@ def description(self) -> str: ) @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [ InputParam("joint_attention_kwargs"), InputParam( @@ -134,11 +136,234 @@ def __call__( return components, block_state +# same as Flux2LoopDenoiser but guidance=None +class Flux2KleinLoopDenoiser(ModularPipelineBlocks): + model_name = "flux2-klein" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Flux2Transformer2DModel)] + + @property + def description(self) -> str: + return ( + "Step within the denoising loop that denoises the latents for Flux2. " + "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " + "object (e.g. `Flux2DenoiseLoopWrapper`)" + ) + + @property + def inputs(self) -> list[tuple[str, Any]]: + return [ + InputParam("joint_attention_kwargs"), + InputParam( + "latents", + required=True, + type_hint=torch.Tensor, + description="The latents to denoise. Shape: (B, seq_len, C)", + ), + InputParam( + "image_latents", + type_hint=torch.Tensor, + description="Packed image latents for conditioning. Shape: (B, img_seq_len, C)", + ), + InputParam( + "image_latent_ids", + type_hint=torch.Tensor, + description="Position IDs for image latents. Shape: (B, img_seq_len, 4)", + ), + InputParam( + "prompt_embeds", + required=True, + type_hint=torch.Tensor, + description="Text embeddings from Qwen3", + ), + InputParam( + "txt_ids", + required=True, + type_hint=torch.Tensor, + description="4D position IDs for text tokens (T, H, W, L)", + ), + InputParam( + "latent_ids", + required=True, + type_hint=torch.Tensor, + description="4D position IDs for latent tokens (T, H, W, L)", + ), + ] + + @torch.no_grad() + def __call__( + self, components: Flux2KleinModularPipeline, block_state: BlockState, i: int, t: torch.Tensor + ) -> PipelineState: + latents = block_state.latents + latent_model_input = latents.to(components.transformer.dtype) + img_ids = block_state.latent_ids + + image_latents = getattr(block_state, "image_latents", None) + if image_latents is not None: + latent_model_input = torch.cat([latents, image_latents], dim=1).to(components.transformer.dtype) + image_latent_ids = block_state.image_latent_ids + img_ids = torch.cat([img_ids, image_latent_ids], dim=1) + + timestep = t.expand(latents.shape[0]).to(latents.dtype) + + noise_pred = components.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + guidance=None, + encoder_hidden_states=block_state.prompt_embeds, + txt_ids=block_state.txt_ids, + img_ids=img_ids, + joint_attention_kwargs=block_state.joint_attention_kwargs, + return_dict=False, + )[0] + + noise_pred = noise_pred[:, : latents.size(1)] + block_state.noise_pred = noise_pred + + return components, block_state + + +# support CFG for Flux2-Klein base model +class Flux2KleinBaseLoopDenoiser(ModularPipelineBlocks): + model_name = "flux2-klein" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", Flux2Transformer2DModel), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 4.0}), + default_creation_method="from_config", + ), + ] + + @property + def expected_configs(self) -> list[ConfigSpec]: + return [ + ConfigSpec(name="is_distilled", default=False), + ] + + @property + def description(self) -> str: + return ( + "Step within the denoising loop that denoises the latents for Flux2. " + "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " + "object (e.g. `Flux2DenoiseLoopWrapper`)" + ) + + @property + def inputs(self) -> list[tuple[str, Any]]: + return [ + InputParam("joint_attention_kwargs"), + InputParam( + "latents", + required=True, + type_hint=torch.Tensor, + description="The latents to denoise. Shape: (B, seq_len, C)", + ), + InputParam( + "image_latents", + type_hint=torch.Tensor, + description="Packed image latents for conditioning. Shape: (B, img_seq_len, C)", + ), + InputParam( + "image_latent_ids", + type_hint=torch.Tensor, + description="Position IDs for image latents. Shape: (B, img_seq_len, 4)", + ), + InputParam( + "prompt_embeds", + required=True, + type_hint=torch.Tensor, + description="Text embeddings from Qwen3", + ), + InputParam( + "negative_prompt_embeds", + required=False, + type_hint=torch.Tensor, + description="Negative text embeddings from Qwen3", + ), + InputParam( + "txt_ids", + required=True, + type_hint=torch.Tensor, + description="4D position IDs for text tokens (T, H, W, L)", + ), + InputParam( + "negative_txt_ids", + required=False, + type_hint=torch.Tensor, + description="4D position IDs for negative text tokens (T, H, W, L)", + ), + InputParam( + "latent_ids", + required=True, + type_hint=torch.Tensor, + description="4D position IDs for latent tokens (T, H, W, L)", + ), + ] + + @torch.no_grad() + def __call__( + self, components: Flux2KleinModularPipeline, block_state: BlockState, i: int, t: torch.Tensor + ) -> PipelineState: + latents = block_state.latents + latent_model_input = latents.to(components.transformer.dtype) + img_ids = block_state.latent_ids + + image_latents = getattr(block_state, "image_latents", None) + if image_latents is not None: + latent_model_input = torch.cat([latents, image_latents], dim=1).to(components.transformer.dtype) + image_latent_ids = block_state.image_latent_ids + img_ids = torch.cat([img_ids, image_latent_ids], dim=1) + + timestep = t.expand(latents.shape[0]).to(latents.dtype) + + guider_inputs = { + "encoder_hidden_states": ( + getattr(block_state, "prompt_embeds", None), + getattr(block_state, "negative_prompt_embeds", None), + ), + "txt_ids": ( + getattr(block_state, "txt_ids", None), + getattr(block_state, "negative_txt_ids", None), + ), + } + + components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) + guider_state = components.guider.prepare_inputs(guider_inputs) + + for guider_state_batch in guider_state: + components.guider.prepare_models(components.transformer) + cond_kwargs = {input_name: getattr(guider_state_batch, input_name) for input_name in guider_inputs.keys()} + + noise_pred = components.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + guidance=None, + img_ids=img_ids, + joint_attention_kwargs=block_state.joint_attention_kwargs, + return_dict=False, + **cond_kwargs, + )[0] + guider_state_batch.noise_pred = noise_pred[:, : latents.size(1)] + components.guider.cleanup_models(components.transformer) + + # perform guidance + block_state.noise_pred = components.guider(guider_state)[0] + + return components, block_state + + class Flux2LoopAfterDenoiser(ModularPipelineBlocks): model_name = "flux2" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] @property @@ -150,15 +375,15 @@ def description(self) -> str: ) @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [] @property - def intermediate_inputs(self) -> List[str]: + def intermediate_inputs(self) -> list[str]: return [InputParam("generator")] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")] @torch.no_grad() @@ -189,14 +414,14 @@ def description(self) -> str: ) @property - def loop_expected_components(self) -> List[ComponentSpec]: + def loop_expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), ComponentSpec("transformer", Flux2Transformer2DModel), ] @property - def loop_inputs(self) -> List[InputParam]: + def loop_inputs(self) -> list[InputParam]: return [ InputParam( "timesteps", @@ -250,3 +475,35 @@ def description(self) -> str: " - `Flux2LoopAfterDenoiser`\n" "This block supports both text-to-image and image-conditioned generation." ) + + +class Flux2KleinDenoiseStep(Flux2DenoiseLoopWrapper): + block_classes = [Flux2KleinLoopDenoiser, Flux2LoopAfterDenoiser] + block_names = ["denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoise step that iteratively denoises the latents for Flux2. \n" + "Its loop logic is defined in `Flux2DenoiseLoopWrapper.__call__` method \n" + "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n" + " - `Flux2KleinLoopDenoiser`\n" + " - `Flux2LoopAfterDenoiser`\n" + "This block supports both text-to-image and image-conditioned generation." + ) + + +class Flux2KleinBaseDenoiseStep(Flux2DenoiseLoopWrapper): + block_classes = [Flux2KleinBaseLoopDenoiser, Flux2LoopAfterDenoiser] + block_names = ["denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoise step that iteratively denoises the latents for Flux2. \n" + "Its loop logic is defined in `Flux2DenoiseLoopWrapper.__call__` method \n" + "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n" + " - `Flux2KleinBaseLoopDenoiser`\n" + " - `Flux2LoopAfterDenoiser`\n" + "This block supports both text-to-image and image-conditioned generation." + ) diff --git a/src/diffusers/modular_pipelines/flux2/encoders.py b/src/diffusers/modular_pipelines/flux2/encoders.py index 6cb0e3bf0a26..81d20a8f4c65 100644 --- a/src/diffusers/modular_pipelines/flux2/encoders.py +++ b/src/diffusers/modular_pipelines/flux2/encoders.py @@ -12,22 +12,23 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Optional, Tuple, Union import torch -from transformers import AutoProcessor, Mistral3ForConditionalGeneration +from transformers import AutoProcessor, Mistral3ForConditionalGeneration, Qwen2TokenizerFast, Qwen3ForCausalLM +from ...configuration_utils import FrozenDict +from ...guiders import ClassifierFreeGuidance from ...models import AutoencoderKLFlux2 from ...utils import logging from ..modular_pipeline import ModularPipelineBlocks, PipelineState -from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam -from .modular_pipeline import Flux2ModularPipeline +from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam +from .modular_pipeline import Flux2KleinModularPipeline, Flux2ModularPipeline logger = logging.get_logger(__name__) # pylint: disable=invalid-name -def format_text_input(prompts: List[str], system_message: str = None): +def format_text_input(prompts: list[str], system_message: str = None): """Format prompts for Mistral3 chat template.""" cleaned_txt = [prompt.replace("[IMG]", "") for prompt in prompts] @@ -45,7 +46,7 @@ def format_text_input(prompts: List[str], system_message: str = None): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -69,24 +70,22 @@ def description(self) -> str: return "Text Encoder step that generates text embeddings using Mistral3 to guide the image generation" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("text_encoder", Mistral3ForConditionalGeneration), ComponentSpec("tokenizer", AutoProcessor), ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("prompt"), - InputParam("prompt_embeds", type_hint=torch.Tensor, required=False), InputParam("max_sequence_length", type_hint=int, default=512, required=False), - InputParam("text_encoder_out_layers", type_hint=Tuple[int], default=(10, 20, 30), required=False), - InputParam("joint_attention_kwargs"), + InputParam("text_encoder_out_layers", type_hint=tuple[int], default=(10, 20, 30), required=False), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "prompt_embeds", @@ -99,28 +98,21 @@ def intermediate_outputs(self) -> List[OutputParam]: @staticmethod def check_inputs(block_state): prompt = block_state.prompt - prompt_embeds = getattr(block_state, "prompt_embeds", None) - - if prompt is not None and prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. " - "Please make sure to only forward one of the two." - ) - elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + if prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") @staticmethod def _get_mistral_3_prompt_embeds( text_encoder: Mistral3ForConditionalGeneration, tokenizer: AutoProcessor, - prompt: Union[str, List[str]], - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, + prompt: str | list[str], + dtype: torch.dtype | None = None, + device: torch.device | None = None, max_sequence_length: int = 512, # fmt: off system_message: str = "You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object attribution and actions without speculation.", # fmt: on - hidden_states_layers: Tuple[int] = (10, 20, 30), + hidden_states_layers: tuple[int] = (10, 20, 30), ): dtype = text_encoder.dtype if dtype is None else dtype device = text_encoder.device if device is None else device @@ -165,10 +157,6 @@ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> Pi block_state.device = components._execution_device - if block_state.prompt_embeds is not None: - self.set_block_state(state, block_state) - return components, state - prompt = block_state.prompt if prompt is None: prompt = "" @@ -198,18 +186,17 @@ def description(self) -> str: return "Text Encoder step that generates text embeddings using a remote API endpoint" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("prompt"), - InputParam("prompt_embeds", type_hint=torch.Tensor, required=False), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "prompt_embeds", @@ -222,15 +209,8 @@ def intermediate_outputs(self) -> List[OutputParam]: @staticmethod def check_inputs(block_state): prompt = block_state.prompt - prompt_embeds = getattr(block_state, "prompt_embeds", None) - - if prompt is not None and prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. " - "Please make sure to only forward one of the two." - ) - elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + if prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(block_state.prompt)}") @torch.no_grad() def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState: @@ -244,10 +224,6 @@ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> Pi block_state.device = components._execution_device - if block_state.prompt_embeds is not None: - self.set_block_state(state, block_state) - return components, state - prompt = block_state.prompt if prompt is None: prompt = "" @@ -270,6 +246,289 @@ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> Pi return components, state +class Flux2KleinTextEncoderStep(ModularPipelineBlocks): + model_name = "flux2-klein" + + @property + def description(self) -> str: + return "Text Encoder step that generates text embeddings using Qwen3 to guide the image generation" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", Qwen3ForCausalLM), + ComponentSpec("tokenizer", Qwen2TokenizerFast), + ] + + @property + def expected_configs(self) -> list[ConfigSpec]: + return [ + ConfigSpec(name="is_distilled", default=True), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("prompt"), + InputParam("max_sequence_length", type_hint=int, default=512, required=False), + InputParam("text_encoder_out_layers", type_hint=tuple[int], default=(9, 18, 27), required=False), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "prompt_embeds", + kwargs_type="denoiser_input_fields", + type_hint=torch.Tensor, + description="Text embeddings from qwen3 used to guide the image generation", + ), + ] + + @staticmethod + def check_inputs(block_state): + prompt = block_state.prompt + + if prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + @staticmethod + # Copied from diffusers.pipelines.flux2.pipeline_flux2_klein.Flux2KleinPipeline._get_qwen3_prompt_embeds + def _get_qwen3_prompt_embeds( + text_encoder: Qwen3ForCausalLM, + tokenizer: Qwen2TokenizerFast, + prompt: str | list[str], + dtype: torch.dtype | None = None, + device: torch.device | None = None, + max_sequence_length: int = 512, + hidden_states_layers: list[int] = (9, 18, 27), + ): + dtype = text_encoder.dtype if dtype is None else dtype + device = text_encoder.device if device is None else device + + prompt = [prompt] if isinstance(prompt, str) else prompt + + all_input_ids = [] + all_attention_masks = [] + + for single_prompt in prompt: + messages = [{"role": "user", "content": single_prompt}] + text = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + inputs = tokenizer( + text, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=max_sequence_length, + ) + + all_input_ids.append(inputs["input_ids"]) + all_attention_masks.append(inputs["attention_mask"]) + + input_ids = torch.cat(all_input_ids, dim=0).to(device) + attention_mask = torch.cat(all_attention_masks, dim=0).to(device) + + # Forward pass through the model + output = text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + use_cache=False, + ) + + # Only use outputs from intermediate layers and stack them + out = torch.stack([output.hidden_states[k] for k in hidden_states_layers], dim=1) + out = out.to(dtype=dtype, device=device) + + batch_size, num_channels, seq_len, hidden_dim = out.shape + prompt_embeds = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_channels * hidden_dim) + + return prompt_embeds + + @torch.no_grad() + def __call__(self, components: Flux2KleinModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs(block_state) + + device = components._execution_device + + prompt = block_state.prompt + if prompt is None: + prompt = "" + prompt = [prompt] if isinstance(prompt, str) else prompt + + block_state.prompt_embeds = self._get_qwen3_prompt_embeds( + text_encoder=components.text_encoder, + tokenizer=components.tokenizer, + prompt=prompt, + device=device, + max_sequence_length=block_state.max_sequence_length, + hidden_states_layers=block_state.text_encoder_out_layers, + ) + + self.set_block_state(state, block_state) + return components, state + + +class Flux2KleinBaseTextEncoderStep(ModularPipelineBlocks): + model_name = "flux2-klein" + + @property + def description(self) -> str: + return "Text Encoder step that generates text embeddings using Qwen3 to guide the image generation" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", Qwen3ForCausalLM), + ComponentSpec("tokenizer", Qwen2TokenizerFast), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 4.0}), + default_creation_method="from_config", + ), + ] + + @property + def expected_configs(self) -> list[ConfigSpec]: + return [ + ConfigSpec(name="is_distilled", default=False), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("prompt"), + InputParam("max_sequence_length", type_hint=int, default=512, required=False), + InputParam("text_encoder_out_layers", type_hint=tuple[int], default=(9, 18, 27), required=False), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "prompt_embeds", + kwargs_type="denoiser_input_fields", + type_hint=torch.Tensor, + description="Text embeddings from qwen3 used to guide the image generation", + ), + OutputParam( + "negative_prompt_embeds", + kwargs_type="denoiser_input_fields", + type_hint=torch.Tensor, + description="Negative text embeddings from qwen3 used to guide the image generation", + ), + ] + + @staticmethod + def check_inputs(block_state): + prompt = block_state.prompt + + if prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + @staticmethod + # Copied from diffusers.pipelines.flux2.pipeline_flux2_klein.Flux2KleinPipeline._get_qwen3_prompt_embeds + def _get_qwen3_prompt_embeds( + text_encoder: Qwen3ForCausalLM, + tokenizer: Qwen2TokenizerFast, + prompt: str | list[str], + dtype: torch.dtype | None = None, + device: torch.device | None = None, + max_sequence_length: int = 512, + hidden_states_layers: list[int] = (9, 18, 27), + ): + dtype = text_encoder.dtype if dtype is None else dtype + device = text_encoder.device if device is None else device + + prompt = [prompt] if isinstance(prompt, str) else prompt + + all_input_ids = [] + all_attention_masks = [] + + for single_prompt in prompt: + messages = [{"role": "user", "content": single_prompt}] + text = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + inputs = tokenizer( + text, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=max_sequence_length, + ) + + all_input_ids.append(inputs["input_ids"]) + all_attention_masks.append(inputs["attention_mask"]) + + input_ids = torch.cat(all_input_ids, dim=0).to(device) + attention_mask = torch.cat(all_attention_masks, dim=0).to(device) + + # Forward pass through the model + output = text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + use_cache=False, + ) + + # Only use outputs from intermediate layers and stack them + out = torch.stack([output.hidden_states[k] for k in hidden_states_layers], dim=1) + out = out.to(dtype=dtype, device=device) + + batch_size, num_channels, seq_len, hidden_dim = out.shape + prompt_embeds = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_channels * hidden_dim) + + return prompt_embeds + + @torch.no_grad() + def __call__(self, components: Flux2KleinModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs(block_state) + + device = components._execution_device + + prompt = block_state.prompt + if prompt is None: + prompt = "" + prompt = [prompt] if isinstance(prompt, str) else prompt + + block_state.prompt_embeds = self._get_qwen3_prompt_embeds( + text_encoder=components.text_encoder, + tokenizer=components.tokenizer, + prompt=prompt, + device=device, + max_sequence_length=block_state.max_sequence_length, + hidden_states_layers=block_state.text_encoder_out_layers, + ) + + if components.requires_unconditional_embeds: + negative_prompt = [""] * len(prompt) + block_state.negative_prompt_embeds = self._get_qwen3_prompt_embeds( + text_encoder=components.text_encoder, + tokenizer=components.tokenizer, + prompt=negative_prompt, + device=device, + max_sequence_length=block_state.max_sequence_length, + hidden_states_layers=block_state.text_encoder_out_layers, + ) + else: + block_state.negative_prompt_embeds = None + + self.set_block_state(state, block_state) + return components, state + + class Flux2VaeEncoderStep(ModularPipelineBlocks): model_name = "flux2" @@ -278,22 +537,22 @@ def description(self) -> str: return "VAE Encoder step that encodes preprocessed images into latent representations for Flux2." @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("vae", AutoencoderKLFlux2)] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ - InputParam("condition_images", type_hint=List[torch.Tensor]), + InputParam("condition_images", type_hint=list[torch.Tensor]), InputParam("generator"), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "image_latents", - type_hint=List[torch.Tensor], + type_hint=list[torch.Tensor], description="List of latent representations for each reference image", ), ] diff --git a/src/diffusers/modular_pipelines/flux2/inputs.py b/src/diffusers/modular_pipelines/flux2/inputs.py index c9e337fb0bf0..7c6881f70a46 100644 --- a/src/diffusers/modular_pipelines/flux2/inputs.py +++ b/src/diffusers/modular_pipelines/flux2/inputs.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List - import torch from ...configuration_utils import FrozenDict @@ -39,7 +37,7 @@ def description(self) -> str: ) @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("num_images_per_prompt", default=1), InputParam( @@ -47,12 +45,12 @@ def inputs(self) -> List[InputParam]: required=True, kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, - description="Pre-generated text embeddings from Mistral3. Can be generated from text_encoder step.", + description="Pre-generated text embeddings. Can be generated from text_encoder step.", ), ] @property - def intermediate_outputs(self) -> List[str]: + def intermediate_outputs(self) -> list[str]: return [ OutputParam( "batch_size", @@ -89,6 +87,90 @@ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> Pi return components, state +class Flux2KleinBaseTextInputStep(ModularPipelineBlocks): + model_name = "flux2-klein" + + @property + def description(self) -> str: + return ( + "This step:\n" + " 1. Determines `batch_size` and `dtype` based on `prompt_embeds`\n" + " 2. Ensures all text embeddings have consistent batch sizes (batch_size * num_images_per_prompt)" + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("num_images_per_prompt", default=1), + InputParam( + "prompt_embeds", + required=True, + kwargs_type="denoiser_input_fields", + type_hint=torch.Tensor, + description="Pre-generated text embeddings. Can be generated from text_encoder step.", + ), + InputParam( + "negative_prompt_embeds", + required=False, + kwargs_type="denoiser_input_fields", + type_hint=torch.Tensor, + description="Pre-generated negative text embeddings. Can be generated from text_encoder step.", + ), + ] + + @property + def intermediate_outputs(self) -> list[str]: + return [ + OutputParam( + "batch_size", + type_hint=int, + description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt", + ), + OutputParam( + "dtype", + type_hint=torch.dtype, + description="Data type of model tensor inputs (determined by `prompt_embeds`)", + ), + OutputParam( + "prompt_embeds", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="Text embeddings used to guide the image generation", + ), + OutputParam( + "negative_prompt_embeds", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="Negative text embeddings used to guide the image generation", + ), + ] + + @torch.no_grad() + def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + block_state.batch_size = block_state.prompt_embeds.shape[0] + block_state.dtype = block_state.prompt_embeds.dtype + + _, seq_len, _ = block_state.prompt_embeds.shape + block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, block_state.num_images_per_prompt, 1) + block_state.prompt_embeds = block_state.prompt_embeds.view( + block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1 + ) + + if block_state.negative_prompt_embeds is not None: + _, seq_len, _ = block_state.negative_prompt_embeds.shape + block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.repeat( + 1, block_state.num_images_per_prompt, 1 + ) + block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.view( + block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1 + ) + + self.set_block_state(state, block_state) + return components, state + + class Flux2ProcessImagesInputStep(ModularPipelineBlocks): model_name = "flux2" @@ -97,7 +179,7 @@ def description(self) -> str: return "Image preprocess step for Flux2. Validates and preprocesses reference images." @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec( "image_processor", @@ -108,7 +190,7 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("image"), InputParam("height"), @@ -116,8 +198,8 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: - return [OutputParam(name="condition_images", type_hint=List[torch.Tensor])] + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam(name="condition_images", type_hint=list[torch.Tensor])] @torch.no_grad() def __call__(self, components: Flux2ModularPipeline, state: PipelineState): diff --git a/src/diffusers/modular_pipelines/flux2/modular_blocks.py b/src/diffusers/modular_pipelines/flux2/modular_blocks.py deleted file mode 100644 index a31673b6e78c..000000000000 --- a/src/diffusers/modular_pipelines/flux2/modular_blocks.py +++ /dev/null @@ -1,166 +0,0 @@ -# Copyright 2025 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from ...utils import logging -from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks -from ..modular_pipeline_utils import InsertableDict -from .before_denoise import ( - Flux2PrepareImageLatentsStep, - Flux2PrepareLatentsStep, - Flux2RoPEInputsStep, - Flux2SetTimestepsStep, -) -from .decoders import Flux2DecodeStep -from .denoise import Flux2DenoiseStep -from .encoders import ( - Flux2RemoteTextEncoderStep, - Flux2TextEncoderStep, - Flux2VaeEncoderStep, -) -from .inputs import ( - Flux2ProcessImagesInputStep, - Flux2TextInputStep, -) - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -Flux2VaeEncoderBlocks = InsertableDict( - [ - ("preprocess", Flux2ProcessImagesInputStep()), - ("encode", Flux2VaeEncoderStep()), - ("prepare_image_latents", Flux2PrepareImageLatentsStep()), - ] -) - - -class Flux2VaeEncoderSequentialStep(SequentialPipelineBlocks): - model_name = "flux2" - - block_classes = Flux2VaeEncoderBlocks.values() - block_names = Flux2VaeEncoderBlocks.keys() - - @property - def description(self) -> str: - return "VAE encoder step that preprocesses, encodes, and prepares image latents for Flux2 conditioning." - - -class Flux2AutoVaeEncoderStep(AutoPipelineBlocks): - block_classes = [Flux2VaeEncoderSequentialStep] - block_names = ["img_conditioning"] - block_trigger_inputs = ["image"] - - @property - def description(self): - return ( - "VAE encoder step that encodes the image inputs into their latent representations.\n" - "This is an auto pipeline block that works for image conditioning tasks.\n" - " - `Flux2VaeEncoderSequentialStep` is used when `image` is provided.\n" - " - If `image` is not provided, step will be skipped." - ) - - -Flux2BeforeDenoiseBlocks = InsertableDict( - [ - ("prepare_latents", Flux2PrepareLatentsStep()), - ("set_timesteps", Flux2SetTimestepsStep()), - ("prepare_rope_inputs", Flux2RoPEInputsStep()), - ] -) - - -class Flux2BeforeDenoiseStep(SequentialPipelineBlocks): - model_name = "flux2" - - block_classes = Flux2BeforeDenoiseBlocks.values() - block_names = Flux2BeforeDenoiseBlocks.keys() - - @property - def description(self): - return "Before denoise step that prepares the inputs for the denoise step in Flux2 generation." - - -AUTO_BLOCKS = InsertableDict( - [ - ("text_encoder", Flux2TextEncoderStep()), - ("text_input", Flux2TextInputStep()), - ("vae_image_encoder", Flux2AutoVaeEncoderStep()), - ("before_denoise", Flux2BeforeDenoiseStep()), - ("denoise", Flux2DenoiseStep()), - ("decode", Flux2DecodeStep()), - ] -) - - -REMOTE_AUTO_BLOCKS = InsertableDict( - [ - ("text_encoder", Flux2RemoteTextEncoderStep()), - ("text_input", Flux2TextInputStep()), - ("vae_image_encoder", Flux2AutoVaeEncoderStep()), - ("before_denoise", Flux2BeforeDenoiseStep()), - ("denoise", Flux2DenoiseStep()), - ("decode", Flux2DecodeStep()), - ] -) - - -class Flux2AutoBlocks(SequentialPipelineBlocks): - model_name = "flux2" - - block_classes = AUTO_BLOCKS.values() - block_names = AUTO_BLOCKS.keys() - - @property - def description(self): - return ( - "Auto Modular pipeline for text-to-image and image-conditioned generation using Flux2.\n" - "- For text-to-image generation, all you need to provide is `prompt`.\n" - "- For image-conditioned generation, you need to provide `image` (list of PIL images)." - ) - - -TEXT2IMAGE_BLOCKS = InsertableDict( - [ - ("text_encoder", Flux2TextEncoderStep()), - ("text_input", Flux2TextInputStep()), - ("prepare_latents", Flux2PrepareLatentsStep()), - ("set_timesteps", Flux2SetTimestepsStep()), - ("prepare_rope_inputs", Flux2RoPEInputsStep()), - ("denoise", Flux2DenoiseStep()), - ("decode", Flux2DecodeStep()), - ] -) - -IMAGE_CONDITIONED_BLOCKS = InsertableDict( - [ - ("text_encoder", Flux2TextEncoderStep()), - ("text_input", Flux2TextInputStep()), - ("preprocess_images", Flux2ProcessImagesInputStep()), - ("vae_encoder", Flux2VaeEncoderStep()), - ("prepare_image_latents", Flux2PrepareImageLatentsStep()), - ("prepare_latents", Flux2PrepareLatentsStep()), - ("set_timesteps", Flux2SetTimestepsStep()), - ("prepare_rope_inputs", Flux2RoPEInputsStep()), - ("denoise", Flux2DenoiseStep()), - ("decode", Flux2DecodeStep()), - ] -) - -ALL_BLOCKS = { - "text2image": TEXT2IMAGE_BLOCKS, - "image_conditioned": IMAGE_CONDITIONED_BLOCKS, - "auto": AUTO_BLOCKS, - "remote": REMOTE_AUTO_BLOCKS, -} diff --git a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py new file mode 100644 index 000000000000..b1033a7dff9e --- /dev/null +++ b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py @@ -0,0 +1,356 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from ...utils import logging +from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import ( + Flux2PrepareGuidanceStep, + Flux2PrepareImageLatentsStep, + Flux2PrepareLatentsStep, + Flux2RoPEInputsStep, + Flux2SetTimestepsStep, +) +from .decoders import Flux2DecodeStep, Flux2UnpackLatentsStep +from .denoise import Flux2DenoiseStep +from .encoders import ( + Flux2TextEncoderStep, + Flux2VaeEncoderStep, +) +from .inputs import ( + Flux2ProcessImagesInputStep, + Flux2TextInputStep, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# auto_docstring +class Flux2VaeEncoderSequentialStep(SequentialPipelineBlocks): + """ + VAE encoder step that preprocesses, encodes, and prepares image latents for Flux2 conditioning. + + Components: + image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) + + Inputs: + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + condition_images (`list`): + TODO: Add description. + image_latents (`list`): + List of latent representations for each reference image + """ + + model_name = "flux2" + + block_classes = [Flux2ProcessImagesInputStep(), Flux2VaeEncoderStep()] + block_names = ["preprocess", "encode"] + + @property + def description(self) -> str: + return "VAE encoder step that preprocesses, encodes, and prepares image latents for Flux2 conditioning." + + +# auto_docstring +class Flux2AutoVaeEncoderStep(AutoPipelineBlocks): + """ + VAE encoder step that encodes the image inputs into their latent representations. + This is an auto pipeline block that works for image conditioning tasks. + - `Flux2VaeEncoderSequentialStep` is used when `image` is provided. + - If `image` is not provided, step will be skipped. + + Components: + image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) + + Inputs: + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + condition_images (`list`): + TODO: Add description. + image_latents (`list`): + List of latent representations for each reference image + """ + + block_classes = [Flux2VaeEncoderSequentialStep] + block_names = ["img_conditioning"] + block_trigger_inputs = ["image"] + + @property + def description(self): + return ( + "VAE encoder step that encodes the image inputs into their latent representations.\n" + "This is an auto pipeline block that works for image conditioning tasks.\n" + " - `Flux2VaeEncoderSequentialStep` is used when `image` is provided.\n" + " - If `image` is not provided, step will be skipped." + ) + + +Flux2CoreDenoiseBlocks = InsertableDict( + [ + ("input", Flux2TextInputStep()), + ("prepare_latents", Flux2PrepareLatentsStep()), + ("set_timesteps", Flux2SetTimestepsStep()), + ("prepare_guidance", Flux2PrepareGuidanceStep()), + ("prepare_rope_inputs", Flux2RoPEInputsStep()), + ("denoise", Flux2DenoiseStep()), + ("after_denoise", Flux2UnpackLatentsStep()), + ] +) + + +# auto_docstring +class Flux2CoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step that performs the denoising process for Flux2-dev. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + guidance_scale (`None`, *optional*, defaults to 4.0): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + image_latents (`Tensor`, *optional*): + Packed image latents for conditioning. Shape: (B, img_seq_len, C) + image_latent_ids (`Tensor`, *optional*): + Position IDs for image latents. Shape: (B, img_seq_len, 4) + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "flux2" + + block_classes = Flux2CoreDenoiseBlocks.values() + block_names = Flux2CoreDenoiseBlocks.keys() + + @property + def description(self): + return "Core denoise step that performs the denoising process for Flux2-dev." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +Flux2ImageConditionedCoreDenoiseBlocks = InsertableDict( + [ + ("input", Flux2TextInputStep()), + ("prepare_image_latents", Flux2PrepareImageLatentsStep()), + ("prepare_latents", Flux2PrepareLatentsStep()), + ("set_timesteps", Flux2SetTimestepsStep()), + ("prepare_guidance", Flux2PrepareGuidanceStep()), + ("prepare_rope_inputs", Flux2RoPEInputsStep()), + ("denoise", Flux2DenoiseStep()), + ("after_denoise", Flux2UnpackLatentsStep()), + ] +) + + +# auto_docstring +class Flux2ImageConditionedCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step that performs the denoising process for Flux2-dev with image conditioning. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + image_latents (`list`, *optional*): + TODO: Add description. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + guidance_scale (`None`, *optional*, defaults to 4.0): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "flux2" + + block_classes = Flux2ImageConditionedCoreDenoiseBlocks.values() + block_names = Flux2ImageConditionedCoreDenoiseBlocks.keys() + + @property + def description(self): + return "Core denoise step that performs the denoising process for Flux2-dev with image conditioning." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +class Flux2AutoCoreDenoiseStep(AutoPipelineBlocks): + model_name = "flux2" + + block_classes = [Flux2ImageConditionedCoreDenoiseStep, Flux2CoreDenoiseStep] + block_names = ["image_conditioned", "text2image"] + block_trigger_inputs = ["image_latents", None] + + @property + def description(self): + return ( + "Auto core denoise step that performs the denoising process for Flux2-dev." + "This is an auto pipeline block that works for text-to-image and image-conditioned generation." + " - `Flux2CoreDenoiseStep` is used for text-to-image generation.\n" + " - `Flux2ImageConditionedCoreDenoiseStep` is used for image-conditioned generation.\n" + ) + + +AUTO_BLOCKS = InsertableDict( + [ + ("text_encoder", Flux2TextEncoderStep()), + ("vae_encoder", Flux2AutoVaeEncoderStep()), + ("denoise", Flux2AutoCoreDenoiseStep()), + ("decode", Flux2DecodeStep()), + ] +) + + +# auto_docstring +class Flux2AutoBlocks(SequentialPipelineBlocks): + """ + Auto Modular pipeline for text-to-image and image-conditioned generation using Flux2. + + Supported workflows: + - `text2image`: requires `prompt` + - `image_conditioned`: requires `image`, `prompt` + + Components: + text_encoder (`Mistral3ForConditionalGeneration`) tokenizer (`AutoProcessor`) image_processor + (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) scheduler (`FlowMatchEulerDiscreteScheduler`) transformer + (`Flux2Transformer2DModel`) + + Inputs: + prompt (`None`, *optional*): + TODO: Add description. + max_sequence_length (`int`, *optional*, defaults to 512): + TODO: Add description. + text_encoder_out_layers (`tuple`, *optional*, defaults to (10, 20, 30)): + TODO: Add description. + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + image_latents (`list`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`): + TODO: Add description. + num_inference_steps (`None`): + TODO: Add description. + timesteps (`None`): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + guidance_scale (`None`, *optional*, defaults to 4.0): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + image_latent_ids (`Tensor`, *optional*): + Position IDs for image latents. Shape: (B, img_seq_len, 4) + output_type (`None`, *optional*, defaults to pil): + TODO: Add description. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "flux2" + + block_classes = AUTO_BLOCKS.values() + block_names = AUTO_BLOCKS.keys() + _workflow_map = { + "text2image": {"prompt": True}, + "image_conditioned": {"image": True, "prompt": True}, + } + + @property + def description(self): + return "Auto Modular pipeline for text-to-image and image-conditioned generation using Flux2." + + @property + def outputs(self): + return [ + OutputParam.template("images"), + ] diff --git a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py new file mode 100644 index 000000000000..5dbae43a5a7f --- /dev/null +++ b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py @@ -0,0 +1,400 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from ...utils import logging +from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import ( + Flux2PrepareImageLatentsStep, + Flux2PrepareLatentsStep, + Flux2RoPEInputsStep, + Flux2SetTimestepsStep, +) +from .decoders import Flux2DecodeStep, Flux2UnpackLatentsStep +from .denoise import Flux2KleinDenoiseStep +from .encoders import ( + Flux2KleinTextEncoderStep, + Flux2VaeEncoderStep, +) +from .inputs import ( + Flux2ProcessImagesInputStep, + Flux2TextInputStep, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +################ +# VAE encoder +################ + + +# auto_docstring +class Flux2KleinVaeEncoderSequentialStep(SequentialPipelineBlocks): + """ + VAE encoder step that preprocesses and encodes the image inputs into their latent representations. + + Components: + image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) + + Inputs: + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + condition_images (`list`): + TODO: Add description. + image_latents (`list`): + List of latent representations for each reference image + """ + + model_name = "flux2-klein" + + block_classes = [Flux2ProcessImagesInputStep(), Flux2VaeEncoderStep()] + block_names = ["preprocess", "encode"] + + @property + def description(self) -> str: + return "VAE encoder step that preprocesses and encodes the image inputs into their latent representations." + + +# auto_docstring +class Flux2KleinAutoVaeEncoderStep(AutoPipelineBlocks): + """ + VAE encoder step that encodes the image inputs into their latent representations. + This is an auto pipeline block that works for image conditioning tasks. + - `Flux2KleinVaeEncoderSequentialStep` is used when `image` is provided. + - If `image` is not provided, step will be skipped. + + Components: + image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) + + Inputs: + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + condition_images (`list`): + TODO: Add description. + image_latents (`list`): + List of latent representations for each reference image + """ + + model_name = "flux2-klein" + + block_classes = [Flux2KleinVaeEncoderSequentialStep] + block_names = ["img_conditioning"] + block_trigger_inputs = ["image"] + + @property + def description(self): + return ( + "VAE encoder step that encodes the image inputs into their latent representations.\n" + "This is an auto pipeline block that works for image conditioning tasks.\n" + " - `Flux2KleinVaeEncoderSequentialStep` is used when `image` is provided.\n" + " - If `image` is not provided, step will be skipped." + ) + + +### +### Core denoise +### + +Flux2KleinCoreDenoiseBlocks = InsertableDict( + [ + ("input", Flux2TextInputStep()), + ("prepare_latents", Flux2PrepareLatentsStep()), + ("set_timesteps", Flux2SetTimestepsStep()), + ("prepare_rope_inputs", Flux2RoPEInputsStep()), + ("denoise", Flux2KleinDenoiseStep()), + ("after_denoise", Flux2UnpackLatentsStep()), + ] +) + + +# auto_docstring +class Flux2KleinCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step that performs the denoising process for Flux2-Klein (distilled model), for text-to-image + generation. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + image_latents (`Tensor`, *optional*): + Packed image latents for conditioning. Shape: (B, img_seq_len, C) + image_latent_ids (`Tensor`, *optional*): + Position IDs for image latents. Shape: (B, img_seq_len, 4) + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "flux2-klein" + + block_classes = Flux2KleinCoreDenoiseBlocks.values() + block_names = Flux2KleinCoreDenoiseBlocks.keys() + + @property + def description(self): + return "Core denoise step that performs the denoising process for Flux2-Klein (distilled model), for text-to-image generation." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +Flux2KleinImageConditionedCoreDenoiseBlocks = InsertableDict( + [ + ("input", Flux2TextInputStep()), + ("prepare_image_latents", Flux2PrepareImageLatentsStep()), + ("prepare_latents", Flux2PrepareLatentsStep()), + ("set_timesteps", Flux2SetTimestepsStep()), + ("prepare_rope_inputs", Flux2RoPEInputsStep()), + ("denoise", Flux2KleinDenoiseStep()), + ("after_denoise", Flux2UnpackLatentsStep()), + ] +) + + +# auto_docstring +class Flux2KleinImageConditionedCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step that performs the denoising process for Flux2-Klein (distilled model) with image conditioning. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + image_latents (`list`, *optional*): + TODO: Add description. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "flux2-klein" + + block_classes = Flux2KleinImageConditionedCoreDenoiseBlocks.values() + block_names = Flux2KleinImageConditionedCoreDenoiseBlocks.keys() + + @property + def description(self): + return "Core denoise step that performs the denoising process for Flux2-Klein (distilled model) with image conditioning." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +# auto_docstring +class Flux2KleinAutoCoreDenoiseStep(AutoPipelineBlocks): + """ + Auto core denoise step that performs the denoising process for Flux2-Klein. + This is an auto pipeline block that works for text-to-image and image-conditioned generation. + - `Flux2KleinCoreDenoiseStep` is used for text-to-image generation. + - `Flux2KleinImageConditionedCoreDenoiseStep` is used for image-conditioned generation. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + image_latents (`list`, *optional*): + TODO: Add description. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`None`): + TODO: Add description. + timesteps (`None`): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + image_latent_ids (`Tensor`, *optional*): + Position IDs for image latents. Shape: (B, img_seq_len, 4) + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "flux2-klein" + block_classes = [Flux2KleinImageConditionedCoreDenoiseStep, Flux2KleinCoreDenoiseStep] + block_names = ["image_conditioned", "text2image"] + block_trigger_inputs = ["image_latents", None] + + @property + def description(self): + return ( + "Auto core denoise step that performs the denoising process for Flux2-Klein.\n" + "This is an auto pipeline block that works for text-to-image and image-conditioned generation.\n" + " - `Flux2KleinCoreDenoiseStep` is used for text-to-image generation.\n" + " - `Flux2KleinImageConditionedCoreDenoiseStep` is used for image-conditioned generation.\n" + ) + + +### +### Auto blocks +### + + +# auto_docstring +class Flux2KleinAutoBlocks(SequentialPipelineBlocks): + """ + Auto blocks that perform the text-to-image and image-conditioned generation using Flux2-Klein. + + Supported workflows: + - `text2image`: requires `prompt` + - `image_conditioned`: requires `image`, `prompt` + + Components: + text_encoder (`Qwen3ForCausalLM`) tokenizer (`Qwen2TokenizerFast`) image_processor (`Flux2ImageProcessor`) + vae (`AutoencoderKLFlux2`) scheduler (`FlowMatchEulerDiscreteScheduler`) transformer + (`Flux2Transformer2DModel`) + + Configs: + is_distilled (default: True) + + Inputs: + prompt (`None`, *optional*): + TODO: Add description. + max_sequence_length (`int`, *optional*, defaults to 512): + TODO: Add description. + text_encoder_out_layers (`tuple`, *optional*, defaults to (9, 18, 27)): + TODO: Add description. + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + image_latents (`list`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`): + TODO: Add description. + num_inference_steps (`None`): + TODO: Add description. + timesteps (`None`): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + image_latent_ids (`Tensor`, *optional*): + Position IDs for image latents. Shape: (B, img_seq_len, 4) + output_type (`None`, *optional*, defaults to pil): + TODO: Add description. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "flux2-klein" + block_classes = [ + Flux2KleinTextEncoderStep(), + Flux2KleinAutoVaeEncoderStep(), + Flux2KleinAutoCoreDenoiseStep(), + Flux2DecodeStep(), + ] + block_names = ["text_encoder", "vae_encoder", "denoise", "decode"] + _workflow_map = { + "text2image": {"prompt": True}, + "image_conditioned": {"image": True, "prompt": True}, + } + + @property + def description(self): + return "Auto blocks that perform the text-to-image and image-conditioned generation using Flux2-Klein." + + @property + def outputs(self): + return [ + OutputParam.template("images"), + ] diff --git a/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py new file mode 100644 index 000000000000..42e025c622b4 --- /dev/null +++ b/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py @@ -0,0 +1,413 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from ...utils import logging +from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import ( + Flux2KleinBaseRoPEInputsStep, + Flux2PrepareImageLatentsStep, + Flux2PrepareLatentsStep, + Flux2SetTimestepsStep, +) +from .decoders import Flux2DecodeStep, Flux2UnpackLatentsStep +from .denoise import Flux2KleinBaseDenoiseStep +from .encoders import ( + Flux2KleinBaseTextEncoderStep, + Flux2VaeEncoderStep, +) +from .inputs import ( + Flux2KleinBaseTextInputStep, + Flux2ProcessImagesInputStep, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +################ +# VAE encoder +################ + + +# auto_docstring +class Flux2KleinBaseVaeEncoderSequentialStep(SequentialPipelineBlocks): + """ + VAE encoder step that preprocesses and encodes the image inputs into their latent representations. + + Components: + image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) + + Inputs: + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + condition_images (`list`): + TODO: Add description. + image_latents (`list`): + List of latent representations for each reference image + """ + + model_name = "flux2" + + block_classes = [Flux2ProcessImagesInputStep(), Flux2VaeEncoderStep()] + block_names = ["preprocess", "encode"] + + @property + def description(self) -> str: + return "VAE encoder step that preprocesses and encodes the image inputs into their latent representations." + + +# auto_docstring +class Flux2KleinBaseAutoVaeEncoderStep(AutoPipelineBlocks): + """ + VAE encoder step that encodes the image inputs into their latent representations. + This is an auto pipeline block that works for image conditioning tasks. + - `Flux2KleinBaseVaeEncoderSequentialStep` is used when `image` is provided. + - If `image` is not provided, step will be skipped. + + Components: + image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) + + Inputs: + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + condition_images (`list`): + TODO: Add description. + image_latents (`list`): + List of latent representations for each reference image + """ + + block_classes = [Flux2KleinBaseVaeEncoderSequentialStep] + block_names = ["img_conditioning"] + block_trigger_inputs = ["image"] + + @property + def description(self): + return ( + "VAE encoder step that encodes the image inputs into their latent representations.\n" + "This is an auto pipeline block that works for image conditioning tasks.\n" + " - `Flux2KleinBaseVaeEncoderSequentialStep` is used when `image` is provided.\n" + " - If `image` is not provided, step will be skipped." + ) + + +### +### Core denoise +### + +Flux2KleinBaseCoreDenoiseBlocks = InsertableDict( + [ + ("input", Flux2KleinBaseTextInputStep()), + ("prepare_latents", Flux2PrepareLatentsStep()), + ("set_timesteps", Flux2SetTimestepsStep()), + ("prepare_rope_inputs", Flux2KleinBaseRoPEInputsStep()), + ("denoise", Flux2KleinBaseDenoiseStep()), + ("after_denoise", Flux2UnpackLatentsStep()), + ] +) + + +# auto_docstring +class Flux2KleinBaseCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step that performs the denoising process for Flux2-Klein (base model), for text-to-image generation. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) guider + (`ClassifierFreeGuidance`) + + Configs: + is_distilled (default: False) + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + Pre-generated negative text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + image_latents (`Tensor`, *optional*): + Packed image latents for conditioning. Shape: (B, img_seq_len, C) + image_latent_ids (`Tensor`, *optional*): + Position IDs for image latents. Shape: (B, img_seq_len, 4) + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "flux2-klein" + block_classes = Flux2KleinBaseCoreDenoiseBlocks.values() + block_names = Flux2KleinBaseCoreDenoiseBlocks.keys() + + @property + def description(self): + return "Core denoise step that performs the denoising process for Flux2-Klein (base model), for text-to-image generation." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +Flux2KleinBaseImageConditionedCoreDenoiseBlocks = InsertableDict( + [ + ("input", Flux2KleinBaseTextInputStep()), + ("prepare_latents", Flux2PrepareLatentsStep()), + ("prepare_image_latents", Flux2PrepareImageLatentsStep()), + ("set_timesteps", Flux2SetTimestepsStep()), + ("prepare_rope_inputs", Flux2KleinBaseRoPEInputsStep()), + ("denoise", Flux2KleinBaseDenoiseStep()), + ("after_denoise", Flux2UnpackLatentsStep()), + ] +) + + +# auto_docstring +class Flux2KleinBaseImageConditionedCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step that performs the denoising process for Flux2-Klein (base model) with image conditioning. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) guider + (`ClassifierFreeGuidance`) + + Configs: + is_distilled (default: False) + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + Pre-generated negative text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + image_latents (`list`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "flux2-klein" + block_classes = Flux2KleinBaseImageConditionedCoreDenoiseBlocks.values() + block_names = Flux2KleinBaseImageConditionedCoreDenoiseBlocks.keys() + + @property + def description(self): + return "Core denoise step that performs the denoising process for Flux2-Klein (base model) with image conditioning." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +# auto_docstring +class Flux2KleinBaseAutoCoreDenoiseStep(AutoPipelineBlocks): + """ + Auto core denoise step that performs the denoising process for Flux2-Klein (base model). + This is an auto pipeline block that works for text-to-image and image-conditioned generation. + - `Flux2KleinBaseCoreDenoiseStep` is used for text-to-image generation. + - `Flux2KleinBaseImageConditionedCoreDenoiseStep` is used for image-conditioned generation. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) guider + (`ClassifierFreeGuidance`) + + Configs: + is_distilled (default: False) + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + Pre-generated negative text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + image_latents (`list`, *optional*): + TODO: Add description. + num_inference_steps (`None`): + TODO: Add description. + timesteps (`None`): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + image_latent_ids (`Tensor`, *optional*): + Position IDs for image latents. Shape: (B, img_seq_len, 4) + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "flux2-klein" + block_classes = [Flux2KleinBaseImageConditionedCoreDenoiseStep, Flux2KleinBaseCoreDenoiseStep] + block_names = ["image_conditioned", "text2image"] + block_trigger_inputs = ["image_latents", None] + + @property + def description(self): + return ( + "Auto core denoise step that performs the denoising process for Flux2-Klein (base model).\n" + "This is an auto pipeline block that works for text-to-image and image-conditioned generation.\n" + " - `Flux2KleinBaseCoreDenoiseStep` is used for text-to-image generation.\n" + " - `Flux2KleinBaseImageConditionedCoreDenoiseStep` is used for image-conditioned generation.\n" + ) + + +### +### Auto blocks +### + + +# auto_docstring +class Flux2KleinBaseAutoBlocks(SequentialPipelineBlocks): + """ + Auto blocks that perform the text-to-image and image-conditioned generation using Flux2-Klein (base model). + + Supported workflows: + - `text2image`: requires `prompt` + - `image_conditioned`: requires `image`, `prompt` + + Components: + text_encoder (`Qwen3ForCausalLM`) tokenizer (`Qwen2TokenizerFast`) guider (`ClassifierFreeGuidance`) + image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) scheduler + (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) + + Configs: + is_distilled (default: False) + + Inputs: + prompt (`None`, *optional*): + TODO: Add description. + max_sequence_length (`int`, *optional*, defaults to 512): + TODO: Add description. + text_encoder_out_layers (`tuple`, *optional*, defaults to (9, 18, 27)): + TODO: Add description. + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + latents (`Tensor | NoneType`): + TODO: Add description. + image_latents (`list`, *optional*): + TODO: Add description. + num_inference_steps (`None`): + TODO: Add description. + timesteps (`None`): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + image_latent_ids (`Tensor`, *optional*): + Position IDs for image latents. Shape: (B, img_seq_len, 4) + output_type (`None`, *optional*, defaults to pil): + TODO: Add description. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "flux2-klein" + block_classes = [ + Flux2KleinBaseTextEncoderStep(), + Flux2KleinBaseAutoVaeEncoderStep(), + Flux2KleinBaseAutoCoreDenoiseStep(), + Flux2DecodeStep(), + ] + block_names = ["text_encoder", "vae_encoder", "denoise", "decode"] + _workflow_map = { + "text2image": {"prompt": True}, + "image_conditioned": {"image": True, "prompt": True}, + } + + @property + def description(self): + return "Auto blocks that perform the text-to-image and image-conditioned generation using Flux2-Klein (base model)." + + @property + def outputs(self): + return [ + OutputParam.template("images"), + ] diff --git a/src/diffusers/modular_pipelines/flux2/modular_pipeline.py b/src/diffusers/modular_pipelines/flux2/modular_pipeline.py index 3e497f3b1e98..31ba5aec7cfb 100644 --- a/src/diffusers/modular_pipelines/flux2/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/flux2/modular_pipeline.py @@ -55,3 +55,45 @@ def num_channels_latents(self): if getattr(self, "transformer", None): num_channels_latents = self.transformer.config.in_channels // 4 return num_channels_latents + + +class Flux2KleinModularPipeline(Flux2ModularPipeline): + """ + A ModularPipeline for Flux2-Klein (distilled model). + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "Flux2KleinAutoBlocks" + + @property + def requires_unconditional_embeds(self): + if hasattr(self.config, "is_distilled") and self.config.is_distilled: + return False + + requires_unconditional_embeds = False + if hasattr(self, "guider") and self.guider is not None: + requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1 + + return requires_unconditional_embeds + + +class Flux2KleinBaseModularPipeline(Flux2ModularPipeline): + """ + A ModularPipeline for Flux2-Klein (base model). + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "Flux2KleinBaseAutoBlocks" + + @property + def requires_unconditional_embeds(self): + if hasattr(self.config, "is_distilled") and self.config.is_distilled: + return False + + requires_unconditional_embeds = False + if hasattr(self, "guider") and self.guider is not None: + requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1 + + return requires_unconditional_embeds diff --git a/src/diffusers/modular_pipelines/mellon_node_utils.py b/src/diffusers/modular_pipelines/mellon_node_utils.py index 4f142a453f3b..f65459dfc990 100644 --- a/src/diffusers/modular_pipelines/mellon_node_utils.py +++ b/src/diffusers/modular_pipelines/mellon_node_utils.py @@ -1,12 +1,13 @@ +import copy import json import logging import os # Simple typed wrapper for parameter overrides from dataclasses import asdict, dataclass -from typing import Any, Dict, Optional, Union +from typing import Any -from huggingface_hub import create_repo, hf_hub_download, upload_folder +from huggingface_hub import create_repo, hf_hub_download, upload_file from huggingface_hub.utils import ( EntryNotFoundError, HfHubHTTPError, @@ -15,243 +16,560 @@ ) from ..utils import HUGGINGFACE_CO_RESOLVE_ENDPOINT +from .modular_pipeline_utils import InputParam, OutputParam logger = logging.getLogger(__name__) +def _name_to_label(name: str) -> str: + """Convert snake_case name to Title Case label.""" + return name.replace("_", " ").title() + + +# Template definitions for standard diffuser pipeline parameters +MELLON_PARAM_TEMPLATES = { + # Image I/O + "image": {"label": "Image", "type": "image", "display": "input", "required_block_params": ["image"]}, + "images": {"label": "Images", "type": "image", "display": "output", "required_block_params": ["images"]}, + "control_image": { + "label": "Control Image", + "type": "image", + "display": "input", + "required_block_params": ["control_image"], + }, + # Latents + "latents": {"label": "Latents", "type": "latents", "display": "input", "required_block_params": ["latents"]}, + "image_latents": { + "label": "Image Latents", + "type": "latents", + "display": "input", + "required_block_params": ["image_latents"], + }, + "first_frame_latents": { + "label": "First Frame Latents", + "type": "latents", + "display": "input", + "required_block_params": ["first_frame_latents"], + }, + "latents_preview": {"label": "Latents Preview", "type": "latent", "display": "output"}, + # Image Latents with Strength + "image_latents_with_strength": { + "name": "image_latents", # name is not same as template key + "label": "Image Latents", + "type": "latents", + "display": "input", + "onChange": {"false": ["height", "width"], "true": ["strength"]}, + "required_block_params": ["image_latents", "strength"], + }, + # Embeddings + "embeddings": {"label": "Text Embeddings", "type": "embeddings", "display": "output"}, + "image_embeds": { + "label": "Image Embeddings", + "type": "image_embeds", + "display": "output", + "required_block_params": ["image_embeds"], + }, + # Text inputs + "prompt": { + "label": "Prompt", + "type": "string", + "display": "textarea", + "default": "", + "required_block_params": ["prompt"], + }, + "negative_prompt": { + "label": "Negative Prompt", + "type": "string", + "display": "textarea", + "default": "", + "required_block_params": ["negative_prompt"], + }, + # Numeric params + "guidance_scale": { + "label": "Guidance Scale", + "type": "float", + "display": "slider", + "default": 5.0, + "min": 1.0, + "max": 30.0, + "step": 0.1, + }, + "strength": { + "label": "Strength", + "type": "float", + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "required_block_params": ["strength"], + }, + "height": { + "label": "Height", + "type": "int", + "default": 1024, + "min": 64, + "step": 8, + "required_block_params": ["height"], + }, + "width": { + "label": "Width", + "type": "int", + "default": 1024, + "min": 64, + "step": 8, + "required_block_params": ["width"], + }, + "seed": { + "label": "Seed", + "type": "int", + "default": 0, + "min": 0, + "max": 4294967295, + "display": "random", + "required_block_params": ["generator"], + }, + "num_inference_steps": { + "label": "Steps", + "type": "int", + "default": 25, + "min": 1, + "max": 100, + "display": "slider", + "required_block_params": ["num_inference_steps"], + }, + "num_frames": { + "label": "Frames", + "type": "int", + "default": 81, + "min": 1, + "max": 480, + "display": "slider", + "required_block_params": ["num_frames"], + }, + "layers": { + "label": "Layers", + "type": "int", + "default": 4, + "min": 1, + "max": 10, + "display": "slider", + "required_block_params": ["layers"], + }, + "output_type": { + "label": "Output Type", + "type": "dropdown", + "default": "np", + "options": ["np", "pil", "pt"], + }, + # ControlNet + "controlnet_conditioning_scale": { + "label": "Controlnet Conditioning Scale", + "type": "float", + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "required_block_params": ["controlnet_conditioning_scale"], + }, + "control_guidance_start": { + "label": "Control Guidance Start", + "type": "float", + "default": 0.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "required_block_params": ["control_guidance_start"], + }, + "control_guidance_end": { + "label": "Control Guidance End", + "type": "float", + "default": 1.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "required_block_params": ["control_guidance_end"], + }, + # Video + "videos": {"label": "Videos", "type": "video", "display": "output", "required_block_params": ["videos"]}, + # Models + "vae": {"label": "VAE", "type": "diffusers_auto_model", "display": "input", "required_block_params": ["vae"]}, + "image_encoder": { + "label": "Image Encoder", + "type": "diffusers_auto_model", + "display": "input", + "required_block_params": ["image_encoder"], + }, + "unet": {"label": "Denoise Model", "type": "diffusers_auto_model", "display": "input"}, + "scheduler": {"label": "Scheduler", "type": "diffusers_auto_model", "display": "input"}, + "controlnet": { + "label": "ControlNet Model", + "type": "diffusers_auto_model", + "display": "input", + "required_block_params": ["controlnet"], + }, + "text_encoders": { + "label": "Text Encoders", + "type": "diffusers_auto_models", + "display": "input", + "required_block_params": ["text_encoder"], + }, + # Bundles/Custom + "controlnet_bundle": { + "label": "ControlNet", + "type": "custom_controlnet", + "display": "input", + "required_block_params": "controlnet_image", + }, + "ip_adapter": {"label": "IP Adapter", "type": "custom_ip_adapter", "display": "input"}, + "guider": { + "label": "Guider", + "type": "custom_guider", + "display": "input", + "onChange": {False: ["guidance_scale"], True: []}, + }, + "doc": {"label": "Doc", "type": "string", "display": "output"}, +} + + +class MellonParamMeta(type): + """Metaclass that enables MellonParam.template_name(**overrides) syntax.""" + + def __getattr__(cls, name: str): + if name in MELLON_PARAM_TEMPLATES: + + def factory(default=None, **overrides): + template = MELLON_PARAM_TEMPLATES[name] + # Use template's name if specified, otherwise use the key + params = {"name": template.get("name", name), **template, **overrides} + if default is not None: + params["default"] = default + return cls(**params) + + return factory + + raise AttributeError(f"type object 'MellonParam' has no attribute '{name}'") + + @dataclass(frozen=True) -class MellonParam: +class MellonParam(metaclass=MellonParamMeta): """ - Parameter definition for Mellon nodes. + Parameter definition for Mellon nodes. + + Usage: + ```python + # From template (standard diffuser params) + MellonParam.seed() + MellonParam.prompt(default="a cat") + MellonParam.latents(display="output") + + # Generic inputs (for custom blocks) + MellonParam.Input.slider("my_scale", default=1.0, min=0.0, max=2.0) + MellonParam.Input.dropdown("mode", options=["fast", "slow"]) + + # Generic outputs + MellonParam.Output.image("result_images") - Use factory methods for common params (e.g., MellonParam.seed()) or create custom ones with MellonParam(name="...", - label="...", type="..."). + # Fully custom + MellonParam(name="custom", label="Custom", type="float", default=0.5) + ``` """ name: str label: str type: str - display: Optional[str] = None + display: str | None = None default: Any = None - min: Optional[float] = None - max: Optional[float] = None - step: Optional[float] = None + min: float | None = None + max: float | None = None + step: float | None = None options: Any = None value: Any = None - fieldOptions: Optional[Dict[str, Any]] = None + fieldOptions: dict[str, Any] | None = None onChange: Any = None onSignal: Any = None + required_block_params: str | list[str] | None = None - def to_dict(self) -> Dict[str, Any]: - """Convert to dict for Mellon schema, excluding None values and name.""" + def to_dict(self) -> dict[str, Any]: + """Convert to dict for Mellon schema, excluding None values and internal fields.""" data = asdict(self) - return {k: v for k, v in data.items() if v is not None and k != "name"} - - @classmethod - def image(cls) -> "MellonParam": - return cls(name="image", label="Image", type="image", display="input") - - @classmethod - def images(cls) -> "MellonParam": - return cls(name="images", label="Images", type="image", display="output") - - @classmethod - def control_image(cls, display: str = "input") -> "MellonParam": - return cls(name="control_image", label="Control Image", type="image", display=display) - - @classmethod - def latents(cls, display: str = "input") -> "MellonParam": - return cls(name="latents", label="Latents", type="latents", display=display) - - @classmethod - def image_latents(cls, display: str = "input") -> "MellonParam": - return cls(name="image_latents", label="Image Latents", type="latents", display=display) - - @classmethod - def image_latents_with_strength(cls) -> "MellonParam": - return cls( - name="image_latents", - label="Image Latents", - type="latents", - display="input", - onChange={"false": ["height", "width"], "true": ["strength"]}, - ) - - @classmethod - def latents_preview(cls) -> "MellonParam": - """ - `Latents Preview` is a special output parameter that is used to preview the latents in the UI. - """ - return cls(name="latents_preview", label="Latents Preview", type="latent", display="output") - - @classmethod - def embeddings(cls, display: str = "output") -> "MellonParam": - return cls(name="embeddings", label="Text Embeddings", type="embeddings", display=display) - - @classmethod - def controlnet_conditioning_scale(cls, default: float = 0.5) -> "MellonParam": - return cls( - name="controlnet_conditioning_scale", - label="Controlnet Conditioning Scale", - type="float", - default=default, - min=0.0, - max=1.0, - step=0.01, - ) - - @classmethod - def control_guidance_start(cls, default: float = 0.0) -> "MellonParam": - return cls( - name="control_guidance_start", - label="Control Guidance Start", - type="float", - default=default, - min=0.0, - max=1.0, - step=0.01, - ) - - @classmethod - def control_guidance_end(cls, default: float = 1.0) -> "MellonParam": - return cls( - name="control_guidance_end", - label="Control Guidance End", - type="float", - default=default, - min=0.0, - max=1.0, - step=0.01, - ) - - @classmethod - def prompt(cls, default: str = "") -> "MellonParam": - return cls(name="prompt", label="Prompt", type="string", default=default, display="textarea") - - @classmethod - def negative_prompt(cls, default: str = "") -> "MellonParam": - return cls(name="negative_prompt", label="Negative Prompt", type="string", default=default, display="textarea") - - @classmethod - def strength(cls, default: float = 0.5) -> "MellonParam": - return cls(name="strength", label="Strength", type="float", default=default, min=0.0, max=1.0, step=0.01) - - @classmethod - def guidance_scale(cls, default: float = 5.0) -> "MellonParam": - return cls( - name="guidance_scale", - label="Guidance Scale", - type="float", - display="slider", - default=default, - min=1.0, - max=30.0, - step=0.1, - ) - - @classmethod - def height(cls, default: int = 1024) -> "MellonParam": - return cls(name="height", label="Height", type="int", default=default, min=64, step=8) - - @classmethod - def width(cls, default: int = 1024) -> "MellonParam": - return cls(name="width", label="Width", type="int", default=default, min=64, step=8) - - @classmethod - def seed(cls, default: int = 0) -> "MellonParam": - return cls(name="seed", label="Seed", type="int", default=default, min=0, max=4294967295, display="random") - - @classmethod - def num_inference_steps(cls, default: int = 25) -> "MellonParam": - return cls( - name="num_inference_steps", label="Steps", type="int", default=default, min=1, max=100, display="slider" - ) - - @classmethod - def vae(cls) -> "MellonParam": - """ - VAE model info dict. - - Contains keys like 'model_id', 'repo_id', 'execution_device' etc. Use components.get_one(model_id) to retrieve - the actual model. - """ - return cls(name="vae", label="VAE", type="diffusers_auto_model", display="input") - - @classmethod - def unet(cls) -> "MellonParam": - """ - Denoising model (UNet/Transformer) info dict. - - Contains keys like 'model_id', 'repo_id', 'execution_device' etc. Use components.get_one(model_id) to retrieve - the actual model. - """ - return cls(name="unet", label="Denoise Model", type="diffusers_auto_model", display="input") - - @classmethod - def scheduler(cls) -> "MellonParam": - """ - Scheduler model info dict. - - Contains keys like 'model_id', 'repo_id' etc. Use components.get_one(model_id) to retrieve the actual - scheduler. - """ - return cls(name="scheduler", label="Scheduler", type="diffusers_auto_model", display="input") - - @classmethod - def controlnet(cls) -> "MellonParam": - """ - ControlNet model info dict. - - Contains keys like 'model_id', 'repo_id', 'execution_device' etc. Use components.get_one(model_id) to retrieve - the actual model. - """ - return cls(name="controlnet", label="ControlNet Model", type="diffusers_auto_model", display="input") + return {k: v for k, v in data.items() if v is not None and k not in ("name", "required_block_params")} + + # ========================================================================= + # Input: Generic input parameter factories (for custom blocks) + # ========================================================================= + class Input: + """input UI elements for custom blocks.""" + + @classmethod + def image(cls, name: str) -> "MellonParam": + """image input.""" + return MellonParam(name=name, label=_name_to_label(name), type="image", display="input") + + @classmethod + def textbox(cls, name: str, default: str = "") -> "MellonParam": + """text input as textarea.""" + return MellonParam( + name=name, label=_name_to_label(name), type="string", display="textarea", default=default + ) - @classmethod - def text_encoders(cls) -> "MellonParam": - """ - Dict of text encoder model info dicts. + @classmethod + def dropdown(cls, name: str, options: list[str] = None, default: str = None) -> "MellonParam": + """dropdown selection.""" + if options and not default: + default = options[0] + if not default: + default = "" + if not options: + options = [default] + return MellonParam(name=name, label=_name_to_label(name), type="string", options=options, value=default) + + @classmethod + def slider( + cls, name: str, default: float = 0, min: float = None, max: float = None, step: float = None + ) -> "MellonParam": + """slider input.""" + is_float = isinstance(default, float) or (step is not None and isinstance(step, float)) + param_type = "float" if is_float else "int" + if min is None: + min = default + if max is None: + max = default + if step is None: + step = 0.01 if is_float else 1 + return MellonParam( + name=name, + label=_name_to_label(name), + type=param_type, + display="slider", + default=default, + min=min, + max=max, + step=step, + ) - Structure: { - 'text_encoder': {'model_id': ..., 'execution_device': ..., ...}, 'tokenizer': {'model_id': ..., ...}, - 'repo_id': '...' - } Use components.get_one(model_id) to retrieve each model. - """ - return cls(name="text_encoders", label="Text Encoders", type="diffusers_auto_models", display="input") + @classmethod + def number( + cls, name: str, default: float = 0, min: float = None, max: float = None, step: float = None + ) -> "MellonParam": + """number input (no slider).""" + is_float = isinstance(default, float) or (step is not None and isinstance(step, float)) + param_type = "float" if is_float else "int" + return MellonParam( + name=name, label=_name_to_label(name), type=param_type, default=default, min=min, max=max, step=step + ) - @classmethod - def controlnet_bundle(cls, display: str = "input") -> "MellonParam": - """ - ControlNet bundle containing model info and processed control inputs. + @classmethod + def seed(cls, name: str = "seed", default: int = 0) -> "MellonParam": + """seed input with randomize button.""" + return MellonParam( + name=name, + label=_name_to_label(name), + type="int", + display="random", + default=default, + min=0, + max=4294967295, + ) - Structure: { - 'controlnet': {'model_id': ..., ...}, # controlnet model info dict 'control_image': ..., # processed - control image/embeddings 'controlnet_conditioning_scale': ..., ... # other inputs expected by denoise - blocks - } + @classmethod + def checkbox(cls, name: str, default: bool = False) -> "MellonParam": + """boolean checkbox.""" + return MellonParam(name=name, label=_name_to_label(name), type="boolean", value=default) + + @classmethod + def custom_type(cls, name: str, type: str) -> "MellonParam": + """custom type input for node connections.""" + return MellonParam(name=name, label=_name_to_label(name), type=type, display="input") + + @classmethod + def model(cls, name: str) -> "MellonParam": + """model input for diffusers components.""" + return MellonParam(name=name, label=_name_to_label(name), type="diffusers_auto_model", display="input") + + # ========================================================================= + # Output: Generic output parameter factories (for custom blocks) + # ========================================================================= + class Output: + """output UI elements for custom blocks.""" + + @classmethod + def image(cls, name: str) -> "MellonParam": + """image output.""" + return MellonParam(name=name, label=_name_to_label(name), type="image", display="output") + + @classmethod + def video(cls, name: str) -> "MellonParam": + """video output.""" + return MellonParam(name=name, label=_name_to_label(name), type="video", display="output") + + @classmethod + def text(cls, name: str) -> "MellonParam": + """text output.""" + return MellonParam(name=name, label=_name_to_label(name), type="string", display="output") + + @classmethod + def custom_type(cls, name: str, type: str) -> "MellonParam": + """custom type output for node connections.""" + return MellonParam(name=name, label=_name_to_label(name), type=type, display="output") + + @classmethod + def model(cls, name: str) -> "MellonParam": + """model output for diffusers components.""" + return MellonParam(name=name, label=_name_to_label(name), type="diffusers_auto_model", display="output") + + +def input_param_to_mellon_param(input_param: "InputParam") -> MellonParam: + """ + Convert an InputParam to a MellonParam using metadata. - Output from Controlnet node, input to Denoise node. - """ - return cls(name="controlnet_bundle", label="ControlNet", type="custom_controlnet", display=display) + Args: + input_param: An InputParam with optional metadata containing either: + - {"mellon": ""} for simple types (image, textbox, slider, etc.) + - {"mellon": MellonParam(...)} for full control over UI configuration - @classmethod - def ip_adapter(cls) -> "MellonParam": - return cls(name="ip_adapter", label="IP Adapter", type="custom_ip_adapter", display="input") + Returns: + MellonParam instance + """ + name = input_param.name + metadata = input_param.metadata + mellon_value = metadata.get("mellon") if metadata else None + default = input_param.default + + # If it's already a MellonParam, return it directly + if isinstance(mellon_value, MellonParam): + return mellon_value + + mellon_type = mellon_value + + if mellon_type == "image": + return MellonParam.Input.image(name) + elif mellon_type == "textbox": + return MellonParam.Input.textbox(name, default=default or "") + elif mellon_type == "dropdown": + return MellonParam.Input.dropdown(name, default=default or "") + elif mellon_type == "slider": + return MellonParam.Input.slider(name, default=default or 0) + elif mellon_type == "number": + return MellonParam.Input.number(name, default=default or 0) + elif mellon_type == "seed": + return MellonParam.Input.seed(name, default=default or 0) + elif mellon_type == "checkbox": + return MellonParam.Input.checkbox(name, default=default or False) + elif mellon_type == "model": + return MellonParam.Input.model(name) + else: + # None or unknown -> custom + return MellonParam.Input.custom_type(name, type="custom") + + +def output_param_to_mellon_param(output_param: "OutputParam") -> MellonParam: + """ + Convert an OutputParam to a MellonParam using metadata. - @classmethod - def guider(cls) -> "MellonParam": - return cls( - name="guider", - label="Guider", - type="custom_guider", - display="input", - onChange={False: ["guidance_scale"], True: []}, - ) + Args: + output_param: An OutputParam with optional metadata={"mellon": ""} where type is one of: + image, video, text, model. If metadata is None or unknown, maps to "custom". - @classmethod - def doc(cls) -> "MellonParam": - return cls(name="doc", label="Doc", type="string", display="output") + Returns: + MellonParam instance + """ + name = output_param.name + metadata = output_param.metadata + mellon_type = metadata.get("mellon") if metadata else None + + if mellon_type == "image": + return MellonParam.Output.image(name) + elif mellon_type == "video": + return MellonParam.Output.video(name) + elif mellon_type == "text": + return MellonParam.Output.text(name) + elif mellon_type == "model": + return MellonParam.Output.model(name) + else: + # None or unknown -> custom + return MellonParam.Output.custom_type(name, type="custom") + + +DEFAULT_NODE_SPECS = { + "controlnet": None, + "denoise": { + "inputs": [ + MellonParam.embeddings(display="input"), + MellonParam.width(), + MellonParam.height(), + MellonParam.seed(), + MellonParam.num_inference_steps(), + MellonParam.num_frames(), + MellonParam.guidance_scale(), + MellonParam.strength(), + MellonParam.image_latents_with_strength(), + MellonParam.image_latents(), + MellonParam.first_frame_latents(), + MellonParam.controlnet_bundle(display="input"), + ], + "model_inputs": [ + MellonParam.unet(), + MellonParam.guider(), + MellonParam.scheduler(), + ], + "outputs": [ + MellonParam.latents(display="output"), + MellonParam.latents_preview(), + MellonParam.doc(), + ], + "required_inputs": ["embeddings"], + "required_model_inputs": ["unet", "scheduler"], + "block_name": "denoise", + }, + "vae_encoder": { + "inputs": [ + MellonParam.image(), + ], + "model_inputs": [ + MellonParam.vae(), + ], + "outputs": [ + MellonParam.image_latents(display="output"), + MellonParam.doc(), + ], + "required_inputs": ["image"], + "required_model_inputs": ["vae"], + "block_name": "vae_encoder", + }, + "text_encoder": { + "inputs": [ + MellonParam.prompt(), + MellonParam.negative_prompt(), + ], + "model_inputs": [ + MellonParam.text_encoders(), + ], + "outputs": [ + MellonParam.embeddings(display="output"), + MellonParam.doc(), + ], + "required_inputs": ["prompt"], + "required_model_inputs": ["text_encoders"], + "block_name": "text_encoder", + }, + "decoder": { + "inputs": [ + MellonParam.latents(display="input"), + ], + "model_inputs": [ + MellonParam.vae(), + ], + "outputs": [ + MellonParam.images(), + MellonParam.videos(), + MellonParam.doc(), + ], + "required_inputs": ["latents"], + "required_model_inputs": ["vae"], + "block_name": "decode", + }, +} def mark_required(label: str, marker: str = " *") -> str: @@ -261,7 +579,7 @@ def mark_required(label: str, marker: str = " *") -> str: return f"{label}{marker}" -def node_spec_to_mellon_dict(node_spec: Dict[str, Any], node_type: str) -> Dict[str, Any]: +def node_spec_to_mellon_dict(node_spec: dict[str, Any], node_type: str) -> dict[str, Any]: """ Convert a node spec dict into Mellon format. @@ -351,10 +669,15 @@ def node_spec_to_mellon_dict(node_spec: Dict[str, Any], node_type: str) -> Dict[ params[p.name] = param_dict model_input_names.append(p.name) - # Process outputs + # Process outputs: add a prefix to the output name if it already exists as an input for p in node_spec.get("outputs", []): - params[p.name] = p.to_dict() - output_names.append(p.name) + if p.name in input_names: + # rename to out_ + output_name = f"out_{p.name}" + else: + output_name = p.name + params[output_name] = p.to_dict() + output_names.append(output_name) return { "params": params, @@ -413,7 +736,7 @@ class MellonPipelineConfig: def __init__( self, - node_specs: Dict[str, Optional[Dict[str, Any]]], + node_specs: dict[str, dict[str, Any] | None], label: str = "", default_repo: str = "", default_dtype: str = "", @@ -428,22 +751,44 @@ def __init__( default_dtype: Default dtype (e.g., "float16", "bfloat16") """ # Convert all node specs to Mellon format immediately - self.node_params = {} - for node_type, spec in node_specs.items(): - if spec is None: - self.node_params[node_type] = None - else: - self.node_params[node_type] = node_spec_to_mellon_dict(spec, node_type) + self.node_specs = node_specs self.label = label self.default_repo = default_repo self.default_dtype = default_dtype - def __repr__(self) -> str: - node_types = list(self.node_params.keys()) - return f"MellonPipelineConfig(label={self.label!r}, default_repo={self.default_repo!r}, default_dtype={self.default_dtype!r}, node_params={node_types})" + @property + def node_params(self) -> dict[str, Any]: + """Lazily compute node_params from node_specs.""" + if self.node_specs is None: + return self._node_params - def to_dict(self) -> Dict[str, Any]: + params = {} + for node_type, spec in self.node_specs.items(): + if spec is None: + params[node_type] = None + else: + params[node_type] = node_spec_to_mellon_dict(spec, node_type) + return params + + def __repr__(self) -> str: + lines = [ + f"MellonPipelineConfig(label={self.label!r}, default_repo={self.default_repo!r}, default_dtype={self.default_dtype!r})" + ] + for node_type, spec in self.node_specs.items(): + if spec is None: + lines.append(f" {node_type}: None") + else: + inputs = [p.name for p in spec.get("inputs", [])] + model_inputs = [p.name for p in spec.get("model_inputs", [])] + outputs = [p.name for p in spec.get("outputs", [])] + lines.append(f" {node_type}:") + lines.append(f" inputs: {inputs}") + lines.append(f" model_inputs: {model_inputs}") + lines.append(f" outputs: {outputs}") + return "\n".join(lines) + + def to_dict(self) -> dict[str, Any]: """Convert to a JSON-serializable dictionary.""" return { "label": self.label, @@ -453,14 +798,15 @@ def to_dict(self) -> Dict[str, Any]: } @classmethod - def from_dict(cls, data: Dict[str, Any]) -> "MellonPipelineConfig": + def from_dict(cls, data: dict[str, Any]) -> "MellonPipelineConfig": """ Create from a dictionary (loaded from JSON). Note: The mellon_params are already in Mellon format when loading from JSON. """ instance = cls.__new__(cls) - instance.node_params = data.get("node_params", {}) + instance.node_specs = None + instance._node_params = data.get("node_params", {}) instance.label = data.get("label", "") instance.default_repo = data.get("default_repo", "") instance.default_dtype = data.get("default_dtype", "") @@ -470,20 +816,20 @@ def to_json_string(self) -> str: """Serialize to JSON string.""" return json.dumps(self.to_dict(), indent=2, sort_keys=False) + "\n" - def to_json_file(self, json_file_path: Union[str, os.PathLike]): + def to_json_file(self, json_file_path: str | os.PathLike): """Save to a JSON file.""" with open(json_file_path, "w", encoding="utf-8") as writer: writer.write(self.to_json_string()) @classmethod - def from_json_file(cls, json_file_path: Union[str, os.PathLike]) -> "MellonPipelineConfig": + def from_json_file(cls, json_file_path: str | os.PathLike) -> "MellonPipelineConfig": """Load from a JSON file.""" with open(json_file_path, "r", encoding="utf-8") as reader: data = json.load(reader) return cls.from_dict(data) - def save(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs): - """Save the pipeline config to a directory.""" + def save(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs): + """Save the mellon pipeline config to a directory.""" if os.path.isfile(save_directory): raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file") @@ -499,22 +845,21 @@ def save(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = Fals token = kwargs.pop("token", None) repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id - subfolder = kwargs.pop("subfolder", None) - upload_folder( + upload_file( + path_or_fileobj=output_path, + path_in_repo=self.config_name, repo_id=repo_id, - folder_path=save_directory, token=token, commit_message=commit_message or "Upload MellonPipelineConfig", create_pr=create_pr, - path_in_repo=subfolder, ) logger.info(f"Pipeline config pushed to hub: {repo_id}") @classmethod def load( cls, - pretrained_model_name_or_path: Union[str, os.PathLike], + pretrained_model_name_or_path: str | os.PathLike, **kwargs, ) -> "MellonPipelineConfig": """Load a pipeline config from a local path or Hugging Face Hub.""" @@ -592,3 +937,165 @@ def load( return cls.from_json_file(config_file) except (json.JSONDecodeError, UnicodeDecodeError): raise EnvironmentError(f"The config file at '{config_file}' is not a valid JSON file.") + + @classmethod + def from_blocks( + cls, + blocks, + template: dict[str, dict[str, Any]] | None = None, + label: str = "", + default_repo: str = "", + default_dtype: str = "bfloat16", + ) -> "MellonPipelineConfig": + """ + Create MellonPipelineConfig by matching template against actual pipeline blocks. + """ + if template is None: + template = DEFAULT_NODE_SPECS + + sub_block_map = dict(blocks.sub_blocks) + + def filter_spec_for_block(template_spec: dict[str, Any], block) -> dict[str, Any] | None: + """Filter template spec params based on what the block actually supports.""" + block_input_names = set(block.input_names) + block_output_names = set(block.intermediate_output_names) + block_component_names = set(block.component_names) + + filtered_inputs = [ + p + for p in template_spec.get("inputs", []) + if p.required_block_params is None + or all(name in block_input_names for name in p.required_block_params) + ] + filtered_model_inputs = [ + p + for p in template_spec.get("model_inputs", []) + if p.required_block_params is None + or all(name in block_component_names for name in p.required_block_params) + ] + filtered_outputs = [ + p + for p in template_spec.get("outputs", []) + if p.required_block_params is None + or all(name in block_output_names for name in p.required_block_params) + ] + + filtered_input_names = {p.name for p in filtered_inputs} + filtered_model_input_names = {p.name for p in filtered_model_inputs} + + filtered_required_inputs = [ + r for r in template_spec.get("required_inputs", []) if r in filtered_input_names + ] + filtered_required_model_inputs = [ + r for r in template_spec.get("required_model_inputs", []) if r in filtered_model_input_names + ] + + return { + "inputs": filtered_inputs, + "model_inputs": filtered_model_inputs, + "outputs": filtered_outputs, + "required_inputs": filtered_required_inputs, + "required_model_inputs": filtered_required_model_inputs, + "block_name": template_spec.get("block_name"), + } + + # Build node specs + node_specs = {} + for node_type, template_spec in template.items(): + if template_spec is None: + node_specs[node_type] = None + continue + + block_name = template_spec.get("block_name") + if block_name is None or block_name not in sub_block_map: + node_specs[node_type] = None + continue + + node_specs[node_type] = filter_spec_for_block(template_spec, sub_block_map[block_name]) + + return cls( + node_specs=node_specs, + label=label or getattr(blocks, "model_name", ""), + default_repo=default_repo, + default_dtype=default_dtype, + ) + + @classmethod + def from_custom_block( + cls, + block, + node_label: str = None, + input_types: dict[str, Any] | None = None, + output_types: dict[str, Any] | None = None, + ) -> "MellonPipelineConfig": + """ + Create a MellonPipelineConfig from a custom block. + + Args: + block: A block instance with `inputs`, `outputs`, and `expected_components`/`component_names` properties. + Each InputParam/OutputParam should have metadata={"mellon": ""} where type is one of: image, + video, text, checkbox, number, slider, dropdown, model. If metadata is None, maps to "custom". + node_label: The display label for the node. Defaults to block class name with spaces. + input_types: + Optional dict mapping input param names to mellon types. Overrides the block's metadata if provided. + Example: {"prompt": "textbox", "image": "image"} + output_types: + Optional dict mapping output param names to mellon types. Overrides the block's metadata if provided. + Example: {"prompt": "text", "images": "image"} + + Returns: + MellonPipelineConfig instance + """ + if node_label is None: + class_name = block.__class__.__name__ + node_label = "".join([" " + c if c.isupper() else c for c in class_name]).strip() + + if input_types is None: + input_types = {} + if output_types is None: + output_types = {} + + inputs = [] + model_inputs = [] + outputs = [] + + # Process block inputs + for input_param in block.inputs: + if input_param.name is None: + continue + if input_param.name in input_types: + input_param = copy.copy(input_param) + input_param.metadata = {"mellon": input_types[input_param.name]} + print(f" processing input: {input_param.name}, metadata: {input_param.metadata}") + inputs.append(input_param_to_mellon_param(input_param)) + + # Process block outputs + for output_param in block.outputs: + if output_param.name is None: + continue + if output_param.name in output_types: + output_param = copy.copy(output_param) + output_param.metadata = {"mellon": output_types[output_param.name]} + outputs.append(output_param_to_mellon_param(output_param)) + + # Process expected components (all map to model inputs) + component_names = block.component_names + for component_name in component_names: + model_inputs.append(MellonParam.Input.model(component_name)) + + # Always add doc output + outputs.append(MellonParam.doc()) + + node_spec = { + "inputs": inputs, + "model_inputs": model_inputs, + "outputs": outputs, + "required_inputs": [], + "required_model_inputs": [], + "block_name": "custom", + } + + return cls( + node_specs={"custom": node_spec}, + label=node_label, + ) diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index c5fa4cf9921f..a563d2aa99eb 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -14,12 +14,13 @@ import importlib import inspect import os +import sys import traceback import warnings from collections import OrderedDict from copy import deepcopy from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch from huggingface_hub import create_repo @@ -28,19 +29,31 @@ from typing_extensions import Self from ..configuration_utils import ConfigMixin, FrozenDict -from ..pipelines.pipeline_loading_utils import _fetch_class_library_tuple, simple_get_class_obj +from ..pipelines.pipeline_loading_utils import ( + LOADABLE_CLASSES, + _fetch_class_library_tuple, + _unwrap_model, + simple_get_class_obj, +) from ..utils import PushToHubMixin, is_accelerate_available, logging from ..utils.dynamic_modules_utils import get_class_from_dynamic_module, resolve_trust_remote_code from ..utils.hub_utils import load_or_create_model_card, populate_model_card +from ..utils.torch_utils import is_compiled_module from .components_manager import ComponentsManager from .modular_pipeline_utils import ( + MODULAR_MODEL_CARD_TEMPLATE, ComponentSpec, ConfigSpec, InputParam, InsertableDict, OutputParam, + _validate_requirements, + combine_inputs, + combine_outputs, format_components, format_configs, + format_workflow, + generate_modular_model_card_content, make_doc_string, ) @@ -52,17 +65,61 @@ # map regular pipeline to modular pipeline class name + + +def _create_default_map_fn(pipeline_class_name: str): + """Create a mapping function that always returns the same pipeline class.""" + + def _map_fn(config_dict=None): + return pipeline_class_name + + return _map_fn + + +def _flux2_klein_map_fn(config_dict=None): + if config_dict is None: + return "Flux2KleinModularPipeline" + + if "is_distilled" in config_dict and config_dict["is_distilled"]: + return "Flux2KleinModularPipeline" + else: + return "Flux2KleinBaseModularPipeline" + + +def _wan_map_fn(config_dict=None): + if config_dict is None: + return "WanModularPipeline" + + if "boundary_ratio" in config_dict and config_dict["boundary_ratio"] is not None: + return "Wan22ModularPipeline" + else: + return "WanModularPipeline" + + +def _wan_i2v_map_fn(config_dict=None): + if config_dict is None: + return "WanImage2VideoModularPipeline" + + if "boundary_ratio" in config_dict and config_dict["boundary_ratio"] is not None: + return "Wan22Image2VideoModularPipeline" + else: + return "WanImage2VideoModularPipeline" + + MODULAR_PIPELINE_MAPPING = OrderedDict( [ - ("stable-diffusion-xl", "StableDiffusionXLModularPipeline"), - ("wan", "WanModularPipeline"), - ("flux", "FluxModularPipeline"), - ("flux-kontext", "FluxKontextModularPipeline"), - ("flux2", "Flux2ModularPipeline"), - ("qwenimage", "QwenImageModularPipeline"), - ("qwenimage-edit", "QwenImageEditModularPipeline"), - ("qwenimage-edit-plus", "QwenImageEditPlusModularPipeline"), - ("z-image", "ZImageModularPipeline"), + ("stable-diffusion-xl", _create_default_map_fn("StableDiffusionXLModularPipeline")), + ("wan", _wan_map_fn), + ("wan-i2v", _wan_i2v_map_fn), + ("flux", _create_default_map_fn("FluxModularPipeline")), + ("flux-kontext", _create_default_map_fn("FluxKontextModularPipeline")), + ("flux2", _create_default_map_fn("Flux2ModularPipeline")), + ("flux2-klein", _flux2_klein_map_fn), + ("qwenimage", _create_default_map_fn("QwenImageModularPipeline")), + ("qwenimage-edit", _create_default_map_fn("QwenImageEditModularPipeline")), + ("qwenimage-edit-plus", _create_default_map_fn("QwenImageEditPlusModularPipeline")), + ("qwenimage-layered", _create_default_map_fn("QwenImageLayeredModularPipeline")), + ("z-image", _create_default_map_fn("ZImageModularPipeline")), ] ) @@ -73,8 +130,8 @@ class PipelineState: [`PipelineState`] stores the state of a pipeline. It is used to pass data between pipeline blocks. """ - values: Dict[str, Any] = field(default_factory=dict) - kwargs_mapping: Dict[str, List[str]] = field(default_factory=dict) + values: dict[str, Any] = field(default_factory=dict) + kwargs_mapping: dict[str, list[str]] = field(default_factory=dict) def set(self, key: str, value: Any, kwargs_type: str = None): """ @@ -93,22 +150,22 @@ def set(self, key: str, value: Any, kwargs_type: str = None): else: self.kwargs_mapping[kwargs_type].append(key) - def get(self, keys: Union[str, List[str]], default: Any = None) -> Union[Any, Dict[str, Any]]: + def get(self, keys: str | list[str], default: Any = None) -> Any | dict[str, Any]: """ Get one or multiple values from the pipeline state. Args: - keys (Union[str, List[str]]): Key or list of keys for the values + keys (str | list[str]): Key or list of keys for the values default (Any): The default value to return if not found Returns: - Union[Any, Dict[str, Any]]: Single value if keys is str, dictionary of values if keys is list + Any | dict[str, Any]: Single value if keys is str, dictionary of values if keys is list """ if isinstance(keys, str): return self.values.get(keys, default) return {key: self.values.get(key, default) for key in keys} - def get_by_kwargs(self, kwargs_type: str) -> Dict[str, Any]: + def get_by_kwargs(self, kwargs_type: str) -> dict[str, Any]: """ Get all values with matching kwargs_type. @@ -116,12 +173,12 @@ def get_by_kwargs(self, kwargs_type: str) -> Dict[str, Any]: kwargs_type (str): The kwargs_type to filter by Returns: - Dict[str, Any]: Dictionary of values with matching kwargs_type + dict[str, Any]: Dictionary of values with matching kwargs_type """ value_names = self.kwargs_mapping.get(kwargs_type, []) return self.get(value_names) - def to_dict(self) -> Dict[str, Any]: + def to_dict(self) -> dict[str, Any]: """ Convert PipelineState to a dictionary. """ @@ -180,7 +237,7 @@ def as_dict(self): Convert BlockState to a dictionary. Returns: - Dict[str, Any]: Dictionary containing all attributes of the BlockState + dict[str, Any]: Dictionary containing all attributes of the BlockState """ return dict(self.__dict__.items()) @@ -194,14 +251,14 @@ def format_value(v): elif isinstance(v, list): if len(v) > 0 and hasattr(v[0], "shape") and hasattr(v[0], "dtype"): shapes = [t.shape for t in v] - return f"List[{len(v)}] of Tensors with shapes {shapes}" + return f"list[{len(v)}] of Tensors with shapes {shapes}" return repr(v) # Handle tuples of tensors elif isinstance(v, tuple): if len(v) > 0 and hasattr(v[0], "shape") and hasattr(v[0], "dtype"): shapes = [t.shape for t in v] - return f"Tuple[{len(v)}] of Tensors with shapes {shapes}" + return f"tuple[{len(v)}] of Tensors with shapes {shapes}" return repr(v) # Handle dicts with tensor values @@ -217,7 +274,7 @@ def format_value(v): and hasattr(val[0], "dtype") ): shapes = [t.shape for t in val] - formatted_dict[k] = f"List[{len(val)}] of Tensors with shapes {shapes}" + formatted_dict[k] = f"list[{len(val)}] of Tensors with shapes {shapes}" else: formatted_dict[k] = repr(val) return formatted_dict @@ -231,7 +288,7 @@ def format_value(v): class ModularPipelineBlocks(ConfigMixin, PushToHubMixin): """ - Base class for all Pipeline Blocks: PipelineBlock, AutoPipelineBlocks, SequentialPipelineBlocks, + Base class for all Pipeline Blocks: ConditionalPipelineBlocks, AutoPipelineBlocks, SequentialPipelineBlocks, LoopSequentialPipelineBlocks [`ModularPipelineBlocks`] provides method to load and save the definition of pipeline blocks. @@ -241,6 +298,8 @@ class ModularPipelineBlocks(ConfigMixin, PushToHubMixin): config_name = "modular_config.json" model_name = None + _requirements: dict[str, str] | None = None + _workflow_map = None @classmethod def _get_signature_keys(cls, obj): @@ -260,16 +319,16 @@ def description(self) -> str: return "" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [] @property - def expected_configs(self) -> List[ConfigSpec]: + def expected_configs(self) -> list[ConfigSpec]: return [] @property - def inputs(self) -> List[InputParam]: - """List of input parameters. Must be implemented by subclasses.""" + def inputs(self) -> list[InputParam]: + """list of input parameters. Must be implemented by subclasses.""" return [] def _get_required_inputs(self): @@ -281,21 +340,50 @@ def _get_required_inputs(self): return input_names @property - def required_inputs(self) -> List[InputParam]: + def required_inputs(self) -> list[InputParam]: return self._get_required_inputs() @property - def intermediate_outputs(self) -> List[OutputParam]: - """List of intermediate output parameters. Must be implemented by subclasses.""" + def intermediate_outputs(self) -> list[OutputParam]: + """list of intermediate output parameters. Must be implemented by subclasses.""" return [] def _get_outputs(self): return self.intermediate_outputs @property - def outputs(self) -> List[OutputParam]: + def outputs(self) -> list[OutputParam]: return self._get_outputs() + # currentlyonly ConditionalPipelineBlocks and SequentialPipelineBlocks support `get_execution_blocks` + def get_execution_blocks(self, **kwargs): + """ + Get the block(s) that would execute given the inputs. Must be implemented by subclasses that support + conditional block selection. + + Args: + **kwargs: Input names and values. Only trigger inputs affect block selection. + """ + raise NotImplementedError(f"`get_execution_blocks` is not implemented for {self.__class__.__name__}") + + # currently only SequentialPipelineBlocks support workflows + @property + def available_workflows(self): + """ + Returns a list of available workflow names. Must be implemented by subclasses that define `_workflow_map`. + """ + raise NotImplementedError(f"`available_workflows` is not implemented for {self.__class__.__name__}") + + def get_workflow(self, workflow_name: str): + """ + Get the execution blocks for a specific workflow. Must be implemented by subclasses that define + `_workflow_map`. + + Args: + workflow_name: Name of the workflow to retrieve. + """ + raise NotImplementedError(f"`get_workflow` is not implemented for {self.__class__.__name__}") + @classmethod def from_pretrained( cls, @@ -325,6 +413,9 @@ def from_pretrained( "Selected model repository does not happear to have any custom code or does not have a valid `config.json` file." ) + if "requirements" in config and config["requirements"] is not None: + _ = _validate_requirements(config["requirements"]) + class_ref = config["auto_map"][cls.__name__] module_file, class_name = class_ref.split(".") module_file = module_file + ".py" @@ -349,22 +440,28 @@ def save_pretrained(self, save_directory, push_to_hub=False, **kwargs): module = full_mod.rsplit(".", 1)[-1].replace("__dynamic__", "") parent_module = self.save_pretrained.__func__.__qualname__.split(".", 1)[0] auto_map = {f"{parent_module}": f"{module}.{cls_name}"} - self.register_to_config(auto_map=auto_map) + + # resolve requirements + requirements = _validate_requirements(getattr(self, "_requirements", None)) + if requirements: + self.register_to_config(requirements=requirements) + self.save_config(save_directory=save_directory, push_to_hub=push_to_hub, **kwargs) config = dict(self.config) self._internal_dict = FrozenDict(config) def init_pipeline( self, - pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None, - components_manager: Optional[ComponentsManager] = None, - collection: Optional[str] = None, + pretrained_model_name_or_path: str | os.PathLike | None = None, + components_manager: ComponentsManager | None = None, + collection: str | None = None, ) -> "ModularPipeline": """ create a ModularPipeline, optionally accept pretrained_model_name_or_path to load from hub. """ - pipeline_class_name = MODULAR_PIPELINE_MAPPING.get(self.model_name, ModularPipeline.__name__) + map_fn = MODULAR_PIPELINE_MAPPING.get(self.model_name, _create_default_map_fn("ModularPipeline")) + pipeline_class_name = map_fn() diffusers_module = importlib.import_module("diffusers") pipeline_class = getattr(diffusers_module, pipeline_class_name) @@ -433,86 +530,20 @@ def set_block_state(self, state: PipelineState, block_state: BlockState): if current_value is not param: # Using identity comparison to check if object was modified state.set(param_name, param, input_param.kwargs_type) - @staticmethod - def combine_inputs(*named_input_lists: List[Tuple[str, List[InputParam]]]) -> List[InputParam]: - """ - Combines multiple lists of InputParam objects from different blocks. For duplicate inputs, updates only if - current default value is None and new default value is not None. Warns if multiple non-None default values - exist for the same input. - - Args: - named_input_lists: List of tuples containing (block_name, input_param_list) pairs - - Returns: - List[InputParam]: Combined list of unique InputParam objects - """ - combined_dict = {} # name -> InputParam - value_sources = {} # name -> block_name - - for block_name, inputs in named_input_lists: - for input_param in inputs: - if input_param.name is None and input_param.kwargs_type is not None: - input_name = "*_" + input_param.kwargs_type - else: - input_name = input_param.name - if input_name in combined_dict: - current_param = combined_dict[input_name] - if ( - current_param.default is not None - and input_param.default is not None - and current_param.default != input_param.default - ): - warnings.warn( - f"Multiple different default values found for input '{input_name}': " - f"{current_param.default} (from block '{value_sources[input_name]}') and " - f"{input_param.default} (from block '{block_name}'). Using {current_param.default}." - ) - if current_param.default is None and input_param.default is not None: - combined_dict[input_name] = input_param - value_sources[input_name] = block_name - else: - combined_dict[input_name] = input_param - value_sources[input_name] = block_name - - return list(combined_dict.values()) - - @staticmethod - def combine_outputs(*named_output_lists: List[Tuple[str, List[OutputParam]]]) -> List[OutputParam]: - """ - Combines multiple lists of OutputParam objects from different blocks. For duplicate outputs, keeps the first - occurrence of each output name. - - Args: - named_output_lists: List of tuples containing (block_name, output_param_list) pairs - - Returns: - List[OutputParam]: Combined list of unique OutputParam objects - """ - combined_dict = {} # name -> OutputParam - - for block_name, outputs in named_output_lists: - for output_param in outputs: - if (output_param.name not in combined_dict) or ( - combined_dict[output_param.name].kwargs_type is None and output_param.kwargs_type is not None - ): - combined_dict[output_param.name] = output_param - - return list(combined_dict.values()) - @property - def input_names(self) -> List[str]: + def input_names(self) -> list[str]: return [input_param.name for input_param in self.inputs if input_param.name is not None] @property - def intermediate_output_names(self) -> List[str]: + def intermediate_output_names(self) -> list[str]: return [output_param.name for output_param in self.intermediate_outputs if output_param.name is not None] @property - def output_names(self) -> List[str]: + def output_names(self) -> list[str]: return [output_param.name for output_param in self.outputs if output_param.name is not None] @property - def component_names(self) -> List[str]: + def component_names(self) -> list[str]: return [component.name for component in self.expected_components] @property @@ -527,9 +558,11 @@ def doc(self): ) -class AutoPipelineBlocks(ModularPipelineBlocks): +class ConditionalPipelineBlocks(ModularPipelineBlocks): """ - A Pipeline Blocks that automatically selects a block to run based on the inputs. + A Pipeline Blocks that conditionally selects a block to run based on the inputs. Subclasses must implement the + `select_block` method to define the logic for selecting the block. Currently, we only support selection logic based + on the presence or absence of inputs (i.e., whether they are `None` or not) This class inherits from [`ModularPipelineBlocks`]. Check the superclass documentation for the generic methods the library implements for all the pipeline blocks (such as loading or saving etc.) @@ -537,14 +570,20 @@ class AutoPipelineBlocks(ModularPipelineBlocks): > [!WARNING] > This is an experimental feature and is likely to change in the future. Attributes: - block_classes: List of block classes to be used - block_names: List of prefixes for each block - block_trigger_inputs: List of input names that trigger specific blocks, with None for default + block_classes: List of block classes to be used. Must have the same length as `block_names`. + block_names: List of names for each block. Must have the same length as `block_classes`. + block_trigger_inputs: List of input names that `select_block()` uses to determine which block to run. + For `ConditionalPipelineBlocks`, this does not need to correspond to `block_names` and `block_classes`. For + `AutoPipelineBlocks`, this must have the same length as `block_names` and `block_classes`, where each + element specifies the trigger input for the corresponding block. + default_block_name: Name of the default block to run when no trigger inputs match. + If None, this block can be skipped entirely when no trigger inputs are provided. """ block_classes = [] block_names = [] block_trigger_inputs = [] + default_block_name = None def __init__(self): sub_blocks = InsertableDict() @@ -554,26 +593,15 @@ def __init__(self): else: sub_blocks[block_name] = block self.sub_blocks = sub_blocks - if not (len(self.block_classes) == len(self.block_names) == len(self.block_trigger_inputs)): + if not (len(self.block_classes) == len(self.block_names)): raise ValueError( - f"In {self.__class__.__name__}, the number of block_classes, block_names, and block_trigger_inputs must be the same." + f"In {self.__class__.__name__}, the number of block_classes and block_names must be the same." ) - default_blocks = [t for t in self.block_trigger_inputs if t is None] - # can only have 1 or 0 default block, and has to put in the last - # the order of blocks matters here because the first block with matching trigger will be dispatched - # e.g. blocks = [inpaint, img2img] and block_trigger_inputs = ["mask", "image"] - # as long as mask is provided, it is inpaint; if only image is provided, it is img2img - if len(default_blocks) > 1 or (len(default_blocks) == 1 and self.block_trigger_inputs[-1] is not None): + if self.default_block_name is not None and self.default_block_name not in self.block_names: raise ValueError( - f"In {self.__class__.__name__}, exactly one None must be specified as the last element " - "in block_trigger_inputs." + f"In {self.__class__.__name__}, default_block_name '{self.default_block_name}' must be one of block_names: {self.block_names}" ) - # Map trigger inputs to block objects - self.trigger_to_block_map = dict(zip(self.block_trigger_inputs, self.sub_blocks.values())) - self.trigger_to_block_name_map = dict(zip(self.block_trigger_inputs, self.sub_blocks.keys())) - self.block_to_trigger_map = dict(zip(self.sub_blocks.keys(), self.block_trigger_inputs)) - @property def model_name(self): return next(iter(self.sub_blocks.values())).model_name @@ -601,9 +629,11 @@ def expected_configs(self): return expected_configs @property - def required_inputs(self) -> List[str]: - if None not in self.block_trigger_inputs: + def required_inputs(self) -> list[str]: + # no default block means this conditional block can be skipped entirely + if self.default_block_name is None: return [] + first_block = next(iter(self.sub_blocks.values())) required_by_all = set(getattr(first_block, "required_inputs", set())) @@ -614,11 +644,10 @@ def required_inputs(self) -> List[str]: return list(required_by_all) - # YiYi TODO: add test for this @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: named_inputs = [(name, block.inputs) for name, block in self.sub_blocks.items()] - combined_inputs = self.combine_inputs(*named_inputs) + combined_inputs = combine_inputs(*named_inputs) # mark Required inputs only if that input is required by all the blocks for input_param in combined_inputs: if input_param.name in self.required_inputs: @@ -628,33 +657,89 @@ def inputs(self) -> List[Tuple[str, Any]]: return combined_inputs @property - def intermediate_outputs(self) -> List[str]: + def intermediate_outputs(self) -> list[str]: named_outputs = [(name, block.intermediate_outputs) for name, block in self.sub_blocks.items()] - combined_outputs = self.combine_outputs(*named_outputs) + combined_outputs = combine_outputs(*named_outputs) return combined_outputs @property - def outputs(self) -> List[str]: + def outputs(self) -> list[str]: named_outputs = [(name, block.outputs) for name, block in self.sub_blocks.items()] - combined_outputs = self.combine_outputs(*named_outputs) + combined_outputs = combine_outputs(*named_outputs) return combined_outputs + @property + # Copied from diffusers.modular_pipelines.modular_pipeline.SequentialPipelineBlocks._requirements + def _requirements(self) -> dict[str, str]: + requirements = {} + for block_name, block in self.sub_blocks.items(): + if getattr(block, "_requirements", None): + requirements[block_name] = block._requirements + return requirements + + # used for `__repr__` + def _get_trigger_inputs(self) -> set: + """ + Returns a set of all unique trigger input values found in this block and nested blocks. + """ + + def fn_recursive_get_trigger(blocks): + trigger_values = set() + + if blocks is not None: + for name, block in blocks.items(): + # Check if current block has block_trigger_inputs + if hasattr(block, "block_trigger_inputs") and block.block_trigger_inputs is not None: + trigger_values.update(t for t in block.block_trigger_inputs if t is not None) + + # If block has sub_blocks, recursively check them + if block.sub_blocks: + nested_triggers = fn_recursive_get_trigger(block.sub_blocks) + trigger_values.update(nested_triggers) + + return trigger_values + + # Start with this block's block_trigger_inputs + all_triggers = {t for t in self.block_trigger_inputs if t is not None} + # Add nested triggers + all_triggers.update(fn_recursive_get_trigger(self.sub_blocks)) + + return all_triggers + + def select_block(self, **kwargs) -> str | None: + """ + Select the block to run based on the trigger inputs. Subclasses must implement this method to define the logic + for selecting the block. + + Note: When trigger inputs include intermediate outputs from earlier blocks, the selection logic should only + depend on the presence or absence of the input (i.e., whether it is None or not), not on its actual value. This + is because `get_execution_blocks()` resolves conditions statically by propagating intermediate output names + without their runtime values. + + Args: + **kwargs: Trigger input names and their values from the state. + + Returns: + str | None: The name of the block to run, or None to use default/skip. + """ + raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement the `select_block` method.") + @torch.no_grad() def __call__(self, pipeline, state: PipelineState) -> PipelineState: - # Find default block first (if any) + trigger_kwargs = {name: state.get(name) for name in self.block_trigger_inputs if name is not None} + block_name = self.select_block(**trigger_kwargs) - block = self.trigger_to_block_map.get(None) - for input_name in self.block_trigger_inputs: - if input_name is not None and state.get(input_name) is not None: - block = self.trigger_to_block_map[input_name] - break + if block_name is None: + block_name = self.default_block_name - if block is None: - logger.info(f"skipping auto block: {self.__class__.__name__}") + if block_name is None: + logger.info(f"skipping conditional block: {self.__class__.__name__}") return pipeline, state + block = self.sub_blocks[block_name] + try: - logger.info(f"Running block: {block.__class__.__name__}, trigger: {input_name}") + logger.info(f"Running block: {block.__class__.__name__}") return block(pipeline, state) except Exception as e: error_msg = ( @@ -665,37 +750,38 @@ def __call__(self, pipeline, state: PipelineState) -> PipelineState: logger.error(error_msg) raise - def _get_trigger_inputs(self): - """ - Returns a set of all unique trigger input values found in the blocks. Returns: Set[str] containing all unique - block_trigger_inputs values + def get_execution_blocks(self, **kwargs) -> ModularPipelineBlocks | None: """ + Get the block(s) that would execute given the inputs. - def fn_recursive_get_trigger(blocks): - trigger_values = set() + Recursively resolves nested ConditionalPipelineBlocks until reaching either: + - A leaf block (no sub_blocks or LoopSequentialPipelineBlocks) → returns single `ModularPipelineBlocks` + - A `SequentialPipelineBlocks` → delegates to its `get_execution_blocks()` which returns + a `SequentialPipelineBlocks` containing the resolved execution blocks - if blocks is not None: - for name, block in blocks.items(): - # Check if current block has trigger inputs(i.e. auto block) - if hasattr(block, "block_trigger_inputs") and block.block_trigger_inputs is not None: - # Add all non-None values from the trigger inputs list - trigger_values.update(t for t in block.block_trigger_inputs if t is not None) + Args: + **kwargs: Input names and values. Only trigger inputs affect block selection. - # If block has sub_blocks, recursively check them - if block.sub_blocks: - nested_triggers = fn_recursive_get_trigger(block.sub_blocks) - trigger_values.update(nested_triggers) + Returns: + - `ModularPipelineBlocks`: A leaf block or resolved `SequentialPipelineBlocks` + - `None`: If this block would be skipped (no trigger matched and no default) + """ + trigger_kwargs = {name: kwargs.get(name) for name in self.block_trigger_inputs if name is not None} + block_name = self.select_block(**trigger_kwargs) - return trigger_values + if block_name is None: + block_name = self.default_block_name + + if block_name is None: + return None - trigger_inputs = set(self.block_trigger_inputs) - trigger_inputs.update(fn_recursive_get_trigger(self.sub_blocks)) + block = self.sub_blocks[block_name] - return trigger_inputs + # Recursively resolve until we hit a leaf block + if block.sub_blocks and not isinstance(block, LoopSequentialPipelineBlocks): + return block.get_execution_blocks(**kwargs) - @property - def trigger_inputs(self): - return self._get_trigger_inputs() + return block def __repr__(self): class_name = self.__class__.__name__ @@ -704,11 +790,11 @@ def __repr__(self): f"{class_name}(\n Class: {base_class}\n" if base_class and base_class != "object" else f"{class_name}(\n" ) - if self.trigger_inputs: + if self._get_trigger_inputs(): header += "\n" header += " " + "=" * 100 + "\n" header += " This pipeline contains blocks that are selected at runtime based on inputs.\n" - header += f" Trigger Inputs: {[inp for inp in self.trigger_inputs if inp is not None]}\n" + header += f" Trigger Inputs: {sorted(self._get_trigger_inputs())}\n" header += " " + "=" * 100 + "\n\n" # Format description with proper indentation @@ -729,31 +815,20 @@ def __repr__(self): expected_configs = getattr(self, "expected_configs", []) configs_str = format_configs(expected_configs, indent_level=2, add_empty_lines=False) - # Blocks section - moved to the end with simplified format + # Blocks section blocks_str = " Sub-Blocks:\n" for i, (name, block) in enumerate(self.sub_blocks.items()): - # Get trigger input for this block - trigger = None - if hasattr(self, "block_to_trigger_map"): - trigger = self.block_to_trigger_map.get(name) - # Format the trigger info - if trigger is None: - trigger_str = "[default]" - elif isinstance(trigger, (list, tuple)): - trigger_str = f"[trigger: {', '.join(str(t) for t in trigger)}]" - else: - trigger_str = f"[trigger: {trigger}]" - # For AutoPipelineBlocks, add bullet points - blocks_str += f" • {name} {trigger_str} ({block.__class__.__name__})\n" + if name == self.default_block_name: + addtional_str = " [default]" else: - # For SequentialPipelineBlocks, show execution order - blocks_str += f" [{i}] {name} ({block.__class__.__name__})\n" + addtional_str = "" + blocks_str += f" • {name}{addtional_str} ({block.__class__.__name__})\n" # Add block description - desc_lines = block.description.split("\n") - indented_desc = desc_lines[0] - if len(desc_lines) > 1: - indented_desc += "\n" + "\n".join(" " + line for line in desc_lines[1:]) + block_desc_lines = block.description.split("\n") + indented_desc = block_desc_lines[0] + if len(block_desc_lines) > 1: + indented_desc += "\n" + "\n".join(" " + line for line in block_desc_lines[1:]) blocks_str += f" Description: {indented_desc}\n\n" # Build the representation with conditional sections @@ -784,6 +859,67 @@ def doc(self): ) +class AutoPipelineBlocks(ConditionalPipelineBlocks): + """ + A Pipeline Blocks that automatically selects a block to run based on the presence of trigger inputs. + + This is a specialized version of `ConditionalPipelineBlocks` where: + - Each block has one corresponding trigger input (1:1 mapping) + - Block selection is automatic: the first block whose trigger input is present gets selected + - `block_trigger_inputs` must have the same length as `block_names` and `block_classes` + - Use `None` in `block_trigger_inputs` to specify the default block, i.e the block that will run if no trigger + inputs are present + + Attributes: + block_classes: + List of block classes to be used. Must have the same length as `block_names` and + `block_trigger_inputs`. + block_names: + List of names for each block. Must have the same length as `block_classes` and `block_trigger_inputs`. + block_trigger_inputs: + List of input names where each element specifies the trigger input for the corresponding block. Use + `None` to mark the default block. + + Example: + ```python + class MyAutoBlock(AutoPipelineBlocks): + block_classes = [InpaintEncoderBlock, ImageEncoderBlock, TextEncoderBlock] + block_names = ["inpaint", "img2img", "text2img"] + block_trigger_inputs = ["mask_image", "image", None] # text2img is the default + ``` + + With this definition: + - As long as `mask_image` is provided, "inpaint" block runs (regardless of `image` being provided or not) + - If `mask_image` is not provided but `image` is provided, "img2img" block runs + - Otherwise, "text2img" block runs (default, trigger is `None`) + """ + + def __init__(self): + super().__init__() + + if self.default_block_name is not None: + raise ValueError( + f"In {self.__class__.__name__}, do not set `default_block_name` for AutoPipelineBlocks. " + f"Use `None` in `block_trigger_inputs` to specify the default block." + ) + + if not (len(self.block_classes) == len(self.block_names) == len(self.block_trigger_inputs)): + raise ValueError( + f"In {self.__class__.__name__}, the number of block_classes, block_names, and block_trigger_inputs must be the same." + ) + + if None in self.block_trigger_inputs: + idx = self.block_trigger_inputs.index(None) + self.default_block_name = self.block_names[idx] + + def select_block(self, **kwargs) -> str | None: + """Select block based on which trigger input is present (not None).""" + for trigger_input, block_name in zip(self.block_trigger_inputs, self.block_names): + if trigger_input is not None and kwargs.get(trigger_input) is not None: + return block_name + return None + + class SequentialPipelineBlocks(ModularPipelineBlocks): """ A Pipeline Blocks that combines multiple pipeline block classes into one. When called, it will call each block in @@ -795,8 +931,8 @@ class SequentialPipelineBlocks(ModularPipelineBlocks): > [!WARNING] > This is an experimental feature and is likely to change in the future. Attributes: - block_classes: List of block classes to be used - block_names: List of prefixes for each block + block_classes: list of block classes to be used + block_names: list of prefixes for each block """ block_classes = [] @@ -828,9 +964,32 @@ def expected_configs(self): expected_configs.append(config) return expected_configs + @property + def available_workflows(self): + if self._workflow_map is None: + raise NotImplementedError( + f"workflows is not supported because _workflow_map is not set for {self.__class__.__name__}" + ) + + return list(self._workflow_map.keys()) + + def get_workflow(self, workflow_name: str): + if self._workflow_map is None: + raise NotImplementedError( + f"workflows is not supported because _workflow_map is not set for {self.__class__.__name__}" + ) + + if workflow_name not in self._workflow_map: + raise ValueError(f"Workflow {workflow_name} not found in {self.__class__.__name__}") + + trigger_inputs = self._workflow_map[workflow_name] + workflow_blocks = self.get_execution_blocks(**trigger_inputs) + + return workflow_blocks + @classmethod def from_blocks_dict( - cls, blocks_dict: Dict[str, Any], description: Optional[str] = None + cls, blocks_dict: dict[str, Any], description: str | None = None ) -> "SequentialPipelineBlocks": """Creates a SequentialPipelineBlocks instance from a dictionary of blocks. @@ -885,7 +1044,8 @@ def _get_inputs(self): # Only add outputs if the block cannot be skipped should_add_outputs = True - if hasattr(block, "block_trigger_inputs") and None not in block.block_trigger_inputs: + if isinstance(block, ConditionalPipelineBlocks) and block.default_block_name is None: + # ConditionalPipelineBlocks without default can be skipped should_add_outputs = False if should_add_outputs: @@ -897,11 +1057,11 @@ def _get_inputs(self): # YiYi TODO: add test for this @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return self._get_inputs() @property - def required_inputs(self) -> List[str]: + def required_inputs(self) -> list[str]: # Get the first block from the dictionary first_block = next(iter(self.sub_blocks.values())) required_by_any = set(getattr(first_block, "required_inputs", set())) @@ -914,7 +1074,7 @@ def required_inputs(self) -> List[str]: return list(required_by_any) @property - def intermediate_outputs(self) -> List[str]: + def intermediate_outputs(self) -> list[str]: named_outputs = [] for name, block in self.sub_blocks.items(): inp_names = {inp.name for inp in block.inputs} @@ -922,12 +1082,12 @@ def intermediate_outputs(self) -> List[str]: # filter out them here so they do not end up as intermediate_outputs if name not in inp_names: named_outputs.append((name, block.intermediate_outputs)) - combined_outputs = self.combine_outputs(*named_outputs) + combined_outputs = combine_outputs(*named_outputs) return combined_outputs # YiYi TODO: I think we can remove the outputs property @property - def outputs(self) -> List[str]: + def outputs(self) -> list[str]: # return next(reversed(self.sub_blocks.values())).intermediate_outputs return self.intermediate_outputs @@ -946,10 +1106,10 @@ def __call__(self, pipeline, state: PipelineState) -> PipelineState: raise return pipeline, state + # used for `__repr__` def _get_trigger_inputs(self): """ - Returns a set of all unique trigger input values found in the blocks. Returns: Set[str] containing all unique - block_trigger_inputs values + Returns a set of all unique trigger input values found in the blocks. """ def fn_recursive_get_trigger(blocks): @@ -957,9 +1117,8 @@ def fn_recursive_get_trigger(blocks): if blocks is not None: for name, block in blocks.items(): - # Check if current block has trigger inputs(i.e. auto block) + # Check if current block has block_trigger_inputs (ConditionalPipelineBlocks) if hasattr(block, "block_trigger_inputs") and block.block_trigger_inputs is not None: - # Add all non-None values from the trigger inputs list trigger_values.update(t for t in block.block_trigger_inputs if t is not None) # If block has sub_blocks, recursively check them @@ -971,87 +1130,56 @@ def fn_recursive_get_trigger(blocks): return fn_recursive_get_trigger(self.sub_blocks) - @property - def trigger_inputs(self): - return self._get_trigger_inputs() + def get_execution_blocks(self, **kwargs) -> "SequentialPipelineBlocks": + """ + Get the blocks that would execute given the specified inputs. - def _traverse_trigger_blocks(self, trigger_inputs): - # Convert trigger_inputs to a set for easier manipulation - active_triggers = set(trigger_inputs) + As the traversal walks through sequential blocks, intermediate outputs from resolved blocks are added to the + active inputs. This means conditional blocks that depend on intermediates (e.g., "run img2img if image_latents + is present") will resolve correctly, as long as the condition is based on presence/absence (None or not None), + not on the actual value. - def fn_recursive_traverse(block, block_name, active_triggers): - result_blocks = OrderedDict() - # sequential(include loopsequential) or PipelineBlock - if not hasattr(block, "block_trigger_inputs"): - if block.sub_blocks: - # sequential or LoopSequentialPipelineBlocks (keep traversing) - for sub_block_name, sub_block in block.sub_blocks.items(): - blocks_to_update = fn_recursive_traverse(sub_block, sub_block_name, active_triggers) - blocks_to_update = fn_recursive_traverse(sub_block, sub_block_name, active_triggers) - blocks_to_update = {f"{block_name}.{k}": v for k, v in blocks_to_update.items()} - result_blocks.update(blocks_to_update) - else: - # PipelineBlock - result_blocks[block_name] = block - # Add this block's output names to active triggers if defined - if hasattr(block, "outputs"): - active_triggers.update(out.name for out in block.outputs) - return result_blocks - - # auto - else: - # Find first block_trigger_input that matches any value in our active_triggers - this_block = None - for trigger_input in block.block_trigger_inputs: - if trigger_input is not None and trigger_input in active_triggers: - this_block = block.trigger_to_block_map[trigger_input] - break + Args: + **kwargs: Input names and values. Only trigger inputs affect block selection. - # If no matches found, try to get the default (None) block - if this_block is None and None in block.block_trigger_inputs: - this_block = block.trigger_to_block_map[None] + Returns: + SequentialPipelineBlocks containing only the blocks that would execute + """ + # Copy kwargs so we can add outputs as we traverse + active_inputs = dict(kwargs) - if this_block is not None: - # sequential/auto (keep traversing) - if this_block.sub_blocks: - result_blocks.update(fn_recursive_traverse(this_block, block_name, active_triggers)) - else: - # PipelineBlock - result_blocks[block_name] = this_block - # Add this block's output names to active triggers if defined - # YiYi TODO: do we need outputs here? can it just be intermediate_outputs? can we get rid of outputs attribute? - if hasattr(this_block, "outputs"): - active_triggers.update(out.name for out in this_block.outputs) + def fn_recursive_traverse(block, block_name, active_inputs): + result_blocks = OrderedDict() + + # ConditionalPipelineBlocks (includes AutoPipelineBlocks) + if isinstance(block, ConditionalPipelineBlocks): + block = block.get_execution_blocks(**active_inputs) + if block is None: + return result_blocks + + # Has sub_blocks (SequentialPipelineBlocks/ConditionalPipelineBlocks) + if block.sub_blocks and not isinstance(block, LoopSequentialPipelineBlocks): + for sub_block_name, sub_block in block.sub_blocks.items(): + nested_blocks = fn_recursive_traverse(sub_block, sub_block_name, active_inputs) + nested_blocks = {f"{block_name}.{k}": v for k, v in nested_blocks.items()} + result_blocks.update(nested_blocks) + else: + # Leaf block: single ModularPipelineBlocks or LoopSequentialPipelineBlocks + result_blocks[block_name] = block + # Add outputs to active_inputs so subsequent blocks can use them as triggers + if hasattr(block, "intermediate_outputs"): + for out in block.intermediate_outputs: + active_inputs[out.name] = True return result_blocks all_blocks = OrderedDict() for block_name, block in self.sub_blocks.items(): - blocks_to_update = fn_recursive_traverse(block, block_name, active_triggers) - all_blocks.update(blocks_to_update) - return all_blocks - - def get_execution_blocks(self, *trigger_inputs): - trigger_inputs_all = self.trigger_inputs - - if trigger_inputs is not None: - if not isinstance(trigger_inputs, (list, tuple, set)): - trigger_inputs = [trigger_inputs] - invalid_inputs = [x for x in trigger_inputs if x not in trigger_inputs_all] - if invalid_inputs: - logger.warning( - f"The following trigger inputs will be ignored as they are not supported: {invalid_inputs}" - ) - trigger_inputs = [x for x in trigger_inputs if x in trigger_inputs_all] + nested_blocks = fn_recursive_traverse(block, block_name, active_inputs) + all_blocks.update(nested_blocks) - if trigger_inputs is None: - if None in trigger_inputs_all: - trigger_inputs = [None] - else: - trigger_inputs = [trigger_inputs_all[0]] - blocks_triggered = self._traverse_trigger_blocks(trigger_inputs) - return SequentialPipelineBlocks.from_blocks_dict(blocks_triggered) + return SequentialPipelineBlocks.from_blocks_dict(all_blocks) def __repr__(self): class_name = self.__class__.__name__ @@ -1060,18 +1188,23 @@ def __repr__(self): f"{class_name}(\n Class: {base_class}\n" if base_class and base_class != "object" else f"{class_name}(\n" ) - if self.trigger_inputs: + if self._workflow_map is None and self._get_trigger_inputs(): header += "\n" header += " " + "=" * 100 + "\n" header += " This pipeline contains blocks that are selected at runtime based on inputs.\n" - header += f" Trigger Inputs: {[inp for inp in self.trigger_inputs if inp is not None]}\n" + header += f" Trigger Inputs: {[inp for inp in self._get_trigger_inputs() if inp is not None]}\n" # Get first trigger input as example - example_input = next(t for t in self.trigger_inputs if t is not None) - header += f" Use `get_execution_blocks()` with input names to see selected blocks (e.g. `get_execution_blocks('{example_input}')`).\n" + example_input = next(t for t in self._get_trigger_inputs() if t is not None) + header += f" Use `get_execution_blocks()` to see selected blocks (e.g. `get_execution_blocks({example_input}=...)`).\n" header += " " + "=" * 100 + "\n\n" + description = self.description + if self._workflow_map is not None: + workflow_str = format_workflow(self._workflow_map) + description = f"{self.description}\n\n{workflow_str}" + # Format description with proper indentation - desc_lines = self.description.split("\n") + desc_lines = description.split("\n") desc = [] # First line with "Description:" label desc.append(f" Description: {desc_lines[0]}") @@ -1091,22 +1224,8 @@ def __repr__(self): # Blocks section - moved to the end with simplified format blocks_str = " Sub-Blocks:\n" for i, (name, block) in enumerate(self.sub_blocks.items()): - # Get trigger input for this block - trigger = None - if hasattr(self, "block_to_trigger_map"): - trigger = self.block_to_trigger_map.get(name) - # Format the trigger info - if trigger is None: - trigger_str = "[default]" - elif isinstance(trigger, (list, tuple)): - trigger_str = f"[trigger: {', '.join(str(t) for t in trigger)}]" - else: - trigger_str = f"[trigger: {trigger}]" - # For AutoPipelineBlocks, add bullet points - blocks_str += f" • {name} {trigger_str} ({block.__class__.__name__})\n" - else: - # For SequentialPipelineBlocks, show execution order - blocks_str += f" [{i}] {name} ({block.__class__.__name__})\n" + # show execution order + blocks_str += f" [{i}] {name} ({block.__class__.__name__})\n" # Add block description desc_lines = block.description.split("\n") @@ -1133,15 +1252,28 @@ def __repr__(self): @property def doc(self): + description = self.description + if self._workflow_map is not None: + workflow_str = format_workflow(self._workflow_map) + description = f"{self.description}\n\n{workflow_str}" + return make_doc_string( self.inputs, self.outputs, - self.description, + description=description, class_name=self.__class__.__name__, expected_components=self.expected_components, expected_configs=self.expected_configs, ) + @property + def _requirements(self) -> dict[str, str]: + requirements = {} + for block_name, block in self.sub_blocks.items(): + if getattr(block, "_requirements", None): + requirements[block_name] = block._requirements + return requirements + class LoopSequentialPipelineBlocks(ModularPipelineBlocks): """ @@ -1154,8 +1286,8 @@ class LoopSequentialPipelineBlocks(ModularPipelineBlocks): > [!WARNING] > This is an experimental feature and is likely to change in the future. Attributes: - block_classes: List of block classes to be used - block_names: List of prefixes for each block + block_classes: list of block classes to be used + block_names: list of prefixes for each block """ model_name = None @@ -1168,20 +1300,20 @@ def description(self) -> str: raise NotImplementedError("description method must be implemented in subclasses") @property - def loop_expected_components(self) -> List[ComponentSpec]: + def loop_expected_components(self) -> list[ComponentSpec]: return [] @property - def loop_expected_configs(self) -> List[ConfigSpec]: + def loop_expected_configs(self) -> list[ConfigSpec]: return [] @property - def loop_inputs(self) -> List[InputParam]: - """List of input parameters. Must be implemented by subclasses.""" + def loop_inputs(self) -> list[InputParam]: + """list of input parameters. Must be implemented by subclasses.""" return [] @property - def loop_required_inputs(self) -> List[str]: + def loop_required_inputs(self) -> list[str]: input_names = [] for input_param in self.loop_inputs: if input_param.required: @@ -1189,8 +1321,8 @@ def loop_required_inputs(self) -> List[str]: return input_names @property - def loop_intermediate_outputs(self) -> List[OutputParam]: - """List of intermediate output parameters. Must be implemented by subclasses.""" + def loop_intermediate_outputs(self) -> list[OutputParam]: + """list of intermediate output parameters. Must be implemented by subclasses.""" return [] # modified from SequentialPipelineBlocks to include loop_expected_components @@ -1230,15 +1362,9 @@ def _get_inputs(self): if inp.name not in outputs and inp not in inputs: inputs.append(inp) - # Only add outputs if the block cannot be skipped - should_add_outputs = True - if hasattr(block, "block_trigger_inputs") and None not in block.block_trigger_inputs: - should_add_outputs = False - - if should_add_outputs: - # Add this block's outputs - block_intermediate_outputs = [out.name for out in block.intermediate_outputs] - outputs.update(block_intermediate_outputs) + # Add this block's outputs + block_intermediate_outputs = [out.name for out in block.intermediate_outputs] + outputs.update(block_intermediate_outputs) for input_param in inputs: if input_param.name in self.required_inputs: @@ -1255,7 +1381,7 @@ def inputs(self): # modified from SequentialPipelineBlocks, if any additionan input required by the loop is required by the block @property - def required_inputs(self) -> List[str]: + def required_inputs(self) -> list[str]: # Get the first block from the dictionary first_block = next(iter(self.sub_blocks.values())) required_by_any = set(getattr(first_block, "required_inputs", set())) @@ -1273,9 +1399,9 @@ def required_inputs(self) -> List[str]: # YiYi TODO: this need to be thought about more # modified from SequentialPipelineBlocks to include loop_intermediate_outputs @property - def intermediate_outputs(self) -> List[str]: + def intermediate_outputs(self) -> list[str]: named_outputs = [(name, block.intermediate_outputs) for name, block in self.sub_blocks.items()] - combined_outputs = self.combine_outputs(*named_outputs) + combined_outputs = combine_outputs(*named_outputs) for output in self.loop_intermediate_outputs: if output.name not in {output.name for output in combined_outputs}: combined_outputs.append(output) @@ -1283,9 +1409,18 @@ def intermediate_outputs(self) -> List[str]: # YiYi TODO: this need to be thought about more @property - def outputs(self) -> List[str]: + def outputs(self) -> list[str]: return next(reversed(self.sub_blocks.values())).intermediate_outputs + @property + # Copied from diffusers.modular_pipelines.modular_pipeline.SequentialPipelineBlocks._requirements + def _requirements(self) -> dict[str, str]: + requirements = {} + for block_name, block in self.sub_blocks.items(): + if getattr(block, "_requirements", None): + requirements[block_name] = block._requirements + return requirements + def __init__(self): sub_blocks = InsertableDict() for block_name, block in zip(self.block_names, self.block_classes): @@ -1295,8 +1430,16 @@ def __init__(self): sub_blocks[block_name] = block self.sub_blocks = sub_blocks + # Validate that sub_blocks are only leaf blocks + for block_name, block in self.sub_blocks.items(): + if block.sub_blocks: + raise ValueError( + f"In {self.__class__.__name__}, sub_blocks must be leaf blocks (no sub_blocks). " + f"Block '{block_name}' ({block.__class__.__name__}) has sub_blocks." + ) + @classmethod - def from_blocks_dict(cls, blocks_dict: Dict[str, Any]) -> "LoopSequentialPipelineBlocks": + def from_blocks_dict(cls, blocks_dict: dict[str, Any]) -> "LoopSequentialPipelineBlocks": """ Creates a LoopSequentialPipelineBlocks instance from a dictionary of blocks. @@ -1447,12 +1590,12 @@ class ModularPipeline(ConfigMixin, PushToHubMixin): # YiYi TODO: add warning for passing multiple ComponentSpec/ConfigSpec with the same name def __init__( self, - blocks: Optional[ModularPipelineBlocks] = None, - pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None, - components_manager: Optional[ComponentsManager] = None, - collection: Optional[str] = None, - modular_config_dict: Optional[Dict[str, Any]] = None, - config_dict: Optional[Dict[str, Any]] = None, + blocks: ModularPipelineBlocks | None = None, + pretrained_model_name_or_path: str | os.PathLike | None = None, + components_manager: ComponentsManager | None = None, + collection: str | None = None, + modular_config_dict: dict[str, Any] | None = None, + config_dict: dict[str, Any] | None = None, **kwargs, ): """ @@ -1530,19 +1673,26 @@ def __init__( if modular_config_dict is not None: blocks_class_name = modular_config_dict.get("_blocks_class_name") else: - blocks_class_name = self.get_default_blocks_name(config_dict) + blocks_class_name = self.default_blocks_name if blocks_class_name is not None: diffusers_module = importlib.import_module("diffusers") - blocks_class = getattr(diffusers_module, blocks_class_name) + blocks_class = getattr(diffusers_module, blocks_class_name, None) + # If the blocks_class is not found or is a base class (e.g. SequentialPipelineBlocks saved by from_blocks_dict) with empty block_classes + # fall back to default_blocks_name + if blocks_class is None or not blocks_class.block_classes: + blocks_class_name = self.default_blocks_name + blocks_class = getattr(diffusers_module, blocks_class_name) + + if blocks_class is not None: blocks = blocks_class() else: logger.warning(f"`blocks` is `None`, no default blocks class found for {self.__class__.__name__}") - self.blocks = blocks + self._blocks = blocks self._components_manager = components_manager self._collection = collection - self._component_specs = {spec.name: deepcopy(spec) for spec in self.blocks.expected_components} - self._config_specs = {spec.name: deepcopy(spec) for spec in self.blocks.expected_configs} + self._component_specs = {spec.name: deepcopy(spec) for spec in self._blocks.expected_components} + self._config_specs = {spec.name: deepcopy(spec) for spec in self._blocks.expected_configs} # update component_specs and config_specs based on modular_model_index.json if modular_config_dict is not None: @@ -1589,26 +1739,27 @@ def __init__( for name, config_spec in self._config_specs.items(): default_configs[name] = config_spec.default self.register_to_config(**default_configs) - self.register_to_config(_blocks_class_name=self.blocks.__class__.__name__ if self.blocks is not None else None) + self.register_to_config( + _blocks_class_name=self._blocks.__class__.__name__ if self._blocks is not None else None + ) + + self._pretrained_model_name_or_path = pretrained_model_name_or_path @property - def default_call_parameters(self) -> Dict[str, Any]: + def default_call_parameters(self) -> dict[str, Any]: """ Returns: - Dictionary mapping input names to their default values """ params = {} - for input_param in self.blocks.inputs: + for input_param in self._blocks.inputs: params[input_param.name] = input_param.default return params - def get_default_blocks_name(self, config_dict: Optional[Dict[str, Any]]) -> Optional[str]: - return self.default_blocks_name - @classmethod def _load_pipeline_config( cls, - pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], + pretrained_model_name_or_path: str | os.PathLike | None, **load_config_kwargs, ): try: @@ -1638,10 +1789,10 @@ def _load_pipeline_config( @validate_hf_hub_args def from_pretrained( cls, - pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], - trust_remote_code: Optional[bool] = None, - components_manager: Optional[ComponentsManager] = None, - collection: Optional[str] = None, + pretrained_model_name_or_path: str | os.PathLike | None, + trust_remote_code: bool | None = None, + components_manager: ComponentsManager | None = None, + collection: str | None = None, **kwargs, ): """ @@ -1700,7 +1851,8 @@ def from_pretrained( logger.debug(" try to determine the modular pipeline class from model_index.json") standard_pipeline_class = _get_pipeline_class(cls, config=config_dict) model_name = _get_model(standard_pipeline_class.__name__) - pipeline_class_name = MODULAR_PIPELINE_MAPPING.get(model_name, ModularPipeline.__name__) + map_fn = MODULAR_PIPELINE_MAPPING.get(model_name, _create_default_map_fn("ModularPipeline")) + pipeline_class_name = map_fn(config_dict) diffusers_module = importlib.import_module("diffusers") pipeline_class = getattr(diffusers_module, pipeline_class_name) else: @@ -1719,34 +1871,136 @@ def from_pretrained( ) return pipeline - def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs): + def save_pretrained( + self, + save_directory: str | os.PathLike, + safe_serialization: bool = True, + variant: str | None = None, + max_shard_size: int | str | None = None, + push_to_hub: bool = False, + **kwargs, + ): """ - Save the pipeline to a directory. It does not save components, you need to save them separately. + Save the pipeline and all its components to a directory, so that it can be re-loaded using the + [`~ModularPipeline.from_pretrained`] class method. Args: save_directory (`str` or `os.PathLike`): - Path to the directory where the pipeline will be saved. - push_to_hub (`bool`, optional): - Whether to push the pipeline to the huggingface hub. - **kwargs: Additional arguments passed to `save_config()` method + Directory to save the pipeline to. Will be created if it doesn't exist. + safe_serialization (`bool`, *optional*, defaults to `True`): + Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. + variant (`str`, *optional*): + If specified, weights are saved in the format `pytorch_model..bin`. + max_shard_size (`int` or `str`, defaults to `None`): + The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size + lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5GB"`). + If expressed as an integer, the unit is bytes. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether to push the pipeline to the Hugging Face model hub after saving it. + **kwargs: Additional keyword arguments: + - `overwrite_modular_index` (`bool`, *optional*, defaults to `False`): + When saving a Modular Pipeline, its components in `modular_model_index.json` may reference repos + different from the destination repo. Setting this to `True` updates all component references in + `modular_model_index.json` so they point to the repo specified by `repo_id`. + - `repo_id` (`str`, *optional*): + The repository ID to push the pipeline to. Defaults to the last component of `save_directory`. + - `commit_message` (`str`, *optional*): + Commit message for the push to hub operation. + - `private` (`bool`, *optional*): + Whether the repository should be private. + - `create_pr` (`bool`, *optional*, defaults to `False`): + Whether to create a pull request instead of pushing directly. + - `token` (`str`, *optional*): + The Hugging Face token to use for authentication. """ + overwrite_modular_index = kwargs.pop("overwrite_modular_index", False) + repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) + if push_to_hub: commit_message = kwargs.pop("commit_message", None) private = kwargs.pop("private", None) create_pr = kwargs.pop("create_pr", False) token = kwargs.pop("token", None) - repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) + update_model_card = kwargs.pop("update_model_card", False) repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id - # Create a new empty model card and eventually tag it - model_card = load_or_create_model_card(repo_id, token=token, is_pipeline=True) - model_card = populate_model_card(model_card) - model_card.save(os.path.join(save_directory, "README.md")) + for component_name, component_spec in self._component_specs.items(): + if component_spec.default_creation_method != "from_pretrained": + continue + + component = getattr(self, component_name, None) + if component is None: + continue + + model_cls = component.__class__ + if is_compiled_module(component): + component = _unwrap_model(component) + model_cls = component.__class__ + + save_method_name = None + for library_name, library_classes in LOADABLE_CLASSES.items(): + if library_name in sys.modules: + library = importlib.import_module(library_name) + else: + logger.info( + f"{library_name} is not installed. Cannot save {component_name} as {library_classes} from {library_name}" + ) + continue + + for base_class, save_load_methods in library_classes.items(): + class_candidate = getattr(library, base_class, None) + if class_candidate is not None and issubclass(model_cls, class_candidate): + save_method_name = save_load_methods[0] + break + if save_method_name is not None: + break + + if save_method_name is None: + logger.warning(f"self.{component_name}={component} of type {type(component)} cannot be saved.") + continue + + save_method = getattr(component, save_method_name) + save_method_signature = inspect.signature(save_method) + save_method_accept_safe = "safe_serialization" in save_method_signature.parameters + save_method_accept_variant = "variant" in save_method_signature.parameters + save_method_accept_max_shard_size = "max_shard_size" in save_method_signature.parameters + + save_kwargs = {} + if save_method_accept_safe: + save_kwargs["safe_serialization"] = safe_serialization + if save_method_accept_variant: + save_kwargs["variant"] = variant + if save_method_accept_max_shard_size and max_shard_size is not None: + save_kwargs["max_shard_size"] = max_shard_size + + component_save_path = os.path.join(save_directory, component_name) + save_method(component_save_path, **save_kwargs) + + if component_name not in self.config: + continue + + has_no_load_id = not hasattr(component, "_diffusers_load_id") or component._diffusers_load_id == "null" + if overwrite_modular_index or has_no_load_id: + library, class_name, component_spec_dict = self.config[component_name] + component_spec_dict["pretrained_model_name_or_path"] = repo_id if push_to_hub else save_directory + component_spec_dict["subfolder"] = component_name + self.register_to_config(**{component_name: (library, class_name, component_spec_dict)}) - # YiYi TODO: maybe order the json file to make it more readable: configs first, then components self.save_config(save_directory=save_directory) if push_to_hub: + card_content = generate_modular_model_card_content(self.blocks) + model_card = load_or_create_model_card( + repo_id, + token=token, + is_pipeline=True, + model_description=MODULAR_MODEL_CARD_TEMPLATE.format(**card_content), + is_modular=True, + update_model_card=update_model_card, + ) + model_card = populate_model_card(model_card, tags=card_content["tags"]) + model_card.save(os.path.join(save_directory, "README.md")) + self._upload_folder( save_directory, repo_id, @@ -1761,7 +2015,15 @@ def doc(self): Returns: - The docstring of the pipeline blocks """ - return self.blocks.doc + return self._blocks.doc + + @property + def blocks(self) -> ModularPipelineBlocks: + """ + Returns: + - A copy of the pipeline blocks + """ + return deepcopy(self._blocks) def register_components(self, **kwargs): """ @@ -1919,26 +2181,26 @@ def dtype(self) -> torch.dtype: return torch.float32 @property - def null_component_names(self) -> List[str]: + def null_component_names(self) -> list[str]: """ Returns: - - List of names for components that needs to be loaded + - list of names for components that needs to be loaded """ return [name for name in self._component_specs.keys() if hasattr(self, name) and getattr(self, name) is None] @property - def component_names(self) -> List[str]: + def component_names(self) -> list[str]: """ Returns: - - List of names for all components + - list of names for all components """ return list(self.components.keys()) @property - def pretrained_component_names(self) -> List[str]: + def pretrained_component_names(self) -> list[str]: """ Returns: - - List of names for from_pretrained components + - list of names for from_pretrained components """ return [ name @@ -1947,10 +2209,10 @@ def pretrained_component_names(self) -> List[str]: ] @property - def config_component_names(self) -> List[str]: + def config_component_names(self) -> list[str]: """ Returns: - - List of names for from_config components + - list of names for from_config components """ return [ name @@ -1959,7 +2221,7 @@ def config_component_names(self) -> List[str]: ] @property - def components(self) -> Dict[str, Any]: + def components(self) -> dict[str, Any]: """ Returns: - Dictionary mapping component names to their objects (include both from_pretrained and from_config @@ -1989,58 +2251,30 @@ def update_components(self, **kwargs): - the `config` dict, which will be saved as `modular_model_index.json` during `save_pretrained` Args: - **kwargs: Component objects, ComponentSpec objects, or configuration values to update: - - Component objects: Only supports components we can extract specs using - `ComponentSpec.from_component()` method i.e. components created with ComponentSpec.load() or - ConfigMixin subclasses that aren't nn.Modules (e.g., `unet=new_unet, text_encoder=new_encoder`) - - ComponentSpec objects: Only supports default_creation_method == "from_config", will call create() - method to create a new component (e.g., `guider=ComponentSpec(name="guider", - type_hint=ClassifierFreeGuidance, config={...}, default_creation_method="from_config")`) - - Configuration values: Simple values to update configuration settings (e.g., - `requires_safety_checker=False`) - - Raises: - ValueError: If a component object is not supported in ComponentSpec.from_component() method: - - nn.Module components without a valid `_diffusers_load_id` attribute - - Non-ConfigMixin components without a valid `_diffusers_load_id` attribute + **kwargs: Component objects or configuration values to update: + - Component objects: Models loaded with `AutoModel.from_pretrained()` or `ComponentSpec.load()` + are automatically tagged with loading information. ConfigMixin objects without weights (e.g., + schedulers, guiders) can be passed directly. + - Configuration values: Simple values to update configuration settings + (e.g., `requires_safety_checker=False`) Examples: ```python - # Update multiple components at once + # Update pre-trained model pipeline.update_components(unet=new_unet_model, text_encoder=new_text_encoder) # Update configuration values pipeline.update_components(requires_safety_checker=False) - - # Update both components and configs together - pipeline.update_components(unet=new_unet_model, requires_safety_checker=False) - - # Update with ComponentSpec objects (from_config only) - pipeline.update_components( - guider=ComponentSpec( - name="guider", - type_hint=ClassifierFreeGuidance, - config={"guidance_scale": 5.0}, - default_creation_method="from_config", - ) - ) ``` Notes: - - Components with trained weights must be created using ComponentSpec.load(). If the component has not been - shared in huggingface hub and you don't have loading specs, you can upload it using `push_to_hub()` - - ConfigMixin objects without weights (e.g., schedulers, guiders) can be passed directly - - ComponentSpec objects with default_creation_method="from_pretrained" are not supported in - update_components() + - Components loaded with `AutoModel.from_pretrained()` or `ComponentSpec.load()` will have + loading specs preserved for serialization. Custom or locally loaded components without Hub references will + have their `modular_model_index.json` entries updated automatically during `save_pretrained()`. + - ConfigMixin objects without weights (e.g., schedulers, guiders) can be passed directly. """ - # extract component_specs_updates & config_specs_updates from `specs` - passed_component_specs = { - k: kwargs.pop(k) for k in self._component_specs if k in kwargs and isinstance(kwargs[k], ComponentSpec) - } - passed_components = { - k: kwargs.pop(k) for k in self._component_specs if k in kwargs and not isinstance(kwargs[k], ComponentSpec) - } + passed_components = {k: kwargs.pop(k) for k in self._component_specs if k in kwargs} passed_config_values = {k: kwargs.pop(k) for k in self._config_specs if k in kwargs} for name, component in passed_components.items(): @@ -2058,13 +2292,10 @@ def update_components(self, **kwargs): new_component_spec = current_component_spec if hasattr(self, name) and getattr(self, name) is not None: logger.warning(f"ModularPipeline.update_components: setting {name} to None (spec unchanged)") - elif current_component_spec.default_creation_method == "from_pretrained" and not ( - hasattr(component, "_diffusers_load_id") and component._diffusers_load_id is not None + elif ( + current_component_spec.default_creation_method == "from_pretrained" + and getattr(component, "_diffusers_load_id", None) is None ): - logger.warning( - f"ModularPipeline.update_components: {name} has no valid _diffusers_load_id. " - f"This will result in empty loading spec, use ComponentSpec.load() for proper specs" - ) new_component_spec = ComponentSpec(name=name, type_hint=type(component)) else: new_component_spec = ComponentSpec.from_component(name, component) @@ -2079,39 +2310,20 @@ def update_components(self, **kwargs): if len(kwargs) > 0: logger.warning(f"Unexpected keyword arguments, will be ignored: {kwargs.keys()}") - created_components = {} - for name, component_spec in passed_component_specs.items(): - if component_spec.default_creation_method == "from_pretrained": - raise ValueError( - "ComponentSpec object with default_creation_method == 'from_pretrained' is not supported in update_components() method" - ) - created_components[name] = component_spec.create() - current_component_spec = self._component_specs[name] - # warn if type changed - if current_component_spec.type_hint is not None and not isinstance( - created_components[name], current_component_spec.type_hint - ): - logger.info( - f"ModularPipeline.update_components: adding {name} with new type: {created_components[name].__class__.__name__}, previous type: {current_component_spec.type_hint.__name__}" - ) - # update _component_specs based on the user passed component_spec - self._component_specs[name] = component_spec - self.register_components(**passed_components, **created_components) + self.register_components(**passed_components) config_to_register = {} for name, new_value in passed_config_values.items(): - # e.g. requires_aesthetics_score = False self._config_specs[name].default = new_value config_to_register[name] = new_value self.register_to_config(**config_to_register) - # YiYi TODO: support map for additional from_pretrained kwargs - def load_components(self, names: Optional[Union[List[str], str]] = None, **kwargs): + def load_components(self, names: list[str] | str | None = None, **kwargs): """ Load selected components from specs. Args: - names: List of component names to load. If None, will load all components with + names: list of component names to load. If None, will load all components with default_creation_method == "from_pretrained". If provided as a list or string, will load only the specified components. **kwargs: additional kwargs to be passed to `from_pretrained()`.Can be: @@ -2128,6 +2340,8 @@ def load_components(self, names: Optional[Union[List[str], str]] = None, **kwarg name for name in self._component_specs.keys() if self._component_specs[name].default_creation_method == "from_pretrained" + and self._component_specs[name].pretrained_model_name_or_path is not None + and getattr(self, name, None) is None ] elif isinstance(names, str): names = [names] @@ -2154,24 +2368,56 @@ def load_components(self, names: Optional[Union[List[str], str]] = None, **kwarg elif "default" in value: # check if the default is specified component_load_kwargs[key] = value["default"] + # Only pass trust_remote_code to components from the same repo as the pipeline. + # When a user passes trust_remote_code=True, they intend to trust code from the + # pipeline's repo, not from external repos referenced in modular_model_index.json. + trust_remote_code_stripped = False + if ( + "trust_remote_code" in component_load_kwargs + and self._pretrained_model_name_or_path is not None + and spec.pretrained_model_name_or_path != self._pretrained_model_name_or_path + ): + component_load_kwargs.pop("trust_remote_code") + trust_remote_code_stripped = True + + if not spec.pretrained_model_name_or_path: + logger.info(f"Skipping component `{name}`: no pretrained model path specified.") + continue + try: components_to_register[name] = spec.load(**component_load_kwargs) except Exception: - logger.warning( - f"\nFailed to create component {name}:\n" - f"- Component spec: {spec}\n" - f"- load() called with kwargs: {component_load_kwargs}\n" - "If this component is not required for your workflow you can safely ignore this message.\n\n" - "Traceback:\n" - f"{traceback.format_exc()}" - ) + tb = traceback.format_exc() + if trust_remote_code_stripped and "trust_remote_code" in tb: + warning_msg = ( + f"Failed to load component `{name}` from external repository " + f"`{spec.pretrained_model_name_or_path}`.\n\n" + f"`trust_remote_code=True` was not forwarded to `{name}` because it comes from " + f"a different repository than the pipeline (`{self._pretrained_model_name_or_path}`). " + f"For safety, `trust_remote_code` is only forwarded to components from the same " + f"repository as the pipeline.\n\n" + f"You need to load this component manually with `trust_remote_code=True` and pass it " + f"to the pipeline via `pipe.update_components()`. For example, if it is a custom model:\n\n" + f' {name} = AutoModel.from_pretrained("{spec.pretrained_model_name_or_path}", trust_remote_code=True)\n' + f" pipe.update_components({name}={name})\n" + ) + else: + warning_msg = ( + f"Failed to create component {name}:\n" + f"- Component spec: {spec}\n" + f"- load() called with kwargs: {component_load_kwargs}\n" + "If this component is not required for your workflow you can safely ignore this message.\n\n" + "Traceback:\n" + f"{tb}" + ) + logger.warning(warning_msg) # Register all components at once self.register_components(**components_to_register) # Copied from diffusers.pipelines.pipeline_utils.DiffusionPipeline._maybe_raise_error_if_group_offload_active def _maybe_raise_error_if_group_offload_active( - self, raise_error: bool = False, module: Optional[torch.nn.Module] = None + self, raise_error: bool = False, module: torch.nn.Module | None = None ) -> bool: from ..hooks.group_offloading import _is_group_offload_enabled @@ -2383,16 +2629,16 @@ def _component_spec_to_dict(component_spec: ComponentSpec) -> Any: the `default_creation_method` is not `from_pretrained`, return None. This dict contains: - - "type_hint": Tuple[str, str] + - "type_hint": tuple[str, str] Library name and class name of the component. (e.g. ("diffusers", "UNet2DConditionModel")) - All loading fields defined by `component_spec.loading_fields()`, typically: - - "pretrained_model_name_or_path": Optional[str] + - "pretrained_model_name_or_path": str | None The model pretrained_model_name_or_pathsitory (e.g., "stabilityai/stable-diffusion-xl"). - - "subfolder": Optional[str] + - "subfolder": str | None A subfolder within the pretrained_model_name_or_path where this component lives. - - "variant": Optional[str] + - "variant": str | None An optional variant identifier for the model. - - "revision": Optional[str] + - "revision": str | None A specific git revision (commit hash, tag, or branch). - ... any other loading fields defined on the spec. @@ -2401,7 +2647,7 @@ def _component_spec_to_dict(component_spec: ComponentSpec) -> Any: The spec object describing one pipeline component. Returns: - Dict[str, Any]: A mapping suitable for JSON serialization. + dict[str, Any]: A mapping suitable for JSON serialization. Example: >>> from diffusers.pipelines.modular_pipeline_utils import ComponentSpec >>> from diffusers import @@ -2431,31 +2677,28 @@ def _component_spec_to_dict(component_spec: ComponentSpec) -> Any: } @staticmethod - def _dict_to_component_spec( - name: str, - spec_dict: Dict[str, Any], - ) -> ComponentSpec: + def _dict_to_component_spec(name: str, spec_dict: dict[str, Any]) -> ComponentSpec: """ Reconstruct a ComponentSpec from a loading specdict. This method converts a dictionary representation back into a ComponentSpec object. The dict should contain: - - "type_hint": Tuple[str, str] + - "type_hint": tuple[str, str] Library name and class name of the component. (e.g. ("diffusers", "UNet2DConditionModel")) - All loading fields defined by `component_spec.loading_fields()`, typically: - - "pretrained_model_name_or_path": Optional[str] + - "pretrained_model_name_or_path": str | None The model repository (e.g., "stabilityai/stable-diffusion-xl"). - - "subfolder": Optional[str] + - "subfolder": str | None A subfolder within the pretrained_model_name_or_path where this component lives. - - "variant": Optional[str] + - "variant": str | None An optional variant identifier for the model. - - "revision": Optional[str] + - "revision": str | None A specific git revision (commit hash, tag, or branch). - ... any other loading fields defined on the spec. Args: name (str): The name of the component. - specdict (Dict[str, Any]): + specdict (dict[str, Any]): A dictionary containing the component specification data. Returns: @@ -2495,11 +2738,11 @@ def _dict_to_component_spec( ) def set_progress_bar_config(self, **kwargs): - for sub_block_name, sub_block in self.blocks.sub_blocks.items(): + for sub_block_name, sub_block in self._blocks.sub_blocks.items(): if hasattr(sub_block, "set_progress_bar_config"): sub_block.set_progress_bar_config(**kwargs) - def __call__(self, state: PipelineState = None, output: Union[str, List[str]] = None, **kwargs): + def __call__(self, state: PipelineState = None, output: str | list[str] = None, **kwargs): """ Execute the pipeline by running the pipeline blocks with the given inputs. @@ -2507,11 +2750,11 @@ def __call__(self, state: PipelineState = None, output: Union[str, List[str]] = state (`PipelineState`, optional): PipelineState instance contains inputs and intermediate values. If None, a new `PipelineState` will be created based on the user inputs and the pipeline blocks's requirement. - output (`str` or `List[str]`, optional): + output (`str` or `list[str]`, optional): Optional specification of what to return: - None: Returns the complete `PipelineState` with all inputs and intermediates (default) - str: Returns a specific intermediate value from the state (e.g. `output="image"`) - - List[str]: Returns a dictionary of specific intermediate values (e.g. `output=["image", + - list[str]: Returns a dictionary of specific intermediate values (e.g. `output=["image", "latents"]`) @@ -2536,7 +2779,7 @@ def __call__(self, state: PipelineState = None, output: Union[str, List[str]] = Returns: - If `output` is None: Complete `PipelineState` containing all inputs and intermediates - If `output` is str: The specific intermediate value from the state (e.g. `output="image"`) - - If `output` is List[str]: Dictionary mapping output names to their values from the state (e.g. + - If `output` is list[str]: Dictionary mapping output names to their values from the state (e.g. `output=["image", "latents"]`) """ if state is None: @@ -2549,7 +2792,7 @@ def __call__(self, state: PipelineState = None, output: Union[str, List[str]] = # Add inputs to state, using defaults if not provided in the kwargs or the state # if same input already in the state, will override it if provided in the kwargs - for expected_input_param in self.blocks.inputs: + for expected_input_param in self._blocks.inputs: name = expected_input_param.name default = expected_input_param.default kwargs_type = expected_input_param.kwargs_type @@ -2568,9 +2811,9 @@ def __call__(self, state: PipelineState = None, output: Union[str, List[str]] = # Run the pipeline with torch.no_grad(): try: - _, state = self.blocks(self, state) + _, state = self._blocks(self, state) except Exception: - error_msg = f"Error in block: ({self.blocks.__class__.__name__}):\n" + error_msg = f"Error in block: ({self._blocks.__class__.__name__}):\n" logger.error(error_msg) raise diff --git a/src/diffusers/modular_pipelines/modular_pipeline_utils.py b/src/diffusers/modular_pipelines/modular_pipeline_utils.py index aa421a53727b..fa82f17a9108 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline_utils.py +++ b/src/diffusers/modular_pipelines/modular_pipeline_utils.py @@ -14,15 +14,20 @@ import inspect import re +import warnings from collections import OrderedDict -from dataclasses import dataclass, field, fields -from typing import Any, Dict, List, Literal, Optional, Type, Union +from dataclasses import dataclass, field +from types import UnionType +from typing import Any, Literal, Type, Union, get_args, get_origin +import PIL.Image import torch +from packaging.specifiers import InvalidSpecifier, SpecifierSet from ..configuration_utils import ConfigMixin, FrozenDict from ..loaders.single_file_utils import _is_single_file_path_or_url -from ..utils import is_torch_available, logging +from ..utils import DIFFUSERS_LOAD_ID_FIELDS, is_torch_available, logging +from ..utils.import_utils import _is_package_available if is_torch_available(): @@ -30,6 +35,26 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name +# Template for modular pipeline model card description with placeholders +MODULAR_MODEL_CARD_TEMPLATE = """{model_description} + +## Example Usage + +[TODO] + +## Pipeline Architecture + +This modular pipeline is composed of the following blocks: + +{blocks_description} {trigger_inputs_section} + +## Model Components + +{components_description} {configs_section} + +{io_specification_section} +""" + class InsertableDict(OrderedDict): def insert(self, key, value, index): @@ -88,18 +113,18 @@ class ComponentSpec: default_creation_method: Preferred creation method - "from_config" or "from_pretrained" """ - name: Optional[str] = None - type_hint: Optional[Type] = None - description: Optional[str] = None - config: Optional[FrozenDict] = None - pretrained_model_name_or_path: Optional[Union[str, List[str]]] = field(default=None, metadata={"loading": True}) - subfolder: Optional[str] = field(default="", metadata={"loading": True}) - variant: Optional[str] = field(default=None, metadata={"loading": True}) - revision: Optional[str] = field(default=None, metadata={"loading": True}) + name: str | None = None + type_hint: Type | None = None + description: str | None = None + config: FrozenDict | None = None + pretrained_model_name_or_path: str | list[str] | None = field(default=None, metadata={"loading": True}) + subfolder: str | None = field(default="", metadata={"loading": True}) + variant: str | None = field(default=None, metadata={"loading": True}) + revision: str | None = field(default=None, metadata={"loading": True}) default_creation_method: Literal["from_config", "from_pretrained"] = "from_pretrained" # Deprecated - repo: Optional[Union[str, List[str]]] = field(default=None, metadata={"loading": False}) + repo: str | list[str] | None = field(default=None, metadata={"loading": False}) def __post_init__(self): repo_value = self.repo @@ -181,11 +206,11 @@ def from_component(cls, name: str, component: Any) -> Any: ) @classmethod - def loading_fields(cls) -> List[str]: + def loading_fields(cls) -> list[str]: """ Return the names of all loading‐related fields (i.e. those whose field.metadata["loading"] is True). """ - return [f.name for f in fields(cls) if f.metadata.get("loading", False)] + return DIFFUSERS_LOAD_ID_FIELDS.copy() @property def load_id(self) -> str: @@ -197,10 +222,10 @@ def load_id(self) -> str: return "null" parts = [getattr(self, k) for k in self.loading_fields()] parts = ["null" if p is None else p for p in parts] - return "|".join(p for p in parts if p) + return "|".join(parts) @classmethod - def decode_load_id(cls, load_id: str) -> Dict[str, Optional[str]]: + def decode_load_id(cls, load_id: str) -> dict[str, str | None]: """ Decode a load_id string back into a dictionary of loading fields and values. @@ -238,7 +263,7 @@ def decode_load_id(cls, load_id: str) -> Dict[str, Optional[str]]: # otherwise we cannot do spec -> spec.create() -> component -> ComponentSpec.from_component(component) # the config info is lost in the process # remove error check in from_component spec and ModularPipeline.update_components() if we remove support for non configmixin in `create()` method - def create(self, config: Optional[Union[FrozenDict, Dict[str, Any]]] = None, **kwargs) -> Any: + def create(self, config: FrozenDict | dict[str, Any] | None = None, **kwargs) -> Any: """Create component using from_config with config.""" if self.type_hint is None or not isinstance(self.type_hint, type): @@ -284,6 +309,12 @@ def load(self, **kwargs) -> Any: f"`type_hint` is required when loading a single file model but is missing for component: {self.name}" ) + # `torch_dtype` is not an accepted parameter for tokenizers and processors. + # As a result, it gets stored in `init_kwargs`, which are written to the config + # during save. This causes JSON serialization to fail when saving the component. + if self.type_hint is not None and not issubclass(self.type_hint, torch.nn.Module): + kwargs.pop("torch_dtype", None) + if self.type_hint is None: try: from diffusers import AutoModel @@ -301,6 +332,12 @@ def load(self, **kwargs) -> Any: else getattr(self.type_hint, "from_pretrained") ) + # `torch_dtype` is not an accepted parameter for tokenizers and processors. + # As a result, it gets stored in `init_kwargs`, which are written to the config + # during save. This causes JSON serialization to fail when saving the component. + if not issubclass(self.type_hint, torch.nn.Module): + kwargs.pop("torch_dtype", None) + try: component = load_method(pretrained_model_name_or_path, **load_kwargs, **kwargs) except Exception as e: @@ -320,14 +357,199 @@ class ConfigSpec: name: str default: Any - description: Optional[str] = None + description: str | None = None + + +# ====================================================== +# InputParam and OutputParam templates +# ====================================================== + +INPUT_PARAM_TEMPLATES = { + "prompt": { + "type_hint": str, + "required": True, + "description": "The prompt or prompts to guide image generation.", + }, + "negative_prompt": { + "type_hint": str, + "description": "The prompt or prompts not to guide the image generation.", + }, + "max_sequence_length": { + "type_hint": int, + "default": 512, + "description": "Maximum sequence length for prompt encoding.", + }, + "height": { + "type_hint": int, + "description": "The height in pixels of the generated image.", + }, + "width": { + "type_hint": int, + "description": "The width in pixels of the generated image.", + }, + "num_inference_steps": { + "type_hint": int, + "default": 50, + "description": "The number of denoising steps.", + }, + "num_images_per_prompt": { + "type_hint": int, + "default": 1, + "description": "The number of images to generate per prompt.", + }, + "generator": { + "type_hint": torch.Generator, + "description": "Torch generator for deterministic generation.", + }, + "sigmas": { + "type_hint": list[float], + "description": "Custom sigmas for the denoising process.", + }, + "strength": { + "type_hint": float, + "default": 0.9, + "description": "Strength for img2img/inpainting.", + }, + "image": { + "type_hint": PIL.Image.Image | list[PIL.Image.Image], + "required": True, + "description": "Reference image(s) for denoising. Can be a single image or list of images.", + }, + "latents": { + "type_hint": torch.Tensor, + "description": "Pre-generated noisy latents for image generation.", + }, + "timesteps": { + "type_hint": torch.Tensor, + "description": "Timesteps for the denoising process.", + }, + "output_type": { + "type_hint": str, + "default": "pil", + "description": "Output format: 'pil', 'np', 'pt'.", + }, + "attention_kwargs": { + "type_hint": dict[str, Any], + "description": "Additional kwargs for attention processors.", + }, + "denoiser_input_fields": { + "name": None, + "kwargs_type": "denoiser_input_fields", + "description": "conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.", + }, + # inpainting + "mask_image": { + "type_hint": PIL.Image.Image, + "required": True, + "description": "Mask image for inpainting.", + }, + "padding_mask_crop": { + "type_hint": int, + "description": "Padding for mask cropping in inpainting.", + }, + # controlnet + "control_image": { + "type_hint": PIL.Image.Image, + "required": True, + "description": "Control image for ControlNet conditioning.", + }, + "control_guidance_start": { + "type_hint": float, + "default": 0.0, + "description": "When to start applying ControlNet.", + }, + "control_guidance_end": { + "type_hint": float, + "default": 1.0, + "description": "When to stop applying ControlNet.", + }, + "controlnet_conditioning_scale": { + "type_hint": float, + "default": 1.0, + "description": "Scale for ControlNet conditioning.", + }, + "layers": { + "type_hint": int, + "default": 4, + "description": "Number of layers to extract from the image", + }, + # common intermediate inputs + "prompt_embeds": { + "type_hint": torch.Tensor, + "required": True, + "description": "text embeddings used to guide the image generation. Can be generated from text_encoder step.", + }, + "prompt_embeds_mask": { + "type_hint": torch.Tensor, + "required": True, + "description": "mask for the text embeddings. Can be generated from text_encoder step.", + }, + "negative_prompt_embeds": { + "type_hint": torch.Tensor, + "description": "negative text embeddings used to guide the image generation. Can be generated from text_encoder step.", + }, + "negative_prompt_embeds_mask": { + "type_hint": torch.Tensor, + "description": "mask for the negative text embeddings. Can be generated from text_encoder step.", + }, + "image_latents": { + "type_hint": torch.Tensor, + "required": True, + "description": "image latents used to guide the image generation. Can be generated from vae_encoder step.", + }, + "batch_size": { + "type_hint": int, + "default": 1, + "description": "Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can be generated in input step.", + }, + "dtype": { + "type_hint": torch.dtype, + "default": torch.float32, + "description": "The dtype of the model inputs, can be generated in input step.", + }, +} + +OUTPUT_PARAM_TEMPLATES = { + "images": { + "type_hint": list[PIL.Image.Image], + "description": "Generated images.", + }, + "videos": { + "type_hint": list[PIL.Image.Image], + "description": "The generated videos.", + }, + "latents": { + "type_hint": torch.Tensor, + "description": "Denoised latents.", + }, + # intermediate outputs + "prompt_embeds": { + "type_hint": torch.Tensor, + "kwargs_type": "denoiser_input_fields", + "description": "The prompt embeddings.", + }, + "prompt_embeds_mask": { + "type_hint": torch.Tensor, + "kwargs_type": "denoiser_input_fields", + "description": "The encoder attention mask.", + }, + "negative_prompt_embeds": { + "type_hint": torch.Tensor, + "kwargs_type": "denoiser_input_fields", + "description": "The negative prompt embeddings.", + }, + "negative_prompt_embeds_mask": { + "type_hint": torch.Tensor, + "kwargs_type": "denoiser_input_fields", + "description": "The negative prompt embeddings mask.", + }, + "image_latents": { + "type_hint": torch.Tensor, + "description": "The latent representation of the input image.", + }, +} -# YiYi Notes: both inputs and intermediate_inputs are InputParam objects -# however some fields are not relevant for intermediate_inputs -# e.g. unlike inputs, required only used in docstring for intermediate_inputs, we do not check if a required intermediate inputs is passed -# default is not used for intermediate_inputs, we only use default from inputs, so it is ignored if it is set for intermediate_inputs -# -> should we use different class for inputs and intermediate_inputs? @dataclass class InputParam: """Specification for an input parameter.""" @@ -337,11 +559,32 @@ class InputParam: default: Any = None required: bool = False description: str = "" - kwargs_type: str = None # YiYi Notes: remove this feature (maybe) + kwargs_type: str = None + metadata: dict[str, Any] = None def __repr__(self): return f"<{self.name}: {'required' if self.required else 'optional'}, default={self.default}>" + @classmethod + def template(cls, template_name: str, note: str = None, **overrides) -> "InputParam": + """Get template for name if exists, otherwise raise ValueError.""" + if template_name not in INPUT_PARAM_TEMPLATES: + raise ValueError(f"InputParam template for {template_name} not found") + + template_kwargs = INPUT_PARAM_TEMPLATES[template_name].copy() + + # Determine the actual param name: + # 1. From overrides if provided + # 2. From template if present + # 3. Fall back to template_name + name = overrides.pop("name", template_kwargs.pop("name", template_name)) + + if note and "description" in template_kwargs: + template_kwargs["description"] = f"{template_kwargs['description']} ({note})" + + template_kwargs.update(overrides) + return cls(name=name, **template_kwargs) + @dataclass class OutputParam: @@ -350,20 +593,41 @@ class OutputParam: name: str type_hint: Any = None description: str = "" - kwargs_type: str = None # YiYi notes: remove this feature (maybe) + kwargs_type: str = None + metadata: dict[str, Any] = None def __repr__(self): return ( f"<{self.name}: {self.type_hint.__name__ if hasattr(self.type_hint, '__name__') else str(self.type_hint)}>" ) + @classmethod + def template(cls, template_name: str, note: str = None, **overrides) -> "OutputParam": + """Get template for name if exists, otherwise raise ValueError.""" + if template_name not in OUTPUT_PARAM_TEMPLATES: + raise ValueError(f"OutputParam template for {template_name} not found") + + template_kwargs = OUTPUT_PARAM_TEMPLATES[template_name].copy() + + # Determine the actual param name: + # 1. From overrides if provided + # 2. From template if present + # 3. Fall back to template_name + name = overrides.pop("name", template_kwargs.pop("name", template_name)) + + if note and "description" in template_kwargs: + template_kwargs["description"] = f"{template_kwargs['description']} ({note})" + + template_kwargs.update(overrides) + return cls(name=name, **template_kwargs) + def format_inputs_short(inputs): """ Format input parameters into a string representation, with required params first followed by optional ones. Args: - inputs: List of input parameters with 'required' and 'name' attributes, and 'default' for optional params + inputs: list of input parameters with 'required' and 'name' attributes, and 'default' for optional params Returns: str: Formatted string of input parameters @@ -392,9 +656,9 @@ def format_intermediates_short(intermediate_inputs, required_intermediate_inputs Formats intermediate inputs and outputs of a block into a string representation. Args: - intermediate_inputs: List of intermediate input parameters - required_intermediate_inputs: List of required intermediate input names - intermediate_outputs: List of intermediate output parameters + intermediate_inputs: list of intermediate input parameters + required_intermediate_inputs: list of required intermediate input names + intermediate_outputs: list of intermediate output parameters Returns: str: Formatted string like: @@ -441,7 +705,7 @@ def format_params(params, header="Args", indent_level=4, max_line_length=115): """Format a list of InputParam or OutputParam objects into a readable string representation. Args: - params: List of InputParam or OutputParam objects to format + params: list of InputParam or OutputParam objects to format header: Header text to use (e.g. "Args" or "Returns") indent_level: Number of spaces to indent each parameter line (default: 4) max_line_length: Maximum length for each line before wrapping (default: 115) @@ -458,9 +722,9 @@ def format_params(params, header="Args", indent_level=4, max_line_length=115): formatted_params = [] def get_type_str(type_hint): - if hasattr(type_hint, "__origin__") and type_hint.__origin__ is Union: - types = [t.__name__ if hasattr(t, "__name__") else str(t) for t in type_hint.__args__] - return f"Union[{', '.join(types)}]" + if isinstance(type_hint, UnionType) or get_origin(type_hint) is Union: + type_strs = [t.__name__ if hasattr(t, "__name__") else str(t) for t in get_args(type_hint)] + return " | ".join(type_strs) return type_hint.__name__ if hasattr(type_hint, "__name__") else str(type_hint) def wrap_text(text, indent, max_length): @@ -509,17 +773,19 @@ def wrap_text(text, indent, max_length): desc = re.sub(r"\[(.*?)\]\((https?://[^\s\)]+)\)", r"[\1](\2)", param.description) wrapped_desc = wrap_text(desc, desc_indent, max_line_length) param_str += f"\n{desc_indent}{wrapped_desc}" + else: + param_str += f"\n{desc_indent}TODO: Add description." formatted_params.append(param_str) - return "\n\n".join(formatted_params) + return "\n".join(formatted_params) def format_input_params(input_params, indent_level=4, max_line_length=115): """Format a list of InputParam objects into a readable string representation. Args: - input_params: List of InputParam objects to format + input_params: list of InputParam objects to format indent_level: Number of spaces to indent each parameter line (default: 4) max_line_length: Maximum length for each line before wrapping (default: 115) @@ -533,7 +799,7 @@ def format_output_params(output_params, indent_level=4, max_line_length=115): """Format a list of OutputParam objects into a readable string representation. Args: - output_params: List of OutputParam objects to format + output_params: list of OutputParam objects to format indent_level: Number of spaces to indent each parameter line (default: 4) max_line_length: Maximum length for each line before wrapping (default: 115) @@ -543,11 +809,51 @@ def format_output_params(output_params, indent_level=4, max_line_length=115): return format_params(output_params, "Outputs", indent_level, max_line_length) +def format_params_markdown(params, header="Inputs"): + """Format a list of InputParam or OutputParam objects as a markdown bullet-point list. + + Suitable for model cards rendered on Hugging Face Hub. + + Args: + params: list of InputParam or OutputParam objects to format + header: Header text (e.g. "Inputs" or "Outputs") + + Returns: + A formatted markdown string, or empty string if params is empty. + """ + if not params: + return "" + + def get_type_str(type_hint): + if isinstance(type_hint, UnionType) or get_origin(type_hint) is Union: + type_strs = [t.__name__ if hasattr(t, "__name__") else str(t) for t in get_args(type_hint)] + return " | ".join(type_strs) + return type_hint.__name__ if hasattr(type_hint, "__name__") else str(type_hint) + + lines = [f"**{header}:**\n"] if header else [] + for param in params: + type_str = get_type_str(param.type_hint) if param.type_hint != Any else "" + name = f"**{param.kwargs_type}" if param.name is None and param.kwargs_type is not None else param.name + param_str = f"- `{name}` (`{type_str}`" + + if hasattr(param, "required") and not param.required: + param_str += ", *optional*" + if param.default is not None: + param_str += f", defaults to `{param.default}`" + param_str += ")" + + desc = param.description if param.description else "No description provided" + param_str += f": {desc}" + lines.append(param_str) + + return "\n".join(lines) + + def format_components(components, indent_level=4, max_line_length=115, add_empty_lines=True): """Format a list of ComponentSpec objects into a readable string representation. Args: - components: List of ComponentSpec objects to format + components: list of ComponentSpec objects to format indent_level: Number of spaces to indent each component line (default: 4) max_line_length: Maximum length for each line before wrapping (default: 115) add_empty_lines: Whether to add empty lines between components (default: True) @@ -582,7 +888,7 @@ def format_components(components, indent_level=4, max_line_length=115, add_empty loading_field_values = [] for field_name in component.loading_fields(): field_value = getattr(component, field_name) - if field_value is not None: + if field_value: loading_field_values.append(f"{field_name}={field_value}") # Add loading field information if available @@ -602,7 +908,7 @@ def format_configs(configs, indent_level=4, max_line_length=115, add_empty_lines """Format a list of ConfigSpec objects into a readable string representation. Args: - configs: List of ConfigSpec objects to format + configs: list of ConfigSpec objects to format indent_level: Number of spaces to indent each config line (default: 4) max_line_length: Maximum length for each line before wrapping (default: 115) add_empty_lines: Whether to add empty lines between configs (default: True) @@ -636,6 +942,30 @@ def format_configs(configs, indent_level=4, max_line_length=115, add_empty_lines return "\n".join(formatted_configs) +def format_workflow(workflow_map): + """Format a workflow map into a readable string representation. + + Args: + workflow_map: Dictionary mapping workflow names to trigger inputs + + Returns: + A formatted string representing all workflows + """ + if workflow_map is None: + return "" + + lines = ["Supported workflows:"] + for workflow_name, trigger_inputs in workflow_map.items(): + required_inputs = [k for k, v in trigger_inputs.items() if v] + if required_inputs: + inputs_str = ", ".join(f"`{t}`" for t in required_inputs) + lines.append(f" - `{workflow_name}`: requires {inputs_str}") + else: + lines.append(f" - `{workflow_name}`: default (no additional inputs required)") + + return "\n".join(lines) + + def make_doc_string( inputs, outputs, @@ -648,13 +978,13 @@ def make_doc_string( Generates a formatted documentation string describing the pipeline block's parameters and structure. Args: - inputs: List of input parameters - intermediate_inputs: List of intermediate input parameters - outputs: List of output parameters + inputs: list of input parameters + intermediate_inputs: list of intermediate input parameters + outputs: list of output parameters description (str, *optional*): Description of the block class_name (str, *optional*): Name of the class to include in the documentation - expected_components (List[ComponentSpec], *optional*): List of expected components - expected_configs (List[ConfigSpec], *optional*): List of expected configurations + expected_components (list[ComponentSpec], *optional*): list of expected components + expected_configs (list[ConfigSpec], *optional*): list of expected configurations Returns: str: A formatted string containing information about components, configs, call parameters, @@ -669,17 +999,17 @@ def make_doc_string( # Add description if description: desc_lines = description.strip().split("\n") - aligned_desc = "\n".join(" " + line for line in desc_lines) + aligned_desc = "\n".join(" " + line.rstrip() for line in desc_lines) output += aligned_desc + "\n\n" # Add components section if provided if expected_components and len(expected_components) > 0: - components_str = format_components(expected_components, indent_level=2) + components_str = format_components(expected_components, indent_level=2, add_empty_lines=False) output += components_str + "\n\n" # Add configs section if provided if expected_configs and len(expected_configs) > 0: - configs_str = format_configs(expected_configs, indent_level=2) + configs_str = format_configs(expected_configs, indent_level=2, add_empty_lines=False) output += configs_str + "\n\n" # Add inputs section @@ -690,3 +1020,340 @@ def make_doc_string( output += format_output_params(outputs, indent_level=2) return output + + +def _validate_requirements(reqs): + if reqs is None: + normalized_reqs = {} + else: + if not isinstance(reqs, dict): + raise ValueError( + "Requirements must be provided as a dictionary mapping package names to version specifiers." + ) + normalized_reqs = _normalize_requirements(reqs) + + if not normalized_reqs: + return {} + + final: dict[str, str] = {} + for req, specified_ver in normalized_reqs.items(): + req_available, req_actual_ver = _is_package_available(req) + if not req_available: + logger.warning(f"{req} was specified in the requirements but wasn't found in the current environment.") + + if specified_ver: + try: + specifier = SpecifierSet(specified_ver) + except InvalidSpecifier as err: + raise ValueError(f"Requirement specifier '{specified_ver}' for {req} is invalid.") from err + + if req_actual_ver == "N/A": + logger.warning( + f"Version of {req} could not be determined to validate requirement '{specified_ver}'. Things might work unexpected." + ) + elif not specifier.contains(req_actual_ver, prereleases=True): + logger.warning( + f"{req} requirement '{specified_ver}' is not satisfied by the installed version {req_actual_ver}. Things might work unexpected." + ) + + final[req] = specified_ver + + return final + + +def _normalize_requirements(reqs): + if not reqs: + return {} + + normalized: "OrderedDict[str, str]" = OrderedDict() + + def _accumulate(mapping: dict[str, Any]): + for pkg, spec in mapping.items(): + if isinstance(spec, dict): + # This is recursive because blocks are composable. This way, we can merge requirements + # from multiple blocks. + _accumulate(spec) + continue + + pkg_name = str(pkg).strip() + if not pkg_name: + raise ValueError("Requirement package name cannot be empty.") + + spec_str = "" if spec is None else str(spec).strip() + if spec_str and not spec_str.startswith(("<", ">", "=", "!", "~")): + spec_str = f"=={spec_str}" + + existing_spec = normalized.get(pkg_name) + if existing_spec is not None: + if not existing_spec and spec_str: + normalized[pkg_name] = spec_str + elif existing_spec and spec_str and existing_spec != spec_str: + try: + combined_spec = SpecifierSet(",".join(filter(None, [existing_spec, spec_str]))) + except InvalidSpecifier: + logger.warning( + f"Conflicting requirements for '{pkg_name}' detected: '{existing_spec}' vs '{spec_str}'. Keeping '{existing_spec}'." + ) + else: + normalized[pkg_name] = str(combined_spec) + continue + + normalized[pkg_name] = spec_str + + _accumulate(reqs) + + return normalized + + +def combine_inputs(*named_input_lists: list[tuple[str, list[InputParam]]]) -> list[InputParam]: + """ + Combines multiple lists of InputParam objects from different blocks. For duplicate inputs, updates only if current + default value is None and new default value is not None. Warns if multiple non-None default values exist for the + same input. + + Args: + named_input_lists: List of tuples containing (block_name, input_param_list) pairs + + Returns: + List[InputParam]: Combined list of unique InputParam objects + """ + combined_dict = {} # name -> InputParam + value_sources = {} # name -> block_name + + for block_name, inputs in named_input_lists: + for input_param in inputs: + if input_param.name is None and input_param.kwargs_type is not None: + input_name = "*_" + input_param.kwargs_type + else: + input_name = input_param.name + if input_name in combined_dict: + current_param = combined_dict[input_name] + if ( + current_param.default is not None + and input_param.default is not None + and current_param.default != input_param.default + ): + warnings.warn( + f"Multiple different default values found for input '{input_name}': " + f"{current_param.default} (from block '{value_sources[input_name]}') and " + f"{input_param.default} (from block '{block_name}'). Using {current_param.default}." + ) + if current_param.default is None and input_param.default is not None: + combined_dict[input_name] = input_param + value_sources[input_name] = block_name + else: + combined_dict[input_name] = input_param + value_sources[input_name] = block_name + + return list(combined_dict.values()) + + +def combine_outputs(*named_output_lists: list[tuple[str, list[OutputParam]]]) -> list[OutputParam]: + """ + Combines multiple lists of OutputParam objects from different blocks. For duplicate outputs, keeps the first + occurrence of each output name. + + Args: + named_output_lists: List of tuples containing (block_name, output_param_list) pairs + + Returns: + List[OutputParam]: Combined list of unique OutputParam objects + """ + combined_dict = {} # name -> OutputParam + + for block_name, outputs in named_output_lists: + for output_param in outputs: + if (output_param.name not in combined_dict) or ( + combined_dict[output_param.name].kwargs_type is None and output_param.kwargs_type is not None + ): + combined_dict[output_param.name] = output_param + + return list(combined_dict.values()) + + +def generate_modular_model_card_content(blocks) -> dict[str, Any]: + """ + Generate model card content for a modular pipeline. + + This function creates a comprehensive model card with descriptions of the pipeline's architecture, components, + configurations, inputs, and outputs. + + Args: + blocks: The pipeline's blocks object containing all pipeline specifications + + Returns: + Dict[str, Any]: A dictionary containing formatted content sections: + - pipeline_name: Name of the pipeline + - model_description: Overall description with pipeline type + - blocks_description: Detailed architecture of blocks + - components_description: List of required components + - configs_section: Configuration parameters section + - io_specification_section: Input/Output specification (per-workflow or unified) + - trigger_inputs_section: Conditional execution information + - tags: List of relevant tags for the model card + """ + blocks_class_name = blocks.__class__.__name__ + pipeline_name = blocks_class_name.replace("Blocks", " Pipeline") + description = getattr(blocks, "description", "A modular diffusion pipeline.") + + # generate blocks architecture description + blocks_desc_parts = [] + sub_blocks = getattr(blocks, "sub_blocks", None) or {} + if sub_blocks: + for i, (name, block) in enumerate(sub_blocks.items()): + block_class = block.__class__.__name__ + block_desc = block.description.split("\n")[0] if getattr(block, "description", "") else "" + blocks_desc_parts.append(f"{i + 1}. **{name}** (`{block_class}`)") + if block_desc: + blocks_desc_parts.append(f" - {block_desc}") + + blocks_description = "\n".join(blocks_desc_parts) if blocks_desc_parts else "No blocks defined." + + components = getattr(blocks, "expected_components", []) + if components: + components_str = format_components(components, indent_level=0, add_empty_lines=False) + # remove the "Components:" header since template has its own + components_description = components_str.replace("Components:\n", "").strip() + if components_description: + # Convert to enumerated list + lines = [line.strip() for line in components_description.split("\n") if line.strip()] + enumerated_lines = [f"{i + 1}. {line}" for i, line in enumerate(lines)] + components_description = "\n".join(enumerated_lines) + else: + components_description = "No specific components required." + else: + components_description = "No specific components required. Components can be loaded dynamically." + + configs = getattr(blocks, "expected_configs", []) + configs_section = "" + if configs: + configs_str = format_configs(configs, indent_level=0, add_empty_lines=False) + configs_description = configs_str.replace("Configs:\n", "").strip() + if configs_description: + configs_section = f"\n\n## Configuration Parameters\n\n{configs_description}" + + # Branch on whether workflows are defined + has_workflows = getattr(blocks, "_workflow_map", None) is not None + + if has_workflows: + workflow_map = blocks._workflow_map + parts = [] + + # If blocks overrides outputs (e.g. to return just "images" instead of all intermediates), + # use that as the shared output for all workflows + blocks_outputs = blocks.outputs + blocks_intermediate = getattr(blocks, "intermediate_outputs", None) + shared_outputs = ( + blocks_outputs if blocks_intermediate is not None and blocks_outputs != blocks_intermediate else None + ) + + parts.append("## Workflow Input Specification\n") + + # Per-workflow details: show trigger inputs with full param descriptions + for wf_name, trigger_inputs in workflow_map.items(): + trigger_input_names = set(trigger_inputs.keys()) + try: + workflow_blocks = blocks.get_workflow(wf_name) + except Exception: + parts.append(f"
\n{wf_name}\n") + parts.append("*Could not resolve workflow blocks.*\n") + parts.append("
\n") + continue + + wf_inputs = workflow_blocks.inputs + # Show only trigger inputs with full parameter descriptions + trigger_params = [p for p in wf_inputs if p.name in trigger_input_names] + + parts.append(f"
\n{wf_name}\n") + + inputs_str = format_params_markdown(trigger_params, header=None) + parts.append(inputs_str if inputs_str else "No additional inputs required.") + parts.append("") + + parts.append("
\n") + + # Common Inputs & Outputs section (like non-workflow pipelines) + all_inputs = blocks.inputs + all_outputs = shared_outputs if shared_outputs is not None else blocks.outputs + + inputs_str = format_params_markdown(all_inputs, "Inputs") + outputs_str = format_params_markdown(all_outputs, "Outputs") + inputs_description = inputs_str if inputs_str else "No specific inputs defined." + outputs_description = outputs_str if outputs_str else "Standard pipeline outputs." + + parts.append(f"\n## Input/Output Specification\n\n{inputs_description}\n\n{outputs_description}") + + io_specification_section = "\n".join(parts) + # Suppress trigger_inputs_section when workflows are shown (it's redundant) + trigger_inputs_section = "" + else: + # Unified I/O section (original behavior) + inputs = blocks.inputs + outputs = blocks.outputs + inputs_str = format_params_markdown(inputs, "Inputs") + outputs_str = format_params_markdown(outputs, "Outputs") + inputs_description = inputs_str if inputs_str else "No specific inputs defined." + outputs_description = outputs_str if outputs_str else "Standard pipeline outputs." + io_specification_section = f"## Input/Output Specification\n\n{inputs_description}\n\n{outputs_description}" + + trigger_inputs_section = "" + if hasattr(blocks, "trigger_inputs") and blocks.trigger_inputs: + trigger_inputs_list = sorted([t for t in blocks.trigger_inputs if t is not None]) + if trigger_inputs_list: + trigger_inputs_str = ", ".join(f"`{t}`" for t in trigger_inputs_list) + trigger_inputs_section = f""" +### Conditional Execution + +This pipeline contains blocks that are selected at runtime based on inputs: +- **Trigger Inputs**: {trigger_inputs_str} +""" + + # generate tags based on pipeline characteristics + tags = ["modular-diffusers", "diffusers"] + + if hasattr(blocks, "model_name") and blocks.model_name: + tags.append(blocks.model_name) + + if has_workflows: + # Derive tags from workflow names + workflow_names = set(blocks._workflow_map.keys()) + if any("inpainting" in wf for wf in workflow_names): + tags.append("inpainting") + if any("image2image" in wf for wf in workflow_names): + tags.append("image-to-image") + if any("controlnet" in wf for wf in workflow_names): + tags.append("controlnet") + if any("text2image" in wf for wf in workflow_names): + tags.append("text-to-image") + elif hasattr(blocks, "trigger_inputs") and blocks.trigger_inputs: + triggers = blocks.trigger_inputs + if any(t in triggers for t in ["mask", "mask_image"]): + tags.append("inpainting") + if any(t in triggers for t in ["image", "image_latents"]): + tags.append("image-to-image") + if any(t in triggers for t in ["control_image", "controlnet_cond"]): + tags.append("controlnet") + if not any(t in triggers for t in ["image", "mask", "image_latents", "mask_image"]): + tags.append("text-to-image") + else: + tags.append("text-to-image") + + block_count = len(blocks.sub_blocks) + model_description = f"""This is a modular diffusion pipeline built with 🧨 Diffusers' modular pipeline framework. + +**Pipeline Type**: {blocks_class_name} + +**Description**: {description} + +This pipeline uses a {block_count}-block architecture that can be customized and extended.""" + + return { + "pipeline_name": pipeline_name, + "model_description": model_description, + "blocks_description": blocks_description, + "components_description": components_description, + "configs_section": configs_section, + "io_specification_section": io_specification_section, + "trigger_inputs_section": trigger_inputs_section, + "tags": tags, + } diff --git a/src/diffusers/modular_pipelines/node_utils.py b/src/diffusers/modular_pipelines/node_utils.py deleted file mode 100644 index f7ee1dd3097b..000000000000 --- a/src/diffusers/modular_pipelines/node_utils.py +++ /dev/null @@ -1,661 +0,0 @@ -import json -import logging -import os -from pathlib import Path -from typing import List, Optional, Tuple, Union - -import numpy as np -import PIL -import torch - -from ..configuration_utils import ConfigMixin -from ..image_processor import PipelineImageInput -from .modular_pipeline import ModularPipelineBlocks, SequentialPipelineBlocks -from .modular_pipeline_utils import InputParam - - -logger = logging.getLogger(__name__) - -# YiYi Notes: this is actually for SDXL, put it here for now -SDXL_INPUTS_SCHEMA = { - "prompt": InputParam( - "prompt", type_hint=Union[str, List[str]], description="The prompt or prompts to guide the image generation" - ), - "prompt_2": InputParam( - "prompt_2", - type_hint=Union[str, List[str]], - description="The prompt or prompts to be sent to the tokenizer_2 and text_encoder_2", - ), - "negative_prompt": InputParam( - "negative_prompt", - type_hint=Union[str, List[str]], - description="The prompt or prompts not to guide the image generation", - ), - "negative_prompt_2": InputParam( - "negative_prompt_2", - type_hint=Union[str, List[str]], - description="The negative prompt or prompts for text_encoder_2", - ), - "cross_attention_kwargs": InputParam( - "cross_attention_kwargs", - type_hint=Optional[dict], - description="Kwargs dictionary passed to the AttentionProcessor", - ), - "clip_skip": InputParam( - "clip_skip", type_hint=Optional[int], description="Number of layers to skip in CLIP text encoder" - ), - "image": InputParam( - "image", - type_hint=PipelineImageInput, - required=True, - description="The image(s) to modify for img2img or inpainting", - ), - "mask_image": InputParam( - "mask_image", - type_hint=PipelineImageInput, - required=True, - description="Mask image for inpainting, white pixels will be repainted", - ), - "generator": InputParam( - "generator", - type_hint=Optional[Union[torch.Generator, List[torch.Generator]]], - description="Generator(s) for deterministic generation", - ), - "height": InputParam("height", type_hint=Optional[int], description="Height in pixels of the generated image"), - "width": InputParam("width", type_hint=Optional[int], description="Width in pixels of the generated image"), - "num_images_per_prompt": InputParam( - "num_images_per_prompt", type_hint=int, default=1, description="Number of images to generate per prompt" - ), - "num_inference_steps": InputParam( - "num_inference_steps", type_hint=int, default=50, description="Number of denoising steps" - ), - "timesteps": InputParam( - "timesteps", type_hint=Optional[torch.Tensor], description="Custom timesteps for the denoising process" - ), - "sigmas": InputParam( - "sigmas", type_hint=Optional[torch.Tensor], description="Custom sigmas for the denoising process" - ), - "denoising_end": InputParam( - "denoising_end", - type_hint=Optional[float], - description="Fraction of denoising process to complete before termination", - ), - # YiYi Notes: img2img defaults to 0.3, inpainting defaults to 0.9999 - "strength": InputParam( - "strength", type_hint=float, default=0.3, description="How much to transform the reference image" - ), - "denoising_start": InputParam( - "denoising_start", type_hint=Optional[float], description="Starting point of the denoising process" - ), - "latents": InputParam( - "latents", type_hint=Optional[torch.Tensor], description="Pre-generated noisy latents for image generation" - ), - "padding_mask_crop": InputParam( - "padding_mask_crop", - type_hint=Optional[Tuple[int, int]], - description="Size of margin in crop for image and mask", - ), - "original_size": InputParam( - "original_size", - type_hint=Optional[Tuple[int, int]], - description="Original size of the image for SDXL's micro-conditioning", - ), - "target_size": InputParam( - "target_size", type_hint=Optional[Tuple[int, int]], description="Target size for SDXL's micro-conditioning" - ), - "negative_original_size": InputParam( - "negative_original_size", - type_hint=Optional[Tuple[int, int]], - description="Negative conditioning based on image resolution", - ), - "negative_target_size": InputParam( - "negative_target_size", - type_hint=Optional[Tuple[int, int]], - description="Negative conditioning based on target resolution", - ), - "crops_coords_top_left": InputParam( - "crops_coords_top_left", - type_hint=Tuple[int, int], - default=(0, 0), - description="Top-left coordinates for SDXL's micro-conditioning", - ), - "negative_crops_coords_top_left": InputParam( - "negative_crops_coords_top_left", - type_hint=Tuple[int, int], - default=(0, 0), - description="Negative conditioning crop coordinates", - ), - "aesthetic_score": InputParam( - "aesthetic_score", type_hint=float, default=6.0, description="Simulates aesthetic score of generated image" - ), - "negative_aesthetic_score": InputParam( - "negative_aesthetic_score", type_hint=float, default=2.0, description="Simulates negative aesthetic score" - ), - "eta": InputParam("eta", type_hint=float, default=0.0, description="Parameter η in the DDIM paper"), - "output_type": InputParam( - "output_type", type_hint=str, default="pil", description="Output format (pil/tensor/np.array)" - ), - "ip_adapter_image": InputParam( - "ip_adapter_image", - type_hint=PipelineImageInput, - required=True, - description="Image(s) to be used as IP adapter", - ), - "control_image": InputParam( - "control_image", type_hint=PipelineImageInput, required=True, description="ControlNet input condition" - ), - "control_guidance_start": InputParam( - "control_guidance_start", - type_hint=Union[float, List[float]], - default=0.0, - description="When ControlNet starts applying", - ), - "control_guidance_end": InputParam( - "control_guidance_end", - type_hint=Union[float, List[float]], - default=1.0, - description="When ControlNet stops applying", - ), - "controlnet_conditioning_scale": InputParam( - "controlnet_conditioning_scale", - type_hint=Union[float, List[float]], - default=1.0, - description="Scale factor for ControlNet outputs", - ), - "guess_mode": InputParam( - "guess_mode", - type_hint=bool, - default=False, - description="Enables ControlNet encoder to recognize input without prompts", - ), - "control_mode": InputParam( - "control_mode", type_hint=List[int], required=True, description="Control mode for union controlnet" - ), -} - -SDXL_INTERMEDIATE_INPUTS_SCHEMA = { - "prompt_embeds": InputParam( - "prompt_embeds", - type_hint=torch.Tensor, - required=True, - description="Text embeddings used to guide image generation", - ), - "negative_prompt_embeds": InputParam( - "negative_prompt_embeds", type_hint=torch.Tensor, description="Negative text embeddings" - ), - "pooled_prompt_embeds": InputParam( - "pooled_prompt_embeds", type_hint=torch.Tensor, required=True, description="Pooled text embeddings" - ), - "negative_pooled_prompt_embeds": InputParam( - "negative_pooled_prompt_embeds", type_hint=torch.Tensor, description="Negative pooled text embeddings" - ), - "batch_size": InputParam("batch_size", type_hint=int, required=True, description="Number of prompts"), - "dtype": InputParam("dtype", type_hint=torch.dtype, description="Data type of model tensor inputs"), - "preprocess_kwargs": InputParam( - "preprocess_kwargs", type_hint=Optional[dict], description="Kwargs for ImageProcessor" - ), - "latents": InputParam( - "latents", type_hint=torch.Tensor, required=True, description="Initial latents for denoising process" - ), - "timesteps": InputParam("timesteps", type_hint=torch.Tensor, required=True, description="Timesteps for inference"), - "num_inference_steps": InputParam( - "num_inference_steps", type_hint=int, required=True, description="Number of denoising steps" - ), - "latent_timestep": InputParam( - "latent_timestep", type_hint=torch.Tensor, required=True, description="Initial noise level timestep" - ), - "image_latents": InputParam( - "image_latents", type_hint=torch.Tensor, required=True, description="Latents representing reference image" - ), - "mask": InputParam("mask", type_hint=torch.Tensor, required=True, description="Mask for inpainting"), - "masked_image_latents": InputParam( - "masked_image_latents", type_hint=torch.Tensor, description="Masked image latents for inpainting" - ), - "add_time_ids": InputParam( - "add_time_ids", type_hint=torch.Tensor, required=True, description="Time ids for conditioning" - ), - "negative_add_time_ids": InputParam( - "negative_add_time_ids", type_hint=torch.Tensor, description="Negative time ids" - ), - "timestep_cond": InputParam("timestep_cond", type_hint=torch.Tensor, description="Timestep conditioning for LCM"), - "noise": InputParam("noise", type_hint=torch.Tensor, description="Noise added to image latents"), - "crops_coords": InputParam("crops_coords", type_hint=Optional[Tuple[int]], description="Crop coordinates"), - "ip_adapter_embeds": InputParam( - "ip_adapter_embeds", type_hint=List[torch.Tensor], description="Image embeddings for IP-Adapter" - ), - "negative_ip_adapter_embeds": InputParam( - "negative_ip_adapter_embeds", - type_hint=List[torch.Tensor], - description="Negative image embeddings for IP-Adapter", - ), - "images": InputParam( - "images", - type_hint=Union[List[PIL.Image.Image], List[torch.Tensor], List[np.array]], - required=True, - description="Generated images", - ), -} - -SDXL_PARAM_SCHEMA = {**SDXL_INPUTS_SCHEMA, **SDXL_INTERMEDIATE_INPUTS_SCHEMA} - - -DEFAULT_PARAM_MAPS = { - "prompt": { - "label": "Prompt", - "type": "string", - "default": "a bear sitting in a chair drinking a milkshake", - "display": "textarea", - }, - "negative_prompt": { - "label": "Negative Prompt", - "type": "string", - "default": "deformed, ugly, wrong proportion, low res, bad anatomy, worst quality, low quality", - "display": "textarea", - }, - "num_inference_steps": { - "label": "Steps", - "type": "int", - "default": 25, - "min": 1, - "max": 1000, - }, - "seed": { - "label": "Seed", - "type": "int", - "default": 0, - "min": 0, - "display": "random", - }, - "width": { - "label": "Width", - "type": "int", - "display": "text", - "default": 1024, - "min": 8, - "max": 8192, - "step": 8, - "group": "dimensions", - }, - "height": { - "label": "Height", - "type": "int", - "display": "text", - "default": 1024, - "min": 8, - "max": 8192, - "step": 8, - "group": "dimensions", - }, - "images": { - "label": "Images", - "type": "image", - "display": "output", - }, - "image": { - "label": "Image", - "type": "image", - "display": "input", - }, -} - -DEFAULT_TYPE_MAPS = { - "int": { - "type": "int", - "default": 0, - "min": 0, - }, - "float": { - "type": "float", - "default": 0.0, - "min": 0.0, - }, - "str": { - "type": "string", - "default": "", - }, - "bool": { - "type": "boolean", - "default": False, - }, - "image": { - "type": "image", - }, -} - -DEFAULT_MODEL_KEYS = ["unet", "vae", "text_encoder", "tokenizer", "controlnet", "transformer", "image_encoder"] -DEFAULT_CATEGORY = "Modular Diffusers" -DEFAULT_EXCLUDE_MODEL_KEYS = ["processor", "feature_extractor", "safety_checker"] -DEFAULT_PARAMS_GROUPS_KEYS = { - "text_encoders": ["text_encoder", "tokenizer"], - "ip_adapter_embeds": ["ip_adapter_embeds"], - "prompt_embeddings": ["prompt_embeds"], -} - - -def get_group_name(name, group_params_keys=DEFAULT_PARAMS_GROUPS_KEYS): - """ - Get the group name for a given parameter name, if not part of a group, return None e.g. "prompt_embeds" -> - "text_embeds", "text_encoder" -> "text_encoders", "prompt" -> None - """ - if name is None: - return None - for group_name, group_keys in group_params_keys.items(): - for group_key in group_keys: - if group_key in name: - return group_name - return None - - -class ModularNode(ConfigMixin): - """ - A ModularNode is a base class to build UI nodes using diffusers. Currently only supports Mellon. It is a wrapper - around a ModularPipelineBlocks object. - - > [!WARNING] > This is an experimental feature and is likely to change in the future. - """ - - config_name = "node_config.json" - - @classmethod - def from_pretrained( - cls, - pretrained_model_name_or_path: str, - trust_remote_code: Optional[bool] = None, - **kwargs, - ): - blocks = ModularPipelineBlocks.from_pretrained( - pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs - ) - return cls(blocks, **kwargs) - - def __init__(self, blocks, category=DEFAULT_CATEGORY, label=None, **kwargs): - self.blocks = blocks - - if label is None: - label = self.blocks.__class__.__name__ - # blocks param name -> mellon param name - self.name_mapping = {} - - input_params = {} - # pass or create a default param dict for each input - # e.g. for prompt, - # prompt = { - # "name": "text_input", # the name of the input in node definition, could be different from the input name in diffusers - # "label": "Prompt", - # "type": "string", - # "default": "a bear sitting in a chair drinking a milkshake", - # "display": "textarea"} - # if type is not specified, it'll be a "custom" param of its own type - # e.g. you can pass ModularNode(scheduler = {name :"scheduler"}) - # it will get this spec in node definition {"scheduler": {"label": "Scheduler", "type": "scheduler", "display": "input"}} - # name can be a dict, in that case, it is part of a "dict" input in mellon nodes, e.g. text_encoder= {name: {"text_encoders": "text_encoder"}} - inputs = self.blocks.inputs + self.blocks.intermediate_inputs - for inp in inputs: - param = kwargs.pop(inp.name, None) - if param: - # user can pass a param dict for all inputs, e.g. ModularNode(prompt = {...}) - input_params[inp.name] = param - mellon_name = param.pop("name", inp.name) - if mellon_name != inp.name: - self.name_mapping[inp.name] = mellon_name - continue - - if inp.name not in DEFAULT_PARAM_MAPS and not inp.required and not get_group_name(inp.name): - continue - - if inp.name in DEFAULT_PARAM_MAPS: - # first check if it's in the default param map, if so, directly use that - param = DEFAULT_PARAM_MAPS[inp.name].copy() - elif get_group_name(inp.name): - param = get_group_name(inp.name) - if inp.name not in self.name_mapping: - self.name_mapping[inp.name] = param - else: - # if not, check if it's in the SDXL input schema, if so, - # 1. use the type hint to determine the type - # 2. use the default param dict for the type e.g. if "steps" is a "int" type, {"steps": {"type": "int", "default": 0, "min": 0}} - if inp.type_hint is not None: - type_str = str(inp.type_hint).lower() - else: - inp_spec = SDXL_PARAM_SCHEMA.get(inp.name, None) - type_str = str(inp_spec.type_hint).lower() if inp_spec else "" - for type_key, type_param in DEFAULT_TYPE_MAPS.items(): - if type_key in type_str: - param = type_param.copy() - param["label"] = inp.name - param["display"] = "input" - break - else: - param = inp.name - # add the param dict to the inp_params dict - input_params[inp.name] = param - - component_params = {} - for comp in self.blocks.expected_components: - param = kwargs.pop(comp.name, None) - if param: - component_params[comp.name] = param - mellon_name = param.pop("name", comp.name) - if mellon_name != comp.name: - self.name_mapping[comp.name] = mellon_name - continue - - to_exclude = False - for exclude_key in DEFAULT_EXCLUDE_MODEL_KEYS: - if exclude_key in comp.name: - to_exclude = True - break - if to_exclude: - continue - - if get_group_name(comp.name): - param = get_group_name(comp.name) - if comp.name not in self.name_mapping: - self.name_mapping[comp.name] = param - elif comp.name in DEFAULT_MODEL_KEYS: - param = {"label": comp.name, "type": "diffusers_auto_model", "display": "input"} - else: - param = comp.name - # add the param dict to the model_params dict - component_params[comp.name] = param - - output_params = {} - if isinstance(self.blocks, SequentialPipelineBlocks): - last_block_name = list(self.blocks.sub_blocks.keys())[-1] - outputs = self.blocks.sub_blocks[last_block_name].intermediate_outputs - else: - outputs = self.blocks.intermediate_outputs - - for out in outputs: - param = kwargs.pop(out.name, None) - if param: - output_params[out.name] = param - mellon_name = param.pop("name", out.name) - if mellon_name != out.name: - self.name_mapping[out.name] = mellon_name - continue - - if out.name in DEFAULT_PARAM_MAPS: - param = DEFAULT_PARAM_MAPS[out.name].copy() - param["display"] = "output" - else: - group_name = get_group_name(out.name) - if group_name: - param = group_name - if out.name not in self.name_mapping: - self.name_mapping[out.name] = param - else: - param = out.name - # add the param dict to the outputs dict - output_params[out.name] = param - - if len(kwargs) > 0: - logger.warning(f"Unused kwargs: {kwargs}") - - register_dict = { - "category": category, - "label": label, - "input_params": input_params, - "component_params": component_params, - "output_params": output_params, - "name_mapping": self.name_mapping, - } - self.register_to_config(**register_dict) - - def setup(self, components_manager, collection=None): - self.pipeline = self.blocks.init_pipeline(components_manager=components_manager, collection=collection) - self._components_manager = components_manager - - @property - def mellon_config(self): - return self._convert_to_mellon_config() - - def _convert_to_mellon_config(self): - node = {} - node["label"] = self.config.label - node["category"] = self.config.category - - node_param = {} - for inp_name, inp_param in self.config.input_params.items(): - if inp_name in self.name_mapping: - mellon_name = self.name_mapping[inp_name] - else: - mellon_name = inp_name - if isinstance(inp_param, str): - param = { - "label": inp_param, - "type": inp_param, - "display": "input", - } - else: - param = inp_param - - if mellon_name not in node_param: - node_param[mellon_name] = param - else: - logger.debug(f"Input param {mellon_name} already exists in node_param, skipping {inp_name}") - - for comp_name, comp_param in self.config.component_params.items(): - if comp_name in self.name_mapping: - mellon_name = self.name_mapping[comp_name] - else: - mellon_name = comp_name - if isinstance(comp_param, str): - param = { - "label": comp_param, - "type": comp_param, - "display": "input", - } - else: - param = comp_param - - if mellon_name not in node_param: - node_param[mellon_name] = param - else: - logger.debug(f"Component param {comp_param} already exists in node_param, skipping {comp_name}") - - for out_name, out_param in self.config.output_params.items(): - if out_name in self.name_mapping: - mellon_name = self.name_mapping[out_name] - else: - mellon_name = out_name - if isinstance(out_param, str): - param = { - "label": out_param, - "type": out_param, - "display": "output", - } - else: - param = out_param - - if mellon_name not in node_param: - node_param[mellon_name] = param - else: - logger.debug(f"Output param {out_param} already exists in node_param, skipping {out_name}") - node["params"] = node_param - return node - - def save_mellon_config(self, file_path): - """ - Save the Mellon configuration to a JSON file. - - Args: - file_path (str or Path): Path where the JSON file will be saved - - Returns: - Path: Path to the saved config file - """ - file_path = Path(file_path) - - # Create directory if it doesn't exist - os.makedirs(file_path.parent, exist_ok=True) - - # Create a combined dictionary with module definition and name mapping - config = {"module": self.mellon_config, "name_mapping": self.name_mapping} - - # Save the config to file - with open(file_path, "w", encoding="utf-8") as f: - json.dump(config, f, indent=2) - - logger.info(f"Mellon config and name mapping saved to {file_path}") - - return file_path - - @classmethod - def load_mellon_config(cls, file_path): - """ - Load a Mellon configuration from a JSON file. - - Args: - file_path (str or Path): Path to the JSON file containing Mellon config - - Returns: - dict: The loaded combined configuration containing 'module' and 'name_mapping' - """ - file_path = Path(file_path) - - if not file_path.exists(): - raise FileNotFoundError(f"Config file not found: {file_path}") - - with open(file_path, "r", encoding="utf-8") as f: - config = json.load(f) - - logger.info(f"Mellon config loaded from {file_path}") - - return config - - def process_inputs(self, **kwargs): - params_components = {} - for comp_name, comp_param in self.config.component_params.items(): - logger.debug(f"component: {comp_name}") - mellon_comp_name = self.name_mapping.get(comp_name, comp_name) - if mellon_comp_name in kwargs: - if isinstance(kwargs[mellon_comp_name], dict) and comp_name in kwargs[mellon_comp_name]: - comp = kwargs[mellon_comp_name].pop(comp_name) - else: - comp = kwargs.pop(mellon_comp_name) - if comp: - params_components[comp_name] = self._components_manager.get_one(comp["model_id"]) - - params_run = {} - for inp_name, inp_param in self.config.input_params.items(): - logger.debug(f"input: {inp_name}") - mellon_inp_name = self.name_mapping.get(inp_name, inp_name) - if mellon_inp_name in kwargs: - if isinstance(kwargs[mellon_inp_name], dict) and inp_name in kwargs[mellon_inp_name]: - inp = kwargs[mellon_inp_name].pop(inp_name) - else: - inp = kwargs.pop(mellon_inp_name) - if inp is not None: - params_run[inp_name] = inp - - return_output_names = list(self.config.output_params.keys()) - - return params_components, params_run, return_output_names - - def execute(self, **kwargs): - params_components, params_run, return_output_names = self.process_inputs(**kwargs) - - self.pipeline.update_components(**params_components) - output = self.pipeline(**params_run, output=return_output_names) - return output diff --git a/src/diffusers/modular_pipelines/qwenimage/__init__.py b/src/diffusers/modular_pipelines/qwenimage/__init__.py index ae4ec4799fbc..2e6af4495b37 100644 --- a/src/diffusers/modular_pipelines/qwenimage/__init__.py +++ b/src/diffusers/modular_pipelines/qwenimage/__init__.py @@ -21,26 +21,14 @@ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) else: - _import_structure["encoders"] = ["QwenImageTextEncoderStep"] - _import_structure["modular_blocks"] = [ - "ALL_BLOCKS", - "AUTO_BLOCKS", - "CONTROLNET_BLOCKS", - "EDIT_AUTO_BLOCKS", - "EDIT_BLOCKS", - "EDIT_INPAINT_BLOCKS", - "EDIT_PLUS_AUTO_BLOCKS", - "EDIT_PLUS_BLOCKS", - "IMAGE2IMAGE_BLOCKS", - "INPAINT_BLOCKS", - "TEXT2IMAGE_BLOCKS", - "QwenImageAutoBlocks", - "QwenImageEditAutoBlocks", - "QwenImageEditPlusAutoBlocks", - ] + _import_structure["modular_blocks_qwenimage"] = ["QwenImageAutoBlocks"] + _import_structure["modular_blocks_qwenimage_edit"] = ["QwenImageEditAutoBlocks"] + _import_structure["modular_blocks_qwenimage_edit_plus"] = ["QwenImageEditPlusAutoBlocks"] + _import_structure["modular_blocks_qwenimage_layered"] = ["QwenImageLayeredAutoBlocks"] _import_structure["modular_pipeline"] = [ "QwenImageEditModularPipeline", "QwenImageEditPlusModularPipeline", + "QwenImageLayeredModularPipeline", "QwenImageModularPipeline", ] @@ -51,28 +39,14 @@ except OptionalDependencyNotAvailable: from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 else: - from .encoders import ( - QwenImageTextEncoderStep, - ) - from .modular_blocks import ( - ALL_BLOCKS, - AUTO_BLOCKS, - CONTROLNET_BLOCKS, - EDIT_AUTO_BLOCKS, - EDIT_BLOCKS, - EDIT_INPAINT_BLOCKS, - EDIT_PLUS_AUTO_BLOCKS, - EDIT_PLUS_BLOCKS, - IMAGE2IMAGE_BLOCKS, - INPAINT_BLOCKS, - TEXT2IMAGE_BLOCKS, - QwenImageAutoBlocks, - QwenImageEditAutoBlocks, - QwenImageEditPlusAutoBlocks, - ) + from .modular_blocks_qwenimage import QwenImageAutoBlocks + from .modular_blocks_qwenimage_edit import QwenImageEditAutoBlocks + from .modular_blocks_qwenimage_edit_plus import QwenImageEditPlusAutoBlocks + from .modular_blocks_qwenimage_layered import QwenImageLayeredAutoBlocks from .modular_pipeline import ( QwenImageEditModularPipeline, QwenImageEditPlusModularPipeline, + QwenImageLayeredModularPipeline, QwenImageModularPipeline, ) else: diff --git a/src/diffusers/modular_pipelines/qwenimage/before_denoise.py b/src/diffusers/modular_pipelines/qwenimage/before_denoise.py index bd92d403539e..51b5c6ac8c3d 100644 --- a/src/diffusers/modular_pipelines/qwenimage/before_denoise.py +++ b/src/diffusers/modular_pipelines/qwenimage/before_denoise.py @@ -13,7 +13,6 @@ # limitations under the License. import inspect -from typing import List, Optional, Tuple, Union import numpy as np import torch @@ -23,7 +22,7 @@ from ...utils.torch_utils import randn_tensor, unwrap_module from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam -from .modular_pipeline import QwenImageModularPipeline, QwenImagePachifier +from .modular_pipeline import QwenImageLayeredPachifier, QwenImageModularPipeline, QwenImagePachifier # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.calculate_shift @@ -43,10 +42,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -61,15 +60,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -113,10 +112,45 @@ def get_timesteps(scheduler, num_inference_steps, strength): return timesteps, num_inference_steps - t_start -# Prepare Latents steps +# ==================== +# 1. PREPARE LATENTS +# ==================== +# auto_docstring class QwenImagePrepareLatentsStep(ModularPipelineBlocks): + """ + Prepare initial random noise for the generation process + + Components: + pachifier (`QwenImagePachifier`) + + Inputs: + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + batch_size (`int`, *optional*, defaults to 1): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + dtype (`dtype`, *optional*, defaults to torch.float32): + The dtype of the model inputs, can be generated in input step. + + Outputs: + height (`int`): + if not set, updated to default value + width (`int`): + if not set, updated to default value + latents (`Tensor`): + The initial latents to use for the denoising process + """ + model_name = "qwenimage" @property @@ -124,36 +158,28 @@ def description(self) -> str: return "Prepare initial random noise for the generation process" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("pachifier", QwenImagePachifier, default_creation_method="from_config"), ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ - InputParam("latents"), - InputParam(name="height"), - InputParam(name="width"), - InputParam(name="num_images_per_prompt", default=1), - InputParam(name="generator"), - InputParam( - name="batch_size", - required=True, - type_hint=int, - description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can be generated in input step.", - ), - InputParam( - name="dtype", - required=True, - type_hint=torch.dtype, - description="The dtype of the model inputs, can be generated in input step.", - ), + InputParam.template("latents"), + InputParam.template("height"), + InputParam.template("width"), + InputParam.template("num_images_per_prompt"), + InputParam.template("generator"), + InputParam.template("batch_size"), + InputParam.template("dtype"), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ + OutputParam(name="height", type_hint=int, description="if not set, updated to default value"), + OutputParam(name="width", type_hint=int, description="if not set, updated to default value"), OutputParam( name="latents", type_hint=torch.Tensor, @@ -207,7 +233,150 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) - return components, state +# auto_docstring +class QwenImageLayeredPrepareLatentsStep(ModularPipelineBlocks): + """ + Prepare initial random noise (B, layers+1, C, H, W) for the generation process + + Components: + pachifier (`QwenImageLayeredPachifier`) + + Inputs: + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + layers (`int`, *optional*, defaults to 4): + Number of layers to extract from the image + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + batch_size (`int`, *optional*, defaults to 1): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + dtype (`dtype`, *optional*, defaults to torch.float32): + The dtype of the model inputs, can be generated in input step. + + Outputs: + height (`int`): + if not set, updated to default value + width (`int`): + if not set, updated to default value + latents (`Tensor`): + The initial latents to use for the denoising process + """ + + model_name = "qwenimage-layered" + + @property + def description(self) -> str: + return "Prepare initial random noise (B, layers+1, C, H, W) for the generation process" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("pachifier", QwenImageLayeredPachifier, default_creation_method="from_config"), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("latents"), + InputParam.template("height"), + InputParam.template("width"), + InputParam.template("layers"), + InputParam.template("num_images_per_prompt"), + InputParam.template("generator"), + InputParam.template("batch_size"), + InputParam.template("dtype"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam(name="height", type_hint=int, description="if not set, updated to default value"), + OutputParam(name="width", type_hint=int, description="if not set, updated to default value"), + OutputParam( + name="latents", + type_hint=torch.Tensor, + description="The initial latents to use for the denoising process", + ), + ] + + @staticmethod + def check_inputs(height, width, vae_scale_factor): + if height is not None and height % (vae_scale_factor * 2) != 0: + raise ValueError(f"Height must be divisible by {vae_scale_factor * 2} but is {height}") + + if width is not None and width % (vae_scale_factor * 2) != 0: + raise ValueError(f"Width must be divisible by {vae_scale_factor * 2} but is {width}") + + @torch.no_grad() + def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + self.check_inputs( + height=block_state.height, + width=block_state.width, + vae_scale_factor=components.vae_scale_factor, + ) + + device = components._execution_device + batch_size = block_state.batch_size * block_state.num_images_per_prompt + + # we can update the height and width here since it's used to generate the initial + block_state.height = block_state.height or components.default_height + block_state.width = block_state.width or components.default_width + + # VAE applies 8x compression on images but we must also account for packing which requires + # latent height and width to be divisible by 2. + latent_height = 2 * (int(block_state.height) // (components.vae_scale_factor * 2)) + latent_width = 2 * (int(block_state.width) // (components.vae_scale_factor * 2)) + + shape = (batch_size, block_state.layers + 1, components.num_channels_latents, latent_height, latent_width) + if isinstance(block_state.generator, list) and len(block_state.generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(block_state.generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + if block_state.latents is None: + block_state.latents = randn_tensor( + shape, generator=block_state.generator, device=device, dtype=block_state.dtype + ) + block_state.latents = components.pachifier.pack_latents(block_state.latents) + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring class QwenImagePrepareLatentsWithStrengthStep(ModularPipelineBlocks): + """ + Step that adds noise to image latents for image-to-image/inpainting. Should be run after set_timesteps, + prepare_latents. Both noise and image latents should alreadybe patchified. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + latents (`Tensor`): + The initial random noised, can be generated in prepare latent step. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. (Can be + generated from vae encoder and updated in input step.) + timesteps (`Tensor`): + The timesteps to use for the denoising process. Can be generated in set_timesteps step. + + Outputs: + initial_noise (`Tensor`): + The initial random noised used for inpainting denoising. + latents (`Tensor`): + The scaled noisy latents to use for inpainting/image-to-image denoising. + """ + model_name = "qwenimage" @property @@ -215,13 +384,13 @@ def description(self) -> str: return "Step that adds noise to image latents for image-to-image/inpainting. Should be run after set_timesteps, prepare_latents. Both noise and image latents should alreadybe patchified." @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam( name="latents", @@ -229,12 +398,7 @@ def inputs(self) -> List[InputParam]: type_hint=torch.Tensor, description="The initial random noised, can be generated in prepare latent step.", ), - InputParam( - name="image_latents", - required=True, - type_hint=torch.Tensor, - description="The image latents to use for the denoising process. Can be generated in vae encoder and packed in input step.", - ), + InputParam.template("image_latents", note="Can be generated from vae encoder and updated in input step."), InputParam( name="timesteps", required=True, @@ -244,13 +408,18 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( name="initial_noise", type_hint=torch.Tensor, description="The initial random noised used for inpainting denoising.", ), + OutputParam( + name="latents", + type_hint=torch.Tensor, + description="The scaled noisy latents to use for inpainting/image-to-image denoising.", + ), ] @staticmethod @@ -288,7 +457,29 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) - return components, state +# auto_docstring class QwenImageCreateMaskLatentsStep(ModularPipelineBlocks): + """ + Step that creates mask latents from preprocessed mask_image by interpolating to latent space. + + Components: + pachifier (`QwenImagePachifier`) + + Inputs: + processed_mask_image (`Tensor`): + The processed mask to use for the inpainting process. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + dtype (`dtype`, *optional*, defaults to torch.float32): + The dtype of the model inputs, can be generated in input step. + + Outputs: + mask (`Tensor`): + The mask to use for the inpainting process. + """ + model_name = "qwenimage" @property @@ -296,13 +487,13 @@ def description(self) -> str: return "Step that creates mask latents from preprocessed mask_image by interpolating to latent space." @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("pachifier", QwenImagePachifier, default_creation_method="from_config"), ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam( name="processed_mask_image", @@ -310,13 +501,13 @@ def inputs(self) -> List[InputParam]: type_hint=torch.Tensor, description="The processed mask to use for the inpainting process.", ), - InputParam(name="height", required=True), - InputParam(name="width", required=True), - InputParam(name="dtype", required=True), + InputParam.template("height", required=True), + InputParam.template("width", required=True), + InputParam.template("dtype"), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( name="mask", type_hint=torch.Tensor, description="The mask to use for the inpainting process." @@ -351,37 +542,59 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) - return components, state -# Set Timesteps steps +# ==================== +# 2. SET TIMESTEPS +# ==================== +# auto_docstring class QwenImageSetTimestepsStep(ModularPipelineBlocks): + """ + Step that sets the scheduler's timesteps for text-to-image generation. Should be run after prepare latents step. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + latents (`Tensor`): + The initial random noised latents for the denoising process. Can be generated in prepare latents step. + + Outputs: + timesteps (`Tensor`): + The timesteps to use for the denoising process + """ + model_name = "qwenimage" @property def description(self) -> str: - return "Step that sets the the scheduler's timesteps for text-to-image generation. Should be run after prepare latents step." + return "Step that sets the scheduler's timesteps for text-to-image generation. Should be run after prepare latents step." @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ - InputParam(name="num_inference_steps", default=50), - InputParam(name="sigmas"), + InputParam.template("num_inference_steps"), + InputParam.template("sigmas"), InputParam( name="latents", required=True, type_hint=torch.Tensor, - description="The latents to use for the denoising process, used to calculate the image sequence length.", + description="The initial random noised latents for the denoising process. Can be generated in prepare latents step.", ), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( name="timesteps", type_hint=torch.Tensor, description="The timesteps to use for the denoising process" @@ -420,40 +633,150 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) - return components, state +# auto_docstring +class QwenImageLayeredSetTimestepsStep(ModularPipelineBlocks): + """ + Set timesteps step for QwenImage Layered with custom mu calculation based on image_latents. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + + Outputs: + timesteps (`Tensor`): + The timesteps to use for the denoising process. + """ + + model_name = "qwenimage-layered" + + @property + def description(self) -> str: + return "Set timesteps step for QwenImage Layered with custom mu calculation based on image_latents." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("num_inference_steps"), + InputParam.template("sigmas"), + InputParam.template("image_latents"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="timesteps", type_hint=torch.Tensor, description="The timesteps to use for the denoising process." + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + + # Layered-specific mu calculation + base_seqlen = 256 * 256 / 16 / 16 # = 256 + mu = (block_state.image_latents.shape[1] / base_seqlen) ** 0.5 + + # Default sigmas if not provided + sigmas = ( + np.linspace(1.0, 1 / block_state.num_inference_steps, block_state.num_inference_steps) + if block_state.sigmas is None + else block_state.sigmas + ) + + block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps( + components.scheduler, + block_state.num_inference_steps, + device, + sigmas=sigmas, + mu=mu, + ) + + components.scheduler.set_begin_index(0) + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring class QwenImageSetTimestepsWithStrengthStep(ModularPipelineBlocks): + """ + Step that sets the scheduler's timesteps for image-to-image generation, and inpainting. Should be run after prepare + latents step. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + latents (`Tensor`): + The latents to use for the denoising process. Can be generated in prepare latents step. + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. + + Outputs: + timesteps (`Tensor`): + The timesteps to use for the denoising process. + num_inference_steps (`int`): + The number of denoising steps to perform at inference time. Updated based on strength. + """ + model_name = "qwenimage" @property def description(self) -> str: - return "Step that sets the the scheduler's timesteps for image-to-image generation, and inpainting. Should be run after prepare latents step." + return "Step that sets the scheduler's timesteps for image-to-image generation, and inpainting. Should be run after prepare latents step." @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ - InputParam(name="num_inference_steps", default=50), - InputParam(name="sigmas"), + InputParam.template("num_inference_steps"), + InputParam.template("sigmas"), InputParam( - name="latents", + "latents", required=True, type_hint=torch.Tensor, - description="The latents to use for the denoising process, used to calculate the image sequence length.", + description="The latents to use for the denoising process. Can be generated in prepare latents step.", ), - InputParam(name="strength", default=0.9), + InputParam.template("strength", default=0.9), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( name="timesteps", type_hint=torch.Tensor, - description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.", + description="The timesteps to use for the denoising process.", + ), + OutputParam( + name="num_inference_steps", + type_hint=int, + description="The number of denoising steps to perform at inference time. Updated based on strength.", ), ] @@ -493,12 +816,36 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) - return components, state -# other inputs for denoiser +# ==================== +# 3. OTHER INPUTS FOR DENOISER +# ==================== ## RoPE inputs for denoiser +# auto_docstring class QwenImageRoPEInputsStep(ModularPipelineBlocks): + """ + Step that prepares the RoPE inputs for the denoising process. Should be place after prepare_latents step + + Inputs: + batch_size (`int`, *optional*, defaults to 1): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + + Outputs: + img_shapes (`list`): + The shapes of the images latents, used for RoPE calculation + """ + model_name = "qwenimage" @property @@ -508,34 +855,23 @@ def description(self) -> str: ) @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ - InputParam(name="batch_size", required=True), - InputParam(name="height", required=True), - InputParam(name="width", required=True), - InputParam(name="prompt_embeds_mask"), - InputParam(name="negative_prompt_embeds_mask"), + InputParam.template("batch_size"), + InputParam.template("height", required=True), + InputParam.template("width", required=True), + InputParam.template("prompt_embeds_mask"), + InputParam.template("negative_prompt_embeds_mask"), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( name="img_shapes", - type_hint=List[List[Tuple[int, int, int]]], - description="The shapes of the images latents, used for RoPE calculation", - ), - OutputParam( - name="txt_seq_lens", - kwargs_type="denoiser_input_fields", - type_hint=List[int], - description="The sequence lengths of the prompt embeds, used for RoPE calculation", - ), - OutputParam( - name="negative_txt_seq_lens", kwargs_type="denoiser_input_fields", - type_hint=List[int], - description="The sequence lengths of the negative prompt embeds, used for RoPE calculation", + type_hint=list[list[tuple[int, int, int]]], + description="The shapes of the images latents, used for RoPE calculation", ), ] @@ -551,21 +887,40 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) - ) ] ] * block_state.batch_size - block_state.txt_seq_lens = ( - block_state.prompt_embeds_mask.sum(dim=1).tolist() if block_state.prompt_embeds_mask is not None else None - ) - block_state.negative_txt_seq_lens = ( - block_state.negative_prompt_embeds_mask.sum(dim=1).tolist() - if block_state.negative_prompt_embeds_mask is not None - else None - ) self.set_block_state(state, block_state) return components, state +# auto_docstring class QwenImageEditRoPEInputsStep(ModularPipelineBlocks): + """ + Step that prepares the RoPE inputs for denoising process. This is used in QwenImage Edit. Should be placed after + prepare_latents step + + Inputs: + batch_size (`int`, *optional*, defaults to 1): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + image_height (`int`): + The height of the reference image. Can be generated in input step. + image_width (`int`): + The width of the reference image. Can be generated in input step. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + + Outputs: + img_shapes (`list`): + The shapes of the images latents, used for RoPE calculation + """ + model_name = "qwenimage" @property @@ -573,36 +928,35 @@ def description(self) -> str: return "Step that prepares the RoPE inputs for denoising process. This is used in QwenImage Edit. Should be placed after prepare_latents step" @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ - InputParam(name="batch_size", required=True), - InputParam(name="image_height", required=True), - InputParam(name="image_width", required=True), - InputParam(name="height", required=True), - InputParam(name="width", required=True), - InputParam(name="prompt_embeds_mask"), - InputParam(name="negative_prompt_embeds_mask"), + InputParam.template("batch_size"), + InputParam( + name="image_height", + required=True, + type_hint=int, + description="The height of the reference image. Can be generated in input step.", + ), + InputParam( + name="image_width", + required=True, + type_hint=int, + description="The width of the reference image. Can be generated in input step.", + ), + InputParam.template("height", required=True), + InputParam.template("width", required=True), + InputParam.template("prompt_embeds_mask"), + InputParam.template("negative_prompt_embeds_mask"), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( name="img_shapes", - type_hint=List[List[Tuple[int, int, int]]], - description="The shapes of the images latents, used for RoPE calculation", - ), - OutputParam( - name="txt_seq_lens", kwargs_type="denoiser_input_fields", - type_hint=List[int], - description="The sequence lengths of the prompt embeds, used for RoPE calculation", - ), - OutputParam( - name="negative_txt_seq_lens", - kwargs_type="denoiser_input_fields", - type_hint=List[int], - description="The sequence lengths of the negative prompt embeds, used for RoPE calculation", + type_hint=list[list[tuple[int, int, int]]], + description="The shapes of the images latents, used for RoPE calculation", ), ] @@ -625,33 +979,111 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) - ] ] * block_state.batch_size - block_state.txt_seq_lens = ( - block_state.prompt_embeds_mask.sum(dim=1).tolist() if block_state.prompt_embeds_mask is not None else None - ) - block_state.negative_txt_seq_lens = ( - block_state.negative_prompt_embeds_mask.sum(dim=1).tolist() - if block_state.negative_prompt_embeds_mask is not None - else None - ) - self.set_block_state(state, block_state) return components, state -class QwenImageEditPlusRoPEInputsStep(QwenImageEditRoPEInputsStep): +# auto_docstring +class QwenImageEditPlusRoPEInputsStep(ModularPipelineBlocks): + """ + Step that prepares the RoPE inputs for denoising process. This is used in QwenImage Edit Plus. + Unlike Edit, Edit Plus handles lists of image_height/image_width for multiple reference images. Should be placed + after prepare_latents step. + + Inputs: + batch_size (`int`, *optional*, defaults to 1): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + image_height (`list`): + The heights of the reference images. Can be generated in input step. + image_width (`list`): + The widths of the reference images. Can be generated in input step. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + + Outputs: + img_shapes (`list`): + The shapes of the image latents, used for RoPE calculation + txt_seq_lens (`list`): + The sequence lengths of the prompt embeds, used for RoPE calculation + negative_txt_seq_lens (`list`): + The sequence lengths of the negative prompt embeds, used for RoPE calculation + """ + model_name = "qwenimage-edit-plus" + @property + def description(self) -> str: + return ( + "Step that prepares the RoPE inputs for denoising process. This is used in QwenImage Edit Plus.\n" + "Unlike Edit, Edit Plus handles lists of image_height/image_width for multiple reference images.\n" + "Should be placed after prepare_latents step." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("batch_size"), + InputParam( + name="image_height", + required=True, + type_hint=list[int], + description="The heights of the reference images. Can be generated in input step.", + ), + InputParam( + name="image_width", + required=True, + type_hint=list[int], + description="The widths of the reference images. Can be generated in input step.", + ), + InputParam.template("height", required=True), + InputParam.template("width", required=True), + InputParam.template("prompt_embeds_mask"), + InputParam.template("negative_prompt_embeds_mask"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="img_shapes", + kwargs_type="denoiser_input_fields", + type_hint=list[list[tuple[int, int, int]]], + description="The shapes of the image latents, used for RoPE calculation", + ), + OutputParam( + name="txt_seq_lens", + kwargs_type="denoiser_input_fields", + type_hint=list[int], + description="The sequence lengths of the prompt embeds, used for RoPE calculation", + ), + OutputParam( + name="negative_txt_seq_lens", + kwargs_type="denoiser_input_fields", + type_hint=list[int], + description="The sequence lengths of the negative prompt embeds, used for RoPE calculation", + ), + ] + def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) vae_scale_factor = components.vae_scale_factor + + # Edit Plus: image_height and image_width are lists block_state.img_shapes = [ [ (1, block_state.height // vae_scale_factor // 2, block_state.width // vae_scale_factor // 2), *[ - (1, vae_height // vae_scale_factor // 2, vae_width // vae_scale_factor // 2) - for vae_height, vae_width in zip(block_state.image_height, block_state.image_width) + (1, img_height // vae_scale_factor // 2, img_width // vae_scale_factor // 2) + for img_height, img_width in zip(block_state.image_height, block_state.image_width) ], ] ] * block_state.batch_size @@ -670,12 +1102,150 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) - return components, state +# auto_docstring +class QwenImageLayeredRoPEInputsStep(ModularPipelineBlocks): + """ + Step that prepares the RoPE inputs for the denoising process. Should be place after prepare_latents step + + Inputs: + batch_size (`int`, *optional*, defaults to 1): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + layers (`int`, *optional*, defaults to 4): + Number of layers to extract from the image + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + + Outputs: + img_shapes (`list`): + The shapes of the image latents, used for RoPE calculation + txt_seq_lens (`list`): + The sequence lengths of the prompt embeds, used for RoPE calculation + negative_txt_seq_lens (`list`): + The sequence lengths of the negative prompt embeds, used for RoPE calculation + additional_t_cond (`Tensor`): + The additional t cond, used for RoPE calculation + """ + + model_name = "qwenimage-layered" + + @property + def description(self) -> str: + return ( + "Step that prepares the RoPE inputs for the denoising process. Should be place after prepare_latents step" + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("batch_size"), + InputParam.template("layers"), + InputParam.template("height", required=True), + InputParam.template("width", required=True), + InputParam.template("prompt_embeds_mask"), + InputParam.template("negative_prompt_embeds_mask"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="img_shapes", + type_hint=list[list[tuple[int, int, int]]], + kwargs_type="denoiser_input_fields", + description="The shapes of the image latents, used for RoPE calculation", + ), + OutputParam( + name="txt_seq_lens", + type_hint=list[int], + kwargs_type="denoiser_input_fields", + description="The sequence lengths of the prompt embeds, used for RoPE calculation", + ), + OutputParam( + name="negative_txt_seq_lens", + type_hint=list[int], + kwargs_type="denoiser_input_fields", + description="The sequence lengths of the negative prompt embeds, used for RoPE calculation", + ), + OutputParam( + name="additional_t_cond", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="The additional t cond, used for RoPE calculation", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + + # All shapes are the same for Layered + shape = ( + 1, + block_state.height // components.vae_scale_factor // 2, + block_state.width // components.vae_scale_factor // 2, + ) + + # layers+1 output shapes + 1 condition shape (all same) + block_state.img_shapes = [[shape] * (block_state.layers + 2)] * block_state.batch_size + + # txt_seq_lens + block_state.txt_seq_lens = ( + block_state.prompt_embeds_mask.sum(dim=1).tolist() if block_state.prompt_embeds_mask is not None else None + ) + block_state.negative_txt_seq_lens = ( + block_state.negative_prompt_embeds_mask.sum(dim=1).tolist() + if block_state.negative_prompt_embeds_mask is not None + else None + ) + + block_state.additional_t_cond = torch.tensor([0] * block_state.batch_size).to(device=device, dtype=torch.long) + + self.set_block_state(state, block_state) + return components, state + + ## ControlNet inputs for denoiser + + +# auto_docstring class QwenImageControlNetBeforeDenoiserStep(ModularPipelineBlocks): + """ + step that prepare inputs for controlnet. Insert before the Denoise Step, after set_timesteps step. + + Components: + controlnet (`QwenImageControlNetModel`) + + Inputs: + control_guidance_start (`float`, *optional*, defaults to 0.0): + When to start applying ControlNet. + control_guidance_end (`float`, *optional*, defaults to 1.0): + When to stop applying ControlNet. + controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0): + Scale for ControlNet conditioning. + control_image_latents (`Tensor`): + The control image latents to use for the denoising process. Can be generated in controlnet vae encoder + step. + timesteps (`Tensor`): + The timesteps to use for the denoising process. Can be generated in set_timesteps step. + + Outputs: + controlnet_keep (`list`): + The controlnet keep values + """ + model_name = "qwenimage" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("controlnet", QwenImageControlNetModel), ] @@ -685,14 +1255,19 @@ def description(self) -> str: return "step that prepare inputs for controlnet. Insert before the Denoise Step, after set_timesteps step." @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ - InputParam("control_guidance_start", default=0.0), - InputParam("control_guidance_end", default=1.0), - InputParam("controlnet_conditioning_scale", default=1.0), - InputParam("control_image_latents", required=True), + InputParam.template("control_guidance_start"), + InputParam.template("control_guidance_end"), + InputParam.template("controlnet_conditioning_scale"), InputParam( - "timesteps", + name="control_image_latents", + required=True, + type_hint=torch.Tensor, + description="The control image latents to use for the denoising process. Can be generated in controlnet vae encoder step.", + ), + InputParam( + name="timesteps", required=True, type_hint=torch.Tensor, description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.", @@ -700,9 +1275,9 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ - OutputParam("controlnet_keep", type_hint=List[float], description="The controlnet keep values"), + OutputParam("controlnet_keep", type_hint=list[float], description="The controlnet keep values"), ] @torch.no_grad() diff --git a/src/diffusers/modular_pipelines/qwenimage/decoders.py b/src/diffusers/modular_pipelines/qwenimage/decoders.py index 6e145f18550a..e4ccb6b8e047 100644 --- a/src/diffusers/modular_pipelines/qwenimage/decoders.py +++ b/src/diffusers/modular_pipelines/qwenimage/decoders.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Union -import numpy as np -import PIL +from typing import Any + import torch from ...configuration_utils import FrozenDict @@ -24,13 +23,37 @@ from ...utils import logging from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam -from .modular_pipeline import QwenImageModularPipeline, QwenImagePachifier +from .modular_pipeline import QwenImageLayeredPachifier, QwenImageModularPipeline, QwenImagePachifier logger = logging.get_logger(__name__) +# after denoising loop (unpack latents) + + +# auto_docstring class QwenImageAfterDenoiseStep(ModularPipelineBlocks): + """ + Step that unpack the latents from 3D tensor (batch_size, sequence_length, channels) into 5D tensor (batch_size, + channels, 1, height, width) + + Components: + pachifier (`QwenImagePachifier`) + + Inputs: + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + latents (`Tensor`): + The latents to decode, can be generated in the denoise step. + + Outputs: + latents (`Tensor`): + The denoisedlatents unpacked to B, C, 1, H, W + """ + model_name = "qwenimage" @property @@ -38,7 +61,7 @@ def description(self) -> str: return "Step that unpack the latents from 3D tensor (batch_size, sequence_length, channels) into 5D tensor (batch_size, channels, 1, height, width)" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: components = [ ComponentSpec("pachifier", QwenImagePachifier, default_creation_method="from_config"), ] @@ -46,15 +69,23 @@ def expected_components(self) -> List[ComponentSpec]: return components @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ - InputParam(name="height", required=True), - InputParam(name="width", required=True), + InputParam.template("height", required=True), + InputParam.template("width", required=True), InputParam( name="latents", required=True, type_hint=torch.Tensor, - description="The latents to decode, can be generated in the denoise step", + description="The latents to decode, can be generated in the denoise step.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="latents", type_hint=torch.Tensor, description="The denoisedlatents unpacked to B, C, 1, H, W" ), ] @@ -71,7 +102,99 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) - return components, state +# auto_docstring +class QwenImageLayeredAfterDenoiseStep(ModularPipelineBlocks): + """ + Unpack latents from (B, seq, C*4) to (B, C, layers+1, H, W) after denoising. + + Components: + pachifier (`QwenImageLayeredPachifier`) + + Inputs: + latents (`Tensor`): + The denoised latents to decode, can be generated in the denoise step. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + layers (`int`, *optional*, defaults to 4): + Number of layers to extract from the image + + Outputs: + latents (`Tensor`): + Denoised latents. (unpacked to B, C, layers+1, H, W) + """ + + model_name = "qwenimage-layered" + + @property + def description(self) -> str: + return "Unpack latents from (B, seq, C*4) to (B, C, layers+1, H, W) after denoising." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("pachifier", QwenImageLayeredPachifier, default_creation_method="from_config"), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="latents", + required=True, + type_hint=torch.Tensor, + description="The denoised latents to decode, can be generated in the denoise step.", + ), + InputParam.template("height", required=True), + InputParam.template("width", required=True), + InputParam.template("layers"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("latents", note="unpacked to B, C, layers+1, H, W"), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + # Unpack: (B, seq, C*4) -> (B, C, layers+1, H, W) + block_state.latents = components.pachifier.unpack_latents( + block_state.latents, + block_state.height, + block_state.width, + block_state.layers, + components.vae_scale_factor, + ) + + self.set_block_state(state, block_state) + return components, state + + +# decode step + + +# auto_docstring class QwenImageDecoderStep(ModularPipelineBlocks): + """ + Step that decodes the latents to images + + Components: + vae (`AutoencoderKLQwenImage`) + + Inputs: + latents (`Tensor`): + The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise + step. + + Outputs: + images (`list`): + Generated images. (tensor output of the vae decoder.) + """ + model_name = "qwenimage" @property @@ -79,7 +202,7 @@ def description(self) -> str: return "Step that decodes the latents to images" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: components = [ ComponentSpec("vae", AutoencoderKLQwenImage), ] @@ -87,25 +210,19 @@ def expected_components(self) -> List[ComponentSpec]: return components @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam( name="latents", required=True, type_hint=torch.Tensor, - description="The latents to decode, can be generated in the denoise step", + description="The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise step.", ), ] @property - def intermediate_outputs(self) -> List[str]: - return [ - OutputParam( - "images", - type_hint=Union[List[PIL.Image.Image], List[torch.Tensor], List[np.array]], - description="The generated images, can be a PIL.Image.Image, torch.Tensor or a numpy array", - ) - ] + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam.template("images", note="tensor output of the vae decoder.")] @torch.no_grad() def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState: @@ -135,7 +252,124 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) - return components, state +# auto_docstring +class QwenImageLayeredDecoderStep(ModularPipelineBlocks): + """ + Decode unpacked latents (B, C, layers+1, H, W) into layer images. + + Components: + vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`) + + Inputs: + latents (`Tensor`): + The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise + step. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "qwenimage-layered" + + @property + def description(self) -> str: + return "Decode unpacked latents (B, C, layers+1, H, W) into layer images." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLQwenImage), + ComponentSpec( + "image_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="latents", + required=True, + type_hint=torch.Tensor, + description="The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise step.", + ), + InputParam.template("output_type"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam.template("images")] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + latents = block_state.latents + + # 1. VAE normalization + latents = latents.to(components.vae.dtype) + latents_mean = ( + torch.tensor(components.vae.config.latents_mean) + .view(1, components.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = 1.0 / torch.tensor(components.vae.config.latents_std).view( + 1, components.vae.config.z_dim, 1, 1, 1 + ).to(latents.device, latents.dtype) + latents = latents / latents_std + latents_mean + + # 2. Reshape for batch decoding: (B, C, layers+1, H, W) -> (B*layers, C, 1, H, W) + b, c, f, h, w = latents.shape + # 3. Remove first frame (composite), keep layers frames + latents = latents[:, :, 1:] + latents = latents.permute(0, 2, 1, 3, 4).reshape(-1, c, 1, h, w) + + # 4. Decode: (B*layers, C, 1, H, W) -> (B*layers, C, H, W) + image = components.vae.decode(latents, return_dict=False)[0] + image = image.squeeze(2) + + # 5. Postprocess - returns flat list of B*layers images + image = components.image_processor.postprocess(image, output_type=block_state.output_type) + + # 6. Chunk into list per batch item + images = [] + for bidx in range(b): + images.append(image[bidx * f : (bidx + 1) * f]) + + block_state.images = images + + self.set_block_state(state, block_state) + return components, state + + +# postprocess the decoded images + + +# auto_docstring class QwenImageProcessImagesOutputStep(ModularPipelineBlocks): + """ + postprocess the generated image + + Components: + image_processor (`VaeImageProcessor`) + + Inputs: + images (`Tensor`): + the generated image tensor from decoders step + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + + Outputs: + images (`list`): + Generated images. + """ + model_name = "qwenimage" @property @@ -143,7 +377,7 @@ def description(self) -> str: return "postprocess the generated image" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec( "image_processor", @@ -154,17 +388,21 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ - InputParam("images", required=True, description="the generated image from decoders step"), InputParam( - name="output_type", - default="pil", - type_hint=str, - description="The type of the output images, can be 'pil', 'np', 'pt'", + name="images", + required=True, + type_hint=torch.Tensor, + description="the generated image tensor from decoders step", ), + InputParam.template("output_type"), ] + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam.template("images")] + @staticmethod def check_inputs(output_type): if output_type not in ["pil", "np", "pt"]: @@ -185,7 +423,28 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState): return components, state +# auto_docstring class QwenImageInpaintProcessImagesOutputStep(ModularPipelineBlocks): + """ + postprocess the generated image, optional apply the mask overally to the original image.. + + Components: + image_mask_processor (`InpaintProcessor`) + + Inputs: + images (`Tensor`): + the generated image tensor from decoders step + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + mask_overlay_kwargs (`dict`, *optional*): + The kwargs for the postprocess step to apply the mask overlay. generated in + InpaintProcessImagesInputStep. + + Outputs: + images (`list`): + Generated images. + """ + model_name = "qwenimage" @property @@ -193,7 +452,7 @@ def description(self) -> str: return "postprocess the generated image, optional apply the mask overally to the original image.." @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec( "image_mask_processor", @@ -204,18 +463,26 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ - InputParam("images", required=True, description="the generated image from decoders step"), InputParam( - name="output_type", - default="pil", - type_hint=str, - description="The type of the output images, can be 'pil', 'np', 'pt'", + name="images", + required=True, + type_hint=torch.Tensor, + description="the generated image tensor from decoders step", + ), + InputParam.template("output_type"), + InputParam( + name="mask_overlay_kwargs", + type_hint=dict[str, Any], + description="The kwargs for the postprocess step to apply the mask overlay. generated in InpaintProcessImagesInputStep.", ), - InputParam("mask_overlay_kwargs"), ] + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam.template("images")] + @staticmethod def check_inputs(output_type, mask_overlay_kwargs): if output_type not in ["pil", "np", "pt"]: diff --git a/src/diffusers/modular_pipelines/qwenimage/denoise.py b/src/diffusers/modular_pipelines/qwenimage/denoise.py index 49acd2dc0295..de8ea05c5047 100644 --- a/src/diffusers/modular_pipelines/qwenimage/denoise.py +++ b/src/diffusers/modular_pipelines/qwenimage/denoise.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Tuple +import inspect import torch @@ -28,7 +28,12 @@ logger = logging.get_logger(__name__) +# ==================== +# 1. LOOP STEPS (run at each denoising step) +# ==================== + +# loop step:before denoiser class QwenImageLoopBeforeDenoiser(ModularPipelineBlocks): model_name = "qwenimage" @@ -41,10 +46,10 @@ def description(self) -> str: ) @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam( - "latents", + name="latents", required=True, type_hint=torch.Tensor, description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.", @@ -60,7 +65,7 @@ def __call__(self, components: QwenImageModularPipeline, block_state: BlockState class QwenImageEditLoopBeforeDenoiser(ModularPipelineBlocks): - model_name = "qwenimage" + model_name = "qwenimage-edit" @property def description(self) -> str: @@ -71,20 +76,15 @@ def description(self) -> str: ) @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam( - "latents", + name="latents", required=True, type_hint=torch.Tensor, description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.", ), - InputParam( - "image_latents", - required=True, - type_hint=torch.Tensor, - description="The initial image latents to use for the denoising process. Can be encoded in vae_encoder step and packed in prepare_image_latents step.", - ), + InputParam.template("image_latents"), ] @torch.no_grad() @@ -100,7 +100,7 @@ class QwenImageLoopBeforeDenoiserControlNet(ModularPipelineBlocks): model_name = "qwenimage" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec( "guider", @@ -120,7 +120,7 @@ def description(self) -> str: ) @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam( "control_image_latents", @@ -128,29 +128,12 @@ def inputs(self) -> List[InputParam]: type_hint=torch.Tensor, description="The control image to use for the denoising process. Can be generated in prepare_controlnet_inputs step.", ), + InputParam.template("controlnet_conditioning_scale", note="updated in prepare_controlnet_inputs step."), InputParam( - "controlnet_conditioning_scale", - type_hint=float, - description="The controlnet conditioning scale value to use for the denoising process. Can be generated in prepare_controlnet_inputs step.", - ), - InputParam( - "controlnet_keep", - required=True, - type_hint=List[float], - description="The controlnet keep values to use for the denoising process. Can be generated in prepare_controlnet_inputs step.", - ), - InputParam( - "num_inference_steps", + name="controlnet_keep", required=True, - type_hint=int, - description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.", - ), - InputParam( - kwargs_type="denoiser_input_fields", - description=( - "All conditional model inputs for the denoiser. " - "It should contain prompt_embeds/negative_prompt_embeds, txt_seq_lens/negative_txt_seq_lens." - ), + type_hint=list[float], + description="The controlnet keep values. Can be generated in prepare_controlnet_inputs step.", ), ] @@ -176,7 +159,6 @@ def __call__(self, components: QwenImageModularPipeline, block_state: BlockState img_shapes=block_state.img_shapes, encoder_hidden_states=block_state.prompt_embeds, encoder_hidden_states_mask=block_state.prompt_embeds_mask, - txt_seq_lens=block_state.txt_seq_lens, return_dict=False, ) @@ -185,6 +167,7 @@ def __call__(self, components: QwenImageModularPipeline, block_state: BlockState return components, block_state +# loop step:denoiser class QwenImageLoopDenoiser(ModularPipelineBlocks): model_name = "qwenimage" @@ -197,7 +180,7 @@ def description(self) -> str: ) @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec( "guider", @@ -209,30 +192,15 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ - InputParam("attention_kwargs"), - InputParam( - "latents", - required=True, - type_hint=torch.Tensor, - description="The latents to use for the denoising process. Can be generated in prepare_latents step.", - ), - InputParam( - "num_inference_steps", - required=True, - type_hint=int, - description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.", - ), - InputParam( - kwargs_type="denoiser_input_fields", - description="conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.", - ), + InputParam.template("attention_kwargs"), + InputParam.template("denoiser_input_fields"), InputParam( "img_shapes", required=True, - type_hint=List[Tuple[int, int]], - description="The shape of the image latents for RoPE calculation. Can be generated in prepare_additional_inputs step.", + type_hint=list[tuple[int, int]], + description="The shape of the image latents for RoPE calculation. can be generated in prepare_additional_inputs step.", ), ] @@ -247,12 +215,15 @@ def __call__(self, components: QwenImageModularPipeline, block_state: BlockState getattr(block_state, "prompt_embeds_mask", None), getattr(block_state, "negative_prompt_embeds_mask", None), ), - "txt_seq_lens": ( - getattr(block_state, "txt_seq_lens", None), - getattr(block_state, "negative_txt_seq_lens", None), - ), } + transformer_args = set(inspect.signature(components.transformer.forward).parameters.keys()) + additional_cond_kwargs = {} + for field_name, field_value in block_state.denoiser_input_fields.items(): + if field_name in transformer_args and field_name not in guider_inputs: + additional_cond_kwargs[field_name] = field_value + block_state.additional_cond_kwargs.update(additional_cond_kwargs) + components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) guider_state = components.guider.prepare_inputs(guider_inputs) @@ -264,7 +235,6 @@ def __call__(self, components: QwenImageModularPipeline, block_state: BlockState guider_state_batch.noise_pred = components.transformer( hidden_states=block_state.latent_model_input, timestep=block_state.timestep / 1000, - img_shapes=block_state.img_shapes, attention_kwargs=block_state.attention_kwargs, return_dict=False, **cond_kwargs, @@ -284,7 +254,7 @@ def __call__(self, components: QwenImageModularPipeline, block_state: BlockState class QwenImageEditLoopDenoiser(ModularPipelineBlocks): - model_name = "qwenimage" + model_name = "qwenimage-edit" @property def description(self) -> str: @@ -295,7 +265,7 @@ def description(self) -> str: ) @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec( "guider", @@ -307,29 +277,14 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ - InputParam("attention_kwargs"), - InputParam( - "latents", - required=True, - type_hint=torch.Tensor, - description="The latents to use for the denoising process. Can be generated in prepare_latents step.", - ), - InputParam( - "num_inference_steps", - required=True, - type_hint=int, - description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.", - ), - InputParam( - kwargs_type="denoiser_input_fields", - description="conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.", - ), + InputParam.template("attention_kwargs"), + InputParam.template("denoiser_input_fields"), InputParam( "img_shapes", required=True, - type_hint=List[Tuple[int, int]], + type_hint=list[tuple[int, int]], description="The shape of the image latents for RoPE calculation. Can be generated in prepare_additional_inputs step.", ), ] @@ -345,12 +300,15 @@ def __call__(self, components: QwenImageModularPipeline, block_state: BlockState getattr(block_state, "prompt_embeds_mask", None), getattr(block_state, "negative_prompt_embeds_mask", None), ), - "txt_seq_lens": ( - getattr(block_state, "txt_seq_lens", None), - getattr(block_state, "negative_txt_seq_lens", None), - ), } + transformer_args = set(inspect.signature(components.transformer.forward).parameters.keys()) + additional_cond_kwargs = {} + for field_name, field_value in block_state.denoiser_input_fields.items(): + if field_name in transformer_args and field_name not in guider_inputs: + additional_cond_kwargs[field_name] = field_value + block_state.additional_cond_kwargs.update(additional_cond_kwargs) + components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) guider_state = components.guider.prepare_inputs(guider_inputs) @@ -362,7 +320,6 @@ def __call__(self, components: QwenImageModularPipeline, block_state: BlockState guider_state_batch.noise_pred = components.transformer( hidden_states=block_state.latent_model_input, timestep=block_state.timestep / 1000, - img_shapes=block_state.img_shapes, attention_kwargs=block_state.attention_kwargs, return_dict=False, **cond_kwargs, @@ -384,6 +341,7 @@ def __call__(self, components: QwenImageModularPipeline, block_state: BlockState return components, block_state +# loop step:after denoiser class QwenImageLoopAfterDenoiser(ModularPipelineBlocks): model_name = "qwenimage" @@ -396,15 +354,15 @@ def description(self) -> str: ) @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ - OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents."), + OutputParam.template("latents"), ] @torch.no_grad() @@ -437,7 +395,7 @@ def description(self) -> str: ) @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam( "mask", @@ -445,24 +403,19 @@ def inputs(self) -> List[InputParam]: type_hint=torch.Tensor, description="The mask to use for the inpainting process. Can be generated in inpaint prepare latents step.", ), - InputParam( - "image_latents", - required=True, - type_hint=torch.Tensor, - description="The image latents to use for the inpainting process. Can be generated in inpaint prepare latents step.", - ), + InputParam.template("image_latents"), InputParam( "initial_noise", required=True, type_hint=torch.Tensor, description="The initial noise to use for the inpainting process. Can be generated in inpaint prepare latents step.", ), - InputParam( - "timesteps", - required=True, - type_hint=torch.Tensor, - description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.", - ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("latents"), ] @torch.no_grad() @@ -481,6 +434,9 @@ def __call__(self, components: QwenImageModularPipeline, block_state: BlockState return components, block_state +# ==================== +# 2. DENOISE LOOP WRAPPER: define the denoising loop logic +# ==================== class QwenImageDenoiseLoopWrapper(LoopSequentialPipelineBlocks): model_name = "qwenimage" @@ -492,26 +448,21 @@ def description(self) -> str: ) @property - def loop_expected_components(self) -> List[ComponentSpec]: + def loop_expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), ] @property - def loop_inputs(self) -> List[InputParam]: + def loop_inputs(self) -> list[InputParam]: return [ InputParam( - "timesteps", + name="timesteps", required=True, type_hint=torch.Tensor, description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.", ), - InputParam( - "num_inference_steps", - required=True, - type_hint=int, - description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.", - ), + InputParam.template("num_inference_steps", required=True), ] @torch.no_grad() @@ -537,8 +488,50 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) - return components, state -# composing the denoising loops +# ==================== +# 3. DENOISE STEPS: compose the denoising loop with loop wrapper + loop steps +# ==================== + + +# Qwen Image (text2image, image2image) + + +# auto_docstring class QwenImageDenoiseStep(QwenImageDenoiseLoopWrapper): + """ + Denoise step that iteratively denoise the latents. + Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method At each iteration, it runs blocks + defined in `sub_blocks` sequencially: + - `QwenImageLoopBeforeDenoiser` + - `QwenImageLoopDenoiser` + - `QwenImageLoopAfterDenoiser` + This block supports text2image and image2image tasks for QwenImage. + + Components: + guider (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`) scheduler + (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + timesteps (`Tensor`): + The timesteps to use for the denoising process. Can be generated in set_timesteps step. + num_inference_steps (`int`): + The number of denoising steps. + latents (`Tensor`): + The initial latents to use for the denoising process. Can be generated in prepare_latent step. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + img_shapes (`list`): + The shape of the image latents for RoPE calculation. can be generated in prepare_additional_inputs step. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage" + block_classes = [ QwenImageLoopBeforeDenoiser, QwenImageLoopDenoiser, @@ -549,8 +542,8 @@ class QwenImageDenoiseStep(QwenImageDenoiseLoopWrapper): @property def description(self) -> str: return ( - "Denoise step that iteratively denoise the latents. \n" - "Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method \n" + "Denoise step that iteratively denoise the latents.\n" + "Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method\n" "At each iteration, it runs blocks defined in `sub_blocks` sequencially:\n" " - `QwenImageLoopBeforeDenoiser`\n" " - `QwenImageLoopDenoiser`\n" @@ -559,8 +552,49 @@ def description(self) -> str: ) -# composing the inpainting denoising loops +# Qwen Image (inpainting) +# auto_docstring class QwenImageInpaintDenoiseStep(QwenImageDenoiseLoopWrapper): + """ + Denoise step that iteratively denoise the latents. + Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method At each iteration, it runs blocks + defined in `sub_blocks` sequencially: + - `QwenImageLoopBeforeDenoiser` + - `QwenImageLoopDenoiser` + - `QwenImageLoopAfterDenoiser` + - `QwenImageLoopAfterDenoiserInpaint` + This block supports inpainting tasks for QwenImage. + + Components: + guider (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`) scheduler + (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + timesteps (`Tensor`): + The timesteps to use for the denoising process. Can be generated in set_timesteps step. + num_inference_steps (`int`): + The number of denoising steps. + latents (`Tensor`): + The initial latents to use for the denoising process. Can be generated in prepare_latent step. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + img_shapes (`list`): + The shape of the image latents for RoPE calculation. can be generated in prepare_additional_inputs step. + mask (`Tensor`): + The mask to use for the inpainting process. Can be generated in inpaint prepare latents step. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + initial_noise (`Tensor`): + The initial noise to use for the inpainting process. Can be generated in inpaint prepare latents step. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage" block_classes = [ QwenImageLoopBeforeDenoiser, QwenImageLoopDenoiser, @@ -583,8 +617,49 @@ def description(self) -> str: ) -# composing the controlnet denoising loops +# Qwen Image (text2image, image2image) with controlnet +# auto_docstring class QwenImageControlNetDenoiseStep(QwenImageDenoiseLoopWrapper): + """ + Denoise step that iteratively denoise the latents. + Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method At each iteration, it runs blocks + defined in `sub_blocks` sequencially: + - `QwenImageLoopBeforeDenoiser` + - `QwenImageLoopBeforeDenoiserControlNet` + - `QwenImageLoopDenoiser` + - `QwenImageLoopAfterDenoiser` + This block supports text2img/img2img tasks with controlnet for QwenImage. + + Components: + guider (`ClassifierFreeGuidance`) controlnet (`QwenImageControlNetModel`) transformer + (`QwenImageTransformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + timesteps (`Tensor`): + The timesteps to use for the denoising process. Can be generated in set_timesteps step. + num_inference_steps (`int`): + The number of denoising steps. + latents (`Tensor`): + The initial latents to use for the denoising process. Can be generated in prepare_latent step. + control_image_latents (`Tensor`): + The control image to use for the denoising process. Can be generated in prepare_controlnet_inputs step. + controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0): + Scale for ControlNet conditioning. (updated in prepare_controlnet_inputs step.) + controlnet_keep (`list`): + The controlnet keep values. Can be generated in prepare_controlnet_inputs step. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + img_shapes (`list`): + The shape of the image latents for RoPE calculation. can be generated in prepare_additional_inputs step. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage" block_classes = [ QwenImageLoopBeforeDenoiser, QwenImageLoopBeforeDenoiserControlNet, @@ -607,8 +682,56 @@ def description(self) -> str: ) -# composing the controlnet denoising loops +# Qwen Image (inpainting) with controlnet +# auto_docstring class QwenImageInpaintControlNetDenoiseStep(QwenImageDenoiseLoopWrapper): + """ + Denoise step that iteratively denoise the latents. + Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method At each iteration, it runs blocks + defined in `sub_blocks` sequencially: + - `QwenImageLoopBeforeDenoiser` + - `QwenImageLoopBeforeDenoiserControlNet` + - `QwenImageLoopDenoiser` + - `QwenImageLoopAfterDenoiser` + - `QwenImageLoopAfterDenoiserInpaint` + This block supports inpainting tasks with controlnet for QwenImage. + + Components: + guider (`ClassifierFreeGuidance`) controlnet (`QwenImageControlNetModel`) transformer + (`QwenImageTransformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + timesteps (`Tensor`): + The timesteps to use for the denoising process. Can be generated in set_timesteps step. + num_inference_steps (`int`): + The number of denoising steps. + latents (`Tensor`): + The initial latents to use for the denoising process. Can be generated in prepare_latent step. + control_image_latents (`Tensor`): + The control image to use for the denoising process. Can be generated in prepare_controlnet_inputs step. + controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0): + Scale for ControlNet conditioning. (updated in prepare_controlnet_inputs step.) + controlnet_keep (`list`): + The controlnet keep values. Can be generated in prepare_controlnet_inputs step. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + img_shapes (`list`): + The shape of the image latents for RoPE calculation. can be generated in prepare_additional_inputs step. + mask (`Tensor`): + The mask to use for the inpainting process. Can be generated in inpaint prepare latents step. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + initial_noise (`Tensor`): + The initial noise to use for the inpainting process. Can be generated in inpaint prepare latents step. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage" block_classes = [ QwenImageLoopBeforeDenoiser, QwenImageLoopBeforeDenoiserControlNet, @@ -639,8 +762,44 @@ def description(self) -> str: ) -# composing the denoising loops +# Qwen Image Edit (image2image) +# auto_docstring class QwenImageEditDenoiseStep(QwenImageDenoiseLoopWrapper): + """ + Denoise step that iteratively denoise the latents. + Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method At each iteration, it runs blocks + defined in `sub_blocks` sequencially: + - `QwenImageEditLoopBeforeDenoiser` + - `QwenImageEditLoopDenoiser` + - `QwenImageLoopAfterDenoiser` + This block supports QwenImage Edit. + + Components: + guider (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`) scheduler + (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + timesteps (`Tensor`): + The timesteps to use for the denoising process. Can be generated in set_timesteps step. + num_inference_steps (`int`): + The number of denoising steps. + latents (`Tensor`): + The initial latents to use for the denoising process. Can be generated in prepare_latent step. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + img_shapes (`list`): + The shape of the image latents for RoPE calculation. Can be generated in prepare_additional_inputs step. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage-edit" block_classes = [ QwenImageEditLoopBeforeDenoiser, QwenImageEditLoopDenoiser, @@ -661,7 +820,49 @@ def description(self) -> str: ) +# Qwen Image Edit (inpainting) +# auto_docstring class QwenImageEditInpaintDenoiseStep(QwenImageDenoiseLoopWrapper): + """ + Denoise step that iteratively denoise the latents. + Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method At each iteration, it runs blocks + defined in `sub_blocks` sequencially: + - `QwenImageEditLoopBeforeDenoiser` + - `QwenImageEditLoopDenoiser` + - `QwenImageLoopAfterDenoiser` + - `QwenImageLoopAfterDenoiserInpaint` + This block supports inpainting tasks for QwenImage Edit. + + Components: + guider (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`) scheduler + (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + timesteps (`Tensor`): + The timesteps to use for the denoising process. Can be generated in set_timesteps step. + num_inference_steps (`int`): + The number of denoising steps. + latents (`Tensor`): + The initial latents to use for the denoising process. Can be generated in prepare_latent step. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + img_shapes (`list`): + The shape of the image latents for RoPE calculation. Can be generated in prepare_additional_inputs step. + mask (`Tensor`): + The mask to use for the inpainting process. Can be generated in inpaint prepare latents step. + initial_noise (`Tensor`): + The initial noise to use for the inpainting process. Can be generated in inpaint prepare latents step. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage-edit" block_classes = [ QwenImageEditLoopBeforeDenoiser, QwenImageEditLoopDenoiser, @@ -682,3 +883,61 @@ def description(self) -> str: " - `QwenImageLoopAfterDenoiserInpaint`\n" "This block supports inpainting tasks for QwenImage Edit." ) + + +# Qwen Image Layered (image2image) +# auto_docstring +class QwenImageLayeredDenoiseStep(QwenImageDenoiseLoopWrapper): + """ + Denoise step that iteratively denoise the latents. + Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method At each iteration, it runs blocks + defined in `sub_blocks` sequencially: + - `QwenImageEditLoopBeforeDenoiser` + - `QwenImageEditLoopDenoiser` + - `QwenImageLoopAfterDenoiser` + This block supports QwenImage Layered. + + Components: + guider (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`) scheduler + (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + timesteps (`Tensor`): + The timesteps to use for the denoising process. Can be generated in set_timesteps step. + num_inference_steps (`int`): + The number of denoising steps. + latents (`Tensor`): + The initial latents to use for the denoising process. Can be generated in prepare_latent step. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + img_shapes (`list`): + The shape of the image latents for RoPE calculation. Can be generated in prepare_additional_inputs step. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage-layered" + block_classes = [ + QwenImageEditLoopBeforeDenoiser, + QwenImageEditLoopDenoiser, + QwenImageLoopAfterDenoiser, + ] + block_names = ["before_denoiser", "denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoise step that iteratively denoise the latents. \n" + "Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method \n" + "At each iteration, it runs blocks defined in `sub_blocks` sequencially:\n" + " - `QwenImageEditLoopBeforeDenoiser`\n" + " - `QwenImageEditLoopDenoiser`\n" + " - `QwenImageLoopAfterDenoiser`\n" + "This block supports QwenImage Layered." + ) diff --git a/src/diffusers/modular_pipelines/qwenimage/encoders.py b/src/diffusers/modular_pipelines/qwenimage/encoders.py index b126a368bfdf..527267dc0d6e 100644 --- a/src/diffusers/modular_pipelines/qwenimage/encoders.py +++ b/src/diffusers/modular_pipelines/qwenimage/encoders.py @@ -12,7 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Dict, List, Optional, Union +""" +Text and VAE encoder blocks for QwenImage pipelines. +""" import PIL import torch @@ -26,8 +28,19 @@ from ...utils import logging from ...utils.torch_utils import unwrap_module from ..modular_pipeline import ModularPipelineBlocks, PipelineState -from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .modular_pipeline import QwenImageModularPipeline +from .prompt_templates import ( + QWENIMAGE_EDIT_PLUS_IMG_TEMPLATE, + QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE, + QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE_START_IDX, + QWENIMAGE_EDIT_PROMPT_TEMPLATE, + QWENIMAGE_EDIT_PROMPT_TEMPLATE_START_IDX, + QWENIMAGE_LAYERED_CAPTION_PROMPT_CN, + QWENIMAGE_LAYERED_CAPTION_PROMPT_EN, + QWENIMAGE_PROMPT_TEMPLATE, + QWENIMAGE_PROMPT_TEMPLATE_START_IDX, +) logger = logging.get_logger(__name__) @@ -44,11 +57,11 @@ def _extract_masked_hidden(hidden_states: torch.Tensor, mask: torch.Tensor): def get_qwen_prompt_embeds( text_encoder, tokenizer, - prompt: Union[str, List[str]] = None, - prompt_template_encode: str = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n", - prompt_template_encode_start_idx: int = 34, + prompt: str | list[str] = None, + prompt_template_encode: str = QWENIMAGE_PROMPT_TEMPLATE, + prompt_template_encode_start_idx: int = QWENIMAGE_PROMPT_TEMPLATE_START_IDX, tokenizer_max_length: int = 1024, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): prompt = [prompt] if isinstance(prompt, str) else prompt @@ -84,11 +97,11 @@ def get_qwen_prompt_embeds( def get_qwen_prompt_embeds_edit( text_encoder, processor, - prompt: Union[str, List[str]] = None, - image: Optional[torch.Tensor] = None, - prompt_template_encode: str = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n", - prompt_template_encode_start_idx: int = 64, - device: Optional[torch.device] = None, + prompt: str | list[str] = None, + image: torch.Tensor | None = None, + prompt_template_encode: str = QWENIMAGE_EDIT_PROMPT_TEMPLATE, + prompt_template_encode_start_idx: int = QWENIMAGE_EDIT_PROMPT_TEMPLATE_START_IDX, + device: torch.device | None = None, ): prompt = [prompt] if isinstance(prompt, str) else prompt @@ -131,12 +144,12 @@ def get_qwen_prompt_embeds_edit( def get_qwen_prompt_embeds_edit_plus( text_encoder, processor, - prompt: Union[str, List[str]] = None, - image: Optional[Union[torch.Tensor, List[PIL.Image.Image], PIL.Image.Image]] = None, - prompt_template_encode: str = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n", - img_template_encode: str = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>", - prompt_template_encode_start_idx: int = 64, - device: Optional[torch.device] = None, + prompt: str | list[str] = None, + image: torch.Tensor | list[PIL.Image.Image, PIL.Image.Image] | None = None, + prompt_template_encode: str = QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE, + img_template_encode: str = QWENIMAGE_EDIT_PLUS_IMG_TEMPLATE, + prompt_template_encode_start_idx: int = QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE_START_IDX, + device: torch.device | None = None, ): prompt = [prompt] if isinstance(prompt, str) else prompt if isinstance(image, list): @@ -185,7 +198,7 @@ def get_qwen_prompt_embeds_edit_plus( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -241,36 +254,129 @@ def encode_vae_image( return image_latents -class QwenImageEditResizeDynamicStep(ModularPipelineBlocks): - model_name = "qwenimage" +# ==================== +# 1. RESIZE +# ==================== +# In QwenImage pipelines, resize is a separate step because the resized image is used in VL encoding and vae encoder blocks: +# +# image (PIL.Image.Image) +# │ +# ▼ +# resized_image ([PIL.Image.Image]) +# │ +# ├──► text_encoder ──► prompt_embeds, prompt_embeds_mask +# │ (VL encoding needs the resized image for vision-language fusion) +# │ +# └──► image_processor ──► processed_image (torch.Tensor, pixel space) +# │ +# ▼ +# vae_encoder ──► image_latents (torch.Tensor, latent space) +# +# In most of our other pipelines, resizing is done as part of the image preprocessing step. +# ==================== - def __init__(self, input_name: str = "image", output_name: str = "resized_image"): - """Create a configurable step for resizing images to the target area (1024 * 1024) while maintaining the aspect ratio. - This block resizes an input image tensor and exposes the resized result under configurable input and output - names. Use this when you need to wire the resize step to different image fields (e.g., "image", - "control_image") +# auto_docstring +class QwenImageEditResizeStep(ModularPipelineBlocks): + """ + Image Resize step that resize the image to target area while maintaining the aspect ratio. - Args: - input_name (str, optional): Name of the image field to read from the - pipeline state. Defaults to "image". - output_name (str, optional): Name of the resized image field to write - back to the pipeline state. Defaults to "resized_image". - """ - if not isinstance(input_name, str) or not isinstance(output_name, str): - raise ValueError( - f"input_name and output_name must be strings but are {type(input_name)} and {type(output_name)}" - ) - self._image_input_name = input_name - self._resized_image_output_name = output_name - super().__init__() + Components: + image_resize_processor (`VaeImageProcessor`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + + Outputs: + resized_image (`list`): + The resized images + """ + + model_name = "qwenimage-edit" + + @property + def description(self) -> str: + return "Image Resize step that resize the image to target area while maintaining the aspect ratio." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "image_resize_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [InputParam.template("image")] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="resized_image", + type_hint=list[PIL.Image.Image], + description="The resized images", + ), + ] + + @torch.no_grad() + def __call__(self, components: QwenImageModularPipeline, state: PipelineState): + block_state = self.get_block_state(state) + + images = block_state.image + + if not is_valid_image_imagelist(images): + raise ValueError(f"Images must be image or list of images but are {type(images)}") + + if is_valid_image(images): + images = [images] + + image_width, image_height = images[0].size + calculated_width, calculated_height, _ = calculate_dimensions(1024 * 1024, image_width / image_height) + + resized_images = [ + components.image_resize_processor.resize(image, height=calculated_height, width=calculated_width) + for image in images + ] + + block_state.resized_image = resized_images + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class QwenImageLayeredResizeStep(ModularPipelineBlocks): + """ + Image Resize step that resize the image to a target area (defined by the resolution parameter from user) while + maintaining the aspect ratio. + + Components: + image_resize_processor (`VaeImageProcessor`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + resolution (`int`, *optional*, defaults to 640): + The target area to resize the image to, can be 1024 or 640 + + Outputs: + resized_image (`list`): + The resized images + """ + + model_name = "qwenimage-layered" @property def description(self) -> str: - return f"Image Resize step that resize the {self._image_input_name} to the target area (1024 * 1024) while maintaining the aspect ratio." + return "Image Resize step that resize the image to a target area (defined by the resolution parameter from user) while maintaining the aspect ratio." @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec( "image_resize_processor", @@ -281,26 +387,39 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ + InputParam.template("image"), InputParam( - name=self._image_input_name, required=True, type_hint=torch.Tensor, description="The image to resize" + name="resolution", + default=640, + type_hint=int, + description="The target area to resize the image to, can be 1024 or 640", ), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( - name=self._resized_image_output_name, type_hint=List[PIL.Image.Image], description="The resized images" - ), + name="resized_image", + type_hint=list[PIL.Image.Image], + description="The resized images", + ) ] + @staticmethod + def check_inputs(resolution: int): + if resolution not in [1024, 640]: + raise ValueError(f"Resolution must be 1024 or 640 but is {resolution}") + @torch.no_grad() def __call__(self, components: QwenImageModularPipeline, state: PipelineState): block_state = self.get_block_state(state) - images = getattr(block_state, self._image_input_name) + self.check_inputs(resolution=block_state.resolution) + + images = block_state.image if not is_valid_image_imagelist(images): raise ValueError(f"Images must be image or list of images but are {type(images)}") @@ -309,59 +428,79 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState): images = [images] image_width, image_height = images[0].size - calculated_width, calculated_height, _ = calculate_dimensions(1024 * 1024, image_width / image_height) + target_area = block_state.resolution * block_state.resolution + calculated_width, calculated_height, _ = calculate_dimensions(target_area, image_width / image_height) resized_images = [ components.image_resize_processor.resize(image, height=calculated_height, width=calculated_width) for image in images ] - setattr(block_state, self._resized_image_output_name, resized_images) + block_state.resized_image = resized_images self.set_block_state(state, block_state) return components, state -class QwenImageEditPlusResizeDynamicStep(QwenImageEditResizeDynamicStep): - model_name = "qwenimage" +# auto_docstring +class QwenImageEditPlusResizeStep(ModularPipelineBlocks): + """ + Resize images for QwenImage Edit Plus pipeline. + Produces two outputs: resized_image (1024x1024) for VAE encoding, resized_cond_image (384x384) for VL text + encoding. Each image is resized independently based on its own aspect ratio. - def __init__( - self, - input_name: str = "image", - output_name: str = "resized_image", - vae_image_output_name: str = "vae_image", - ): - """Create a configurable step for resizing images to the target area (384 * 384) while maintaining the aspect ratio. + Components: + image_resize_processor (`VaeImageProcessor`) - This block resizes an input image or a list input images and exposes the resized result under configurable - input and output names. Use this when you need to wire the resize step to different image fields (e.g., - "image", "control_image") + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. - Args: - input_name (str, optional): Name of the image field to read from the - pipeline state. Defaults to "image". - output_name (str, optional): Name of the resized image field to write - back to the pipeline state. Defaults to "resized_image". - vae_image_output_name (str, optional): Name of the image field - to write back to the pipeline state. This is used by the VAE encoder step later on. QwenImage Edit Plus - processes the input image(s) differently for the VL and the VAE. - """ - if not isinstance(input_name, str) or not isinstance(output_name, str): - raise ValueError( - f"input_name and output_name must be strings but are {type(input_name)} and {type(output_name)}" - ) - self.condition_image_size = 384 * 384 - self._image_input_name = input_name - self._resized_image_output_name = output_name - self._vae_image_output_name = vae_image_output_name - super().__init__() + Outputs: + resized_image (`list`): + Images resized to 1024x1024 target area for VAE encoding + resized_cond_image (`list`): + Images resized to 384x384 target area for VL text encoding + """ + + model_name = "qwenimage-edit-plus" + + @property + def description(self) -> str: + return ( + "Resize images for QwenImage Edit Plus pipeline.\n" + "Produces two outputs: resized_image (1024x1024) for VAE encoding, " + "resized_cond_image (384x384) for VL text encoding.\n" + "Each image is resized independently based on its own aspect ratio." + ) @property - def intermediate_outputs(self) -> List[OutputParam]: - return super().intermediate_outputs + [ + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "image_resize_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + # image + return [InputParam.template("image")] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ OutputParam( - name=self._vae_image_output_name, - type_hint=List[PIL.Image.Image], - description="The images to be processed which will be further used by the VAE encoder.", + name="resized_image", + type_hint=list[PIL.Image.Image], + description="Images resized to 1024x1024 target area for VAE encoding", + ), + OutputParam( + name="resized_cond_image", + type_hint=list[PIL.Image.Image], + description="Images resized to 384x384 target area for VL text encoding", ), ] @@ -369,44 +508,197 @@ def intermediate_outputs(self) -> List[OutputParam]: def __call__(self, components: QwenImageModularPipeline, state: PipelineState): block_state = self.get_block_state(state) - images = getattr(block_state, self._image_input_name) + images = block_state.image if not is_valid_image_imagelist(images): raise ValueError(f"Images must be image or list of images but are {type(images)}") - if ( - not isinstance(images, torch.Tensor) - and isinstance(images, PIL.Image.Image) - and not isinstance(images, list) - ): + if is_valid_image(images): images = [images] - # TODO (sayakpaul): revisit this when the inputs are `torch.Tensor`s - condition_images = [] - vae_images = [] - for img in images: - image_width, image_height = img.size - condition_width, condition_height, _ = calculate_dimensions( - self.condition_image_size, image_width / image_height + # Resize each image independently based on its own aspect ratio + resized_images = [] + resized_cond_images = [] + for image in images: + image_width, image_height = image.size + + # For VAE encoder (1024x1024 target area) + vae_width, vae_height, _ = calculate_dimensions(1024 * 1024, image_width / image_height) + resized_images.append(components.image_resize_processor.resize(image, height=vae_height, width=vae_width)) + + # For VL text encoder (384x384 target area) + vl_width, vl_height, _ = calculate_dimensions(384 * 384, image_width / image_height) + resized_cond_images.append( + components.image_resize_processor.resize(image, height=vl_height, width=vl_width) ) - condition_images.append(components.image_resize_processor.resize(img, condition_height, condition_width)) - vae_images.append(img) - setattr(block_state, self._resized_image_output_name, condition_images) - setattr(block_state, self._vae_image_output_name, vae_images) + block_state.resized_image = resized_images + block_state.resized_cond_image = resized_cond_images + self.set_block_state(state, block_state) + return components, state + + +# ==================== +# 2. GET IMAGE PROMPT +# ==================== + + +# auto_docstring +class QwenImageLayeredGetImagePromptStep(ModularPipelineBlocks): + """ + Auto-caption step that generates a text prompt from the input image if none is provided. + Uses the VL model (text_encoder) to generate a description of the image. If prompt is already provided, this step + passes through unchanged. + + Components: + text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor (`Qwen2VLProcessor`) + + Inputs: + prompt (`str`, *optional*): + The prompt or prompts to guide image generation. + resized_image (`Image`): + The image to generate caption from, should be resized use the resize step + use_en_prompt (`bool`, *optional*, defaults to False): + Whether to use English prompt template + + Outputs: + prompt (`str`): + The prompt or prompts to guide image generation. If not provided, updated using image caption + """ + + model_name = "qwenimage-layered" + + def __init__(self): + self.image_caption_prompt_en = QWENIMAGE_LAYERED_CAPTION_PROMPT_EN + self.image_caption_prompt_cn = QWENIMAGE_LAYERED_CAPTION_PROMPT_CN + super().__init__() + + @property + def description(self) -> str: + return ( + "Auto-caption step that generates a text prompt from the input image if none is provided.\n" + "Uses the VL model (text_encoder) to generate a description of the image.\n" + "If prompt is already provided, this step passes through unchanged." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", Qwen2_5_VLForConditionalGeneration), + ComponentSpec("processor", Qwen2VLProcessor), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template( + "prompt", required=False + ), # it is not required for qwenimage-layered, unlike other pipelines + InputParam( + name="resized_image", + required=True, + type_hint=PIL.Image.Image, + description="The image to generate caption from, should be resized use the resize step", + ), + InputParam( + name="use_en_prompt", + default=False, + type_hint=bool, + description="Whether to use English prompt template", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="prompt", + type_hint=str, + description="The prompt or prompts to guide image generation. If not provided, updated using image caption", + ), + ] + + @torch.no_grad() + def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + + # If prompt is empty or None, generate caption from image + if block_state.prompt is None or block_state.prompt == "" or block_state.prompt == " ": + if block_state.use_en_prompt: + caption_prompt = self.image_caption_prompt_en + else: + caption_prompt = self.image_caption_prompt_cn + + model_inputs = components.processor( + text=caption_prompt, + images=block_state.resized_image, + padding=True, + return_tensors="pt", + ).to(device) + + generated_ids = components.text_encoder.generate(**model_inputs, max_new_tokens=512) + generated_ids_trimmed = [ + out_ids[len(in_ids) :] for in_ids, out_ids in zip(model_inputs.input_ids, generated_ids) + ] + output_text = components.processor.batch_decode( + generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False + )[0] + + block_state.prompt = output_text.strip() + self.set_block_state(state, block_state) return components, state +# ==================== +# 3. TEXT ENCODER +# ==================== + + +# auto_docstring class QwenImageTextEncoderStep(ModularPipelineBlocks): + """ + Text Encoder step that generates text embeddings to guide the image generation. + + Components: + text_encoder (`Qwen2_5_VLForConditionalGeneration`): The text encoder to use tokenizer (`Qwen2Tokenizer`): + The tokenizer to use guider (`ClassifierFreeGuidance`) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + max_sequence_length (`int`, *optional*, defaults to 1024): + Maximum sequence length for prompt encoding. + + Outputs: + prompt_embeds (`Tensor`): + The prompt embeddings. + prompt_embeds_mask (`Tensor`): + The encoder attention mask. + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. + """ + model_name = "qwenimage" + def __init__(self): + self.prompt_template_encode = QWENIMAGE_PROMPT_TEMPLATE + self.prompt_template_encode_start_idx = QWENIMAGE_PROMPT_TEMPLATE_START_IDX + self.tokenizer_max_length = 1024 + super().__init__() + @property def description(self) -> str: - return "Text Encoder step that generate text_embeddings to guide the image generation" + return "Text Encoder step that generates text embeddings to guide the image generation." @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("text_encoder", Qwen2_5_VLForConditionalGeneration, description="The text encoder to use"), ComponentSpec("tokenizer", Qwen2Tokenizer, description="The tokenizer to use"), @@ -419,53 +711,20 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def expected_configs(self) -> List[ConfigSpec]: - return [ - ConfigSpec( - name="prompt_template_encode", - default="<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n", - ), - ConfigSpec(name="prompt_template_encode_start_idx", default=34), - ConfigSpec(name="tokenizer_max_length", default=1024), - ] - - @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ - InputParam(name="prompt", required=True, type_hint=str, description="The prompt to encode"), - InputParam(name="negative_prompt", type_hint=str, description="The negative prompt to encode"), - InputParam( - name="max_sequence_length", type_hint=int, description="The max sequence length to use", default=1024 - ), + InputParam.template("prompt"), + InputParam.template("negative_prompt"), + InputParam.template("max_sequence_length", default=1024), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ - OutputParam( - name="prompt_embeds", - kwargs_type="denoiser_input_fields", - type_hint=torch.Tensor, - description="The prompt embeddings", - ), - OutputParam( - name="prompt_embeds_mask", - kwargs_type="denoiser_input_fields", - type_hint=torch.Tensor, - description="The encoder attention mask", - ), - OutputParam( - name="negative_prompt_embeds", - kwargs_type="denoiser_input_fields", - type_hint=torch.Tensor, - description="The negative prompt embeddings", - ), - OutputParam( - name="negative_prompt_embeds_mask", - kwargs_type="denoiser_input_fields", - type_hint=torch.Tensor, - description="The negative prompt embeddings mask", - ), + OutputParam.template("prompt_embeds"), + OutputParam.template("prompt_embeds_mask"), + OutputParam.template("negative_prompt_embeds"), + OutputParam.template("negative_prompt_embeds_mask"), ] @staticmethod @@ -494,9 +753,9 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState): components.text_encoder, components.tokenizer, prompt=block_state.prompt, - prompt_template_encode=components.config.prompt_template_encode, - prompt_template_encode_start_idx=components.config.prompt_template_encode_start_idx, - tokenizer_max_length=components.config.tokenizer_max_length, + prompt_template_encode=self.prompt_template_encode, + prompt_template_encode_start_idx=self.prompt_template_encode_start_idx, + tokenizer_max_length=self.tokenizer_max_length, device=device, ) @@ -511,9 +770,9 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState): components.text_encoder, components.tokenizer, prompt=negative_prompt, - prompt_template_encode=components.config.prompt_template_encode, - prompt_template_encode_start_idx=components.config.prompt_template_encode_start_idx, - tokenizer_max_length=components.config.tokenizer_max_length, + prompt_template_encode=self.prompt_template_encode, + prompt_template_encode_start_idx=self.prompt_template_encode_start_idx, + tokenizer_max_length=self.tokenizer_max_length, device=device, ) block_state.negative_prompt_embeds = block_state.negative_prompt_embeds[ @@ -527,15 +786,48 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState): return components, state +# auto_docstring class QwenImageEditTextEncoderStep(ModularPipelineBlocks): + """ + Text Encoder step that processes both prompt and image together to generate text embeddings for guiding image + generation. + + Components: + text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor (`Qwen2VLProcessor`) guider + (`ClassifierFreeGuidance`) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + resized_image (`Image`): + The image prompt to encode, should be resized using resize step + + Outputs: + prompt_embeds (`Tensor`): + The prompt embeddings. + prompt_embeds_mask (`Tensor`): + The encoder attention mask. + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. + """ + model_name = "qwenimage" + def __init__(self): + self.prompt_template_encode = QWENIMAGE_EDIT_PROMPT_TEMPLATE + self.prompt_template_encode_start_idx = QWENIMAGE_EDIT_PROMPT_TEMPLATE_START_IDX + super().__init__() + @property def description(self) -> str: - return "Text Encoder step that processes both prompt and image together to generate text embeddings for guiding image generation" + return "Text Encoder step that processes both prompt and image together to generate text embeddings for guiding image generation." @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("text_encoder", Qwen2_5_VLForConditionalGeneration), ComponentSpec("processor", Qwen2VLProcessor), @@ -548,55 +840,25 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def expected_configs(self) -> List[ConfigSpec]: + def inputs(self) -> list[InputParam]: return [ - ConfigSpec( - name="prompt_template_encode", - default="<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n", - ), - ConfigSpec(name="prompt_template_encode_start_idx", default=64), - ] - - @property - def inputs(self) -> List[InputParam]: - return [ - InputParam(name="prompt", required=True, type_hint=str, description="The prompt to encode"), - InputParam(name="negative_prompt", type_hint=str, description="The negative prompt to encode"), + InputParam.template("prompt"), + InputParam.template("negative_prompt"), InputParam( name="resized_image", required=True, - type_hint=torch.Tensor, + type_hint=PIL.Image.Image, description="The image prompt to encode, should be resized using resize step", ), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ - OutputParam( - name="prompt_embeds", - kwargs_type="denoiser_input_fields", - type_hint=torch.Tensor, - description="The prompt embeddings", - ), - OutputParam( - name="prompt_embeds_mask", - kwargs_type="denoiser_input_fields", - type_hint=torch.Tensor, - description="The encoder attention mask", - ), - OutputParam( - name="negative_prompt_embeds", - kwargs_type="denoiser_input_fields", - type_hint=torch.Tensor, - description="The negative prompt embeddings", - ), - OutputParam( - name="negative_prompt_embeds_mask", - kwargs_type="denoiser_input_fields", - type_hint=torch.Tensor, - description="The negative prompt embeddings mask", - ), + OutputParam.template("prompt_embeds"), + OutputParam.template("prompt_embeds_mask"), + OutputParam.template("negative_prompt_embeds"), + OutputParam.template("negative_prompt_embeds_mask"), ] @staticmethod @@ -624,8 +886,8 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState): components.processor, prompt=block_state.prompt, image=block_state.resized_image, - prompt_template_encode=components.config.prompt_template_encode, - prompt_template_encode_start_idx=components.config.prompt_template_encode_start_idx, + prompt_template_encode=self.prompt_template_encode, + prompt_template_encode_start_idx=self.prompt_template_encode_start_idx, device=device, ) @@ -638,8 +900,8 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState): components.processor, prompt=negative_prompt, image=block_state.resized_image, - prompt_template_encode=components.config.prompt_template_encode, - prompt_template_encode_start_idx=components.config.prompt_template_encode_start_idx, + prompt_template_encode=self.prompt_template_encode, + prompt_template_encode_start_idx=self.prompt_template_encode_start_idx, device=device, ) @@ -647,23 +909,98 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState): return components, state -class QwenImageEditPlusTextEncoderStep(QwenImageEditTextEncoderStep): - model_name = "qwenimage" +# auto_docstring +class QwenImageEditPlusTextEncoderStep(ModularPipelineBlocks): + """ + Text Encoder step for QwenImage Edit Plus that processes prompt and multiple images together to generate text + embeddings for guiding image generation. + + Components: + text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor (`Qwen2VLProcessor`) guider + (`ClassifierFreeGuidance`) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + resized_cond_image (`Tensor`): + The image(s) to encode, can be a single image or list of images, should be resized to 384x384 using + resize step + + Outputs: + prompt_embeds (`Tensor`): + The prompt embeddings. + prompt_embeds_mask (`Tensor`): + The encoder attention mask. + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. + """ + + model_name = "qwenimage-edit-plus" + + def __init__(self): + self.prompt_template_encode = QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE + self.img_template_encode = QWENIMAGE_EDIT_PLUS_IMG_TEMPLATE + self.prompt_template_encode_start_idx = QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE_START_IDX + super().__init__() + + @property + def description(self) -> str: + return ( + "Text Encoder step for QwenImage Edit Plus that processes prompt and multiple images together " + "to generate text embeddings for guiding image generation." + ) @property - def expected_configs(self) -> List[ConfigSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ - ConfigSpec( - name="prompt_template_encode", - default="<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n", + ComponentSpec("text_encoder", Qwen2_5_VLForConditionalGeneration), + ComponentSpec("processor", Qwen2VLProcessor), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 4.0}), + default_creation_method="from_config", ), - ConfigSpec( - name="img_template_encode", - default="Picture {}: <|vision_start|><|image_pad|><|vision_end|>", + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("prompt"), + InputParam.template("negative_prompt"), + InputParam( + name="resized_cond_image", + required=True, + type_hint=torch.Tensor, + description="The image(s) to encode, can be a single image or list of images, should be resized to 384x384 using resize step", ), - ConfigSpec(name="prompt_template_encode_start_idx", default=64), ] + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("prompt_embeds"), + OutputParam.template("prompt_embeds_mask"), + OutputParam.template("negative_prompt_embeds"), + OutputParam.template("negative_prompt_embeds_mask"), + ] + + @staticmethod + def check_inputs(prompt, negative_prompt): + if not isinstance(prompt, str) and not isinstance(prompt, list): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if ( + negative_prompt is not None + and not isinstance(negative_prompt, str) + and not isinstance(negative_prompt, list) + ): + raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}") + @torch.no_grad() def __call__(self, components: QwenImageModularPipeline, state: PipelineState): block_state = self.get_block_state(state) @@ -676,10 +1013,10 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState): components.text_encoder, components.processor, prompt=block_state.prompt, - image=block_state.resized_image, - prompt_template_encode=components.config.prompt_template_encode, - img_template_encode=components.config.img_template_encode, - prompt_template_encode_start_idx=components.config.prompt_template_encode_start_idx, + image=block_state.resized_cond_image, + prompt_template_encode=self.prompt_template_encode, + img_template_encode=self.img_template_encode, + prompt_template_encode_start_idx=self.prompt_template_encode_start_idx, device=device, ) @@ -692,10 +1029,10 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState): components.text_encoder, components.processor, prompt=negative_prompt, - image=block_state.resized_image, - prompt_template_encode=components.config.prompt_template_encode, - img_template_encode=components.config.img_template_encode, - prompt_template_encode_start_idx=components.config.prompt_template_encode_start_idx, + image=block_state.resized_cond_image, + prompt_template_encode=self.prompt_template_encode, + img_template_encode=self.img_template_encode, + prompt_template_encode_start_idx=self.prompt_template_encode_start_idx, device=device, ) ) @@ -704,15 +1041,49 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState): return components, state +# ==================== +# 4. IMAGE PREPROCESS +# ==================== + + +# auto_docstring class QwenImageInpaintProcessImagesInputStep(ModularPipelineBlocks): + """ + Image Preprocess step for inpainting task. This processes the image and mask inputs together. Images will be + resized to the given height and width. + + Components: + image_mask_processor (`InpaintProcessor`) + + Inputs: + mask_image (`Image`): + Mask image for inpainting. + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + padding_mask_crop (`int`, *optional*): + Padding for mask cropping in inpainting. + + Outputs: + processed_image (`Tensor`): + The processed image + processed_mask_image (`Tensor`): + The processed mask image + mask_overlay_kwargs (`dict`): + The kwargs for the postprocess step to apply the mask overlay + """ + model_name = "qwenimage" @property def description(self) -> str: - return "Image Preprocess step for inpainting task. This processes the image and mask inputs together. Images can be resized first using QwenImageEditResizeDynamicStep." + return "Image Preprocess step for inpainting task. This processes the image and mask inputs together. Images will be resized to the given height and width." @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec( "image_mask_processor", @@ -723,24 +1094,31 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ - InputParam("mask_image", required=True), - InputParam("resized_image"), - InputParam("image"), - InputParam("height"), - InputParam("width"), - InputParam("padding_mask_crop"), + InputParam.template("mask_image"), + InputParam.template("image"), + InputParam.template("height"), + InputParam.template("width"), + InputParam.template("padding_mask_crop"), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ - OutputParam(name="processed_image"), - OutputParam(name="processed_mask_image"), + OutputParam( + name="processed_image", + type_hint=torch.Tensor, + description="The processed image", + ), + OutputParam( + name="processed_mask_image", + type_hint=torch.Tensor, + description="The processed mask image", + ), OutputParam( name="mask_overlay_kwargs", - type_hint=Dict, + type_hint=dict, description="The kwargs for the postprocess step to apply the mask overlay", ), ] @@ -757,23 +1135,107 @@ def check_inputs(height, width, vae_scale_factor): def __call__(self, components: QwenImageModularPipeline, state: PipelineState): block_state = self.get_block_state(state) - if block_state.resized_image is None and block_state.image is None: - raise ValueError("resized_image and image cannot be None at the same time") + self.check_inputs( + height=block_state.height, width=block_state.width, vae_scale_factor=components.vae_scale_factor + ) + height = block_state.height or components.default_height + width = block_state.width or components.default_width - if block_state.resized_image is None: - image = block_state.image - self.check_inputs( - height=block_state.height, width=block_state.width, vae_scale_factor=components.vae_scale_factor + block_state.processed_image, block_state.processed_mask_image, block_state.mask_overlay_kwargs = ( + components.image_mask_processor.preprocess( + image=block_state.image, + mask=block_state.mask_image, + height=height, + width=width, + padding_mask_crop=block_state.padding_mask_crop, ) - height = block_state.height or components.default_height - width = block_state.width or components.default_width - else: - width, height = block_state.resized_image[0].size - image = block_state.resized_image + ) + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class QwenImageEditInpaintProcessImagesInputStep(ModularPipelineBlocks): + """ + Image Preprocess step for inpainting task. This processes the image and mask inputs together. Images should be + resized first. + + Components: + image_mask_processor (`InpaintProcessor`) + + Inputs: + mask_image (`Image`): + Mask image for inpainting. + resized_image (`Image`): + The resized image. should be generated using a resize step + padding_mask_crop (`int`, *optional*): + Padding for mask cropping in inpainting. + + Outputs: + processed_image (`Tensor`): + The processed image + processed_mask_image (`Tensor`): + The processed mask image + mask_overlay_kwargs (`dict`): + The kwargs for the postprocess step to apply the mask overlay + """ + + model_name = "qwenimage-edit" + + @property + def description(self) -> str: + return "Image Preprocess step for inpainting task. This processes the image and mask inputs together. Images should be resized first." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "image_mask_processor", + InpaintProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("mask_image"), + InputParam( + name="resized_image", + required=True, + type_hint=PIL.Image.Image, + description="The resized image. should be generated using a resize step", + ), + InputParam.template("padding_mask_crop"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam(name="processed_image", type_hint=torch.Tensor, description="The processed image"), + OutputParam( + name="processed_mask_image", + type_hint=torch.Tensor, + description="The processed mask image", + ), + OutputParam( + name="mask_overlay_kwargs", + type_hint=dict, + description="The kwargs for the postprocess step to apply the mask overlay", + ), + ] + + @torch.no_grad() + def __call__(self, components: QwenImageModularPipeline, state: PipelineState): + block_state = self.get_block_state(state) + + width, height = block_state.resized_image[0].size block_state.processed_image, block_state.processed_mask_image, block_state.mask_overlay_kwargs = ( components.image_mask_processor.preprocess( - image=image, + image=block_state.resized_image, mask=block_state.mask_image, height=height, width=width, @@ -785,15 +1247,35 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState): return components, state +# auto_docstring class QwenImageProcessImagesInputStep(ModularPipelineBlocks): + """ + Image Preprocess step. will resize the image to the given height and width. + + Components: + image_processor (`VaeImageProcessor`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + + Outputs: + processed_image (`Tensor`): + The processed image + """ + model_name = "qwenimage" @property def description(self) -> str: - return "Image Preprocess step. Images can be resized first using QwenImageEditResizeDynamicStep." + return "Image Preprocess step. will resize the image to the given height and width." @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec( "image_processor", @@ -804,12 +1286,22 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def inputs(self) -> List[InputParam]: - return [InputParam("resized_image"), InputParam("image"), InputParam("height"), InputParam("width")] + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("image"), + InputParam.template("height"), + InputParam.template("width"), + ] @property - def intermediate_outputs(self) -> List[OutputParam]: - return [OutputParam(name="processed_image")] + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="processed_image", + type_hint=torch.Tensor, + description="The processed image", + ) + ] @staticmethod def check_inputs(height, width, vae_scale_factor): @@ -823,22 +1315,14 @@ def check_inputs(height, width, vae_scale_factor): def __call__(self, components: QwenImageModularPipeline, state: PipelineState): block_state = self.get_block_state(state) - if block_state.resized_image is None and block_state.image is None: - raise ValueError("resized_image and image cannot be None at the same time") - - if block_state.resized_image is None: - image = block_state.image - self.check_inputs( - height=block_state.height, width=block_state.width, vae_scale_factor=components.vae_scale_factor - ) - height = block_state.height or components.default_height - width = block_state.width or components.default_width - else: - width, height = block_state.resized_image[0].size - image = block_state.resized_image + self.check_inputs( + height=block_state.height, width=block_state.width, vae_scale_factor=components.vae_scale_factor + ) + height = block_state.height or components.default_height + width = block_state.width or components.default_width block_state.processed_image = components.image_processor.preprocess( - image=image, + image=block_state.image, height=height, width=width, ) @@ -847,151 +1331,237 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState): return components, state -class QwenImageEditPlusProcessImagesInputStep(QwenImageProcessImagesInputStep): - model_name = "qwenimage-edit-plus" +# auto_docstring +class QwenImageEditProcessImagesInputStep(ModularPipelineBlocks): + """ + Image Preprocess step. Images needs to be resized first. - def __init__(self): - self.vae_image_size = 1024 * 1024 - super().__init__() + Components: + image_processor (`VaeImageProcessor`) + + Inputs: + resized_image (`list`): + The resized image. should be generated using a resize step + + Outputs: + processed_image (`Tensor`): + The processed image + """ + + model_name = "qwenimage-edit" @property def description(self) -> str: - return "Image Preprocess step for QwenImage Edit Plus. Unlike QwenImage Edit, QwenImage Edit Plus doesn't use the same resized image for further preprocessing." + return "Image Preprocess step. Images needs to be resized first." @property - def inputs(self) -> List[InputParam]: - return [InputParam("vae_image"), InputParam("image"), InputParam("height"), InputParam("width")] + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "image_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="resized_image", + required=True, + type_hint=list[PIL.Image.Image], + description="The resized image. should be generated using a resize step", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="processed_image", + type_hint=torch.Tensor, + description="The processed image", + ) + ] @torch.no_grad() def __call__(self, components: QwenImageModularPipeline, state: PipelineState): block_state = self.get_block_state(state) - if block_state.vae_image is None and block_state.image is None: - raise ValueError("`vae_image` and `image` cannot be None at the same time") - - vae_image_sizes = None - if block_state.vae_image is None: - image = block_state.image - self.check_inputs( - height=block_state.height, width=block_state.width, vae_scale_factor=components.vae_scale_factor - ) - height = block_state.height or components.default_height - width = block_state.width or components.default_width - block_state.processed_image = components.image_processor.preprocess( - image=image, height=height, width=width - ) - else: - # QwenImage Edit Plus can allow multiple input images with varied resolutions - processed_images = [] - vae_image_sizes = [] - for img in block_state.vae_image: - width, height = img.size - vae_width, vae_height, _ = calculate_dimensions(self.vae_image_size, width / height) - vae_image_sizes.append((vae_width, vae_height)) - processed_images.append( - components.image_processor.preprocess(image=img, height=vae_height, width=vae_width) - ) - block_state.processed_image = processed_images + width, height = block_state.resized_image[0].size - block_state.vae_image_sizes = vae_image_sizes + block_state.processed_image = components.image_processor.preprocess( + image=block_state.resized_image, + height=height, + width=width, + ) self.set_block_state(state, block_state) return components, state -class QwenImageVaeEncoderDynamicStep(ModularPipelineBlocks): - model_name = "qwenimage" - - def __init__( - self, - input_name: str = "processed_image", - output_name: str = "image_latents", - ): - """Initialize a VAE encoder step for converting images to latent representations. +# auto_docstring +class QwenImageEditPlusProcessImagesInputStep(ModularPipelineBlocks): + """ + Image Preprocess step. Images can be resized first. If a list of images is provided, will return a list of + processed images. - Both the input and output names are configurable so this block can be configured to process to different image - inputs (e.g., "processed_image" -> "image_latents", "processed_control_image" -> "control_image_latents"). + Components: + image_processor (`VaeImageProcessor`) - Args: - input_name (str, optional): Name of the input image tensor. Defaults to "processed_image". - Examples: "processed_image" or "processed_control_image" - output_name (str, optional): Name of the output latent tensor. Defaults to "image_latents". - Examples: "image_latents" or "control_image_latents" + Inputs: + resized_image (`list`): + The resized image. should be generated using a resize step - Examples: - # Basic usage with default settings (includes image processor) QwenImageVaeEncoderDynamicStep() + Outputs: + processed_image (`Tensor`): + The processed image + """ - # Custom input/output names for control image QwenImageVaeEncoderDynamicStep( - input_name="processed_control_image", output_name="control_image_latents" - ) - """ - self._image_input_name = input_name - self._image_latents_output_name = output_name - super().__init__() + model_name = "qwenimage-edit-plus" @property def description(self) -> str: - return f"Dynamic VAE Encoder step that converts {self._image_input_name} into latent representations {self._image_latents_output_name}.\n" + return "Image Preprocess step. Images can be resized first. If a list of images is provided, will return a list of processed images." @property - def expected_components(self) -> List[ComponentSpec]: - components = [ComponentSpec("vae", AutoencoderKLQwenImage)] - return components + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "image_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] @property - def inputs(self) -> List[InputParam]: - inputs = [InputParam(self._image_input_name, required=True), InputParam("generator")] - return inputs + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="resized_image", + required=True, + type_hint=list[PIL.Image.Image], + description="The resized image. should be generated using a resize step", + ) + ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( - self._image_latents_output_name, + name="processed_image", type_hint=torch.Tensor, - description="The latents representing the reference image", + description="The processed image", ) ] @torch.no_grad() - def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState: + def __call__(self, components: QwenImageModularPipeline, state: PipelineState): block_state = self.get_block_state(state) - device = components._execution_device - dtype = components.vae.dtype + image = block_state.resized_image - image = getattr(block_state, self._image_input_name) + is_image_list = isinstance(image, list) + if not is_image_list: + image = [image] - # Encode image into latents - image_latents = encode_vae_image( - image=image, - vae=components.vae, - generator=block_state.generator, - device=device, - dtype=dtype, - latent_channels=components.num_channels_latents, - ) - setattr(block_state, self._image_latents_output_name, image_latents) + processed_images = [] + for img in image: + img_width, img_height = img.size + processed_images.append( + components.image_processor.preprocess(image=img, height=img_height, width=img_width) + ) - self.set_block_state(state, block_state) + if is_image_list: + block_state.processed_image = processed_images + else: + block_state.processed_image = processed_images[0] + self.set_block_state(state, block_state) return components, state -class QwenImageEditPlusVaeEncoderDynamicStep(QwenImageVaeEncoderDynamicStep): - model_name = "qwenimage-edit-plus" +# ==================== +# 5. VAE ENCODER +# ==================== + + +# auto_docstring +class QwenImageVaeEncoderStep(ModularPipelineBlocks): + """ + VAE Encoder step that converts processed_image into latent representations image_latents. + Handles both single images and lists of images with varied resolutions. + + Components: + vae (`AutoencoderKLQwenImage`) + + Inputs: + processed_image (`Tensor`): + The image tensor to encode + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + image_latents (`Tensor`): + The latent representation of the input image. + """ + + model_name = "qwenimage" + + def __init__(self, input: InputParam | None = None, output: OutputParam | None = None): + """Initialize a VAE encoder step for converting images to latent representations. + + Handles both single images and lists of images. When input is a list, outputs a list of latents. When input is + a single tensor, outputs a single latent tensor. + + Args: + input (InputParam, optional): Input parameter for the processed image. Defaults to "processed_image". + output (OutputParam, optional): Output parameter for the image latents. Defaults to "image_latents". + """ + if input is None: + input = InputParam( + name="processed_image", required=True, type_hint=torch.Tensor, description="The image tensor to encode" + ) + + if output is None: + output = OutputParam.template("image_latents") + + if not isinstance(input, InputParam): + raise ValueError(f"input must be InputParam but is {type(input)}") + if not isinstance(output, OutputParam): + raise ValueError(f"output must be OutputParam but is {type(output)}") + + self._input = input + self._output = output + self._image_input_name = input.name + self._image_latents_output_name = output.name + super().__init__() + + @property + def description(self) -> str: + return ( + f"VAE Encoder step that converts {self._image_input_name} into latent representations {self._image_latents_output_name}.\n" + "Handles both single images and lists of images with varied resolutions." + ) @property - def intermediate_outputs(self) -> List[OutputParam]: - # Each reference image latent can have varied resolutions hence we return this as a list. + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("vae", AutoencoderKLQwenImage)] + + @property + def inputs(self) -> list[InputParam]: return [ - OutputParam( - self._image_latents_output_name, - type_hint=List[torch.Tensor], - description="The latents representing the reference image(s).", - ) + self._input, # default is "processed_image" + InputParam.template("generator"), ] + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [self._output] # default is "image_latents" + @torch.no_grad() def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) @@ -1000,8 +1570,11 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) - dtype = components.vae.dtype image = getattr(block_state, self._image_input_name) + is_image_list = isinstance(image, list) + if not is_image_list: + image = [image] - # Encode image into latents + # Handle both single image and list of images image_latents = [] for img in image: image_latents.append( @@ -1014,6 +1587,8 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) - latent_channels=components.num_channels_latents, ) ) + if not is_image_list: + image_latents = image_latents[0] setattr(block_state, self._image_latents_output_name, image_latents) @@ -1022,7 +1597,30 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) - return components, state +# auto_docstring class QwenImageControlNetVaeEncoderStep(ModularPipelineBlocks): + """ + VAE Encoder step that converts `control_image` into latent representations control_image_latents. + + Components: + vae (`AutoencoderKLQwenImage`) controlnet (`QwenImageControlNetModel`) control_image_processor + (`VaeImageProcessor`) + + Inputs: + control_image (`Image`): + Control image for ControlNet conditioning. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + control_image_latents (`Tensor`): + The latents representing the control image + """ + model_name = "qwenimage" @property @@ -1030,7 +1628,7 @@ def description(self) -> str: return "VAE Encoder step that converts `control_image` into latent representations control_image_latents.\n" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: components = [ ComponentSpec("vae", AutoencoderKLQwenImage), ComponentSpec("controlnet", QwenImageControlNetModel), @@ -1044,17 +1642,17 @@ def expected_components(self) -> List[ComponentSpec]: return components @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: inputs = [ - InputParam("control_image", required=True), - InputParam("height"), - InputParam("width"), - InputParam("generator"), + InputParam.template("control_image"), + InputParam.template("height"), + InputParam.template("width"), + InputParam.template("generator"), ] return inputs @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "control_image_latents", @@ -1131,3 +1729,52 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) - self.set_block_state(state, block_state) return components, state + + +# ==================== +# 6. PERMUTE LATENTS +# ==================== + + +# auto_docstring +class QwenImageLayeredPermuteLatentsStep(ModularPipelineBlocks): + """ + Permute image latents from (B, C, 1, H, W) to (B, 1, C, H, W) for Layered packing. + + Inputs: + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + + Outputs: + image_latents (`Tensor`): + The latent representation of the input image. (permuted from [B, C, 1, H, W] to [B, 1, C, H, W]) + """ + + model_name = "qwenimage-layered" + + @property + def description(self) -> str: + return "Permute image latents from (B, C, 1, H, W) to (B, 1, C, H, W) for Layered packing." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("image_latents"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("image_latents", note="permuted from [B, C, 1, H, W] to [B, 1, C, H, W]"), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + # Permute: (B, C, 1, H, W) -> (B, 1, C, H, W) + latents = block_state.image_latents + block_state.image_latents = latents.permute(0, 2, 1, 3, 4) + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/qwenimage/inputs.py b/src/diffusers/modular_pipelines/qwenimage/inputs.py index 6e656e484847..faec7db245df 100644 --- a/src/diffusers/modular_pipelines/qwenimage/inputs.py +++ b/src/diffusers/modular_pipelines/qwenimage/inputs.py @@ -12,14 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Tuple import torch from ...models import QwenImageMultiControlNetModel from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam -from .modular_pipeline import QwenImageModularPipeline, QwenImagePachifier +from .modular_pipeline import QwenImageLayeredPachifier, QwenImageModularPipeline, QwenImagePachifier def repeat_tensor_to_batch_size( @@ -78,7 +77,7 @@ def repeat_tensor_to_batch_size( return input_tensor -def calculate_dimension_from_latents(latents: torch.Tensor, vae_scale_factor: int) -> Tuple[int, int]: +def calculate_dimension_from_latents(latents: torch.Tensor, vae_scale_factor: int) -> tuple[int, int]: """Calculate image dimensions from latent tensor dimensions. This function converts latent space dimensions to image space dimensions by multiplying the latent height and width @@ -91,7 +90,7 @@ def calculate_dimension_from_latents(latents: torch.Tensor, vae_scale_factor: in Typically 8 for most VAEs (image is 8x larger than latents in each dimension) Returns: - Tuple[int, int]: The calculated image dimensions as (height, width) + tuple[int, int]: The calculated image dimensions as (height, width) Raises: ValueError: If latents tensor doesn't have 4 or 5 dimensions @@ -109,7 +108,44 @@ def calculate_dimension_from_latents(latents: torch.Tensor, vae_scale_factor: in return height, width +# auto_docstring class QwenImageTextInputsStep(ModularPipelineBlocks): + """ + Text input processing step that standardizes text embeddings for the pipeline. + This step: + 1. Determines `batch_size` and `dtype` based on `prompt_embeds` + 2. Ensures all text embeddings have consistent batch sizes (batch_size * num_images_per_prompt) + + This block should be placed after all encoder steps to process the text embeddings before they are used in + subsequent pipeline steps. + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + + Outputs: + batch_size (`int`): + The batch size of the prompt embeddings + dtype (`dtype`): + The data type of the prompt embeddings + prompt_embeds (`Tensor`): + The prompt embeddings. (batch-expanded) + prompt_embeds_mask (`Tensor`): + The encoder attention mask. (batch-expanded) + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. (batch-expanded) + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. (batch-expanded) + """ + model_name = "qwenimage" @property @@ -127,28 +163,24 @@ def description(self) -> str: return summary_section + placement_section @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ - InputParam(name="num_images_per_prompt", default=1), - InputParam(name="prompt_embeds", required=True, kwargs_type="denoiser_input_fields"), - InputParam(name="prompt_embeds_mask", required=True, kwargs_type="denoiser_input_fields"), - InputParam(name="negative_prompt_embeds", kwargs_type="denoiser_input_fields"), - InputParam(name="negative_prompt_embeds_mask", kwargs_type="denoiser_input_fields"), + InputParam.template("num_images_per_prompt"), + InputParam.template("prompt_embeds"), + InputParam.template("prompt_embeds_mask"), + InputParam.template("negative_prompt_embeds"), + InputParam.template("negative_prompt_embeds_mask"), ] @property - def intermediate_outputs(self) -> List[str]: + def intermediate_outputs(self) -> list[OutputParam]: return [ - OutputParam( - "batch_size", - type_hint=int, - description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt", - ), - OutputParam( - "dtype", - type_hint=torch.dtype, - description="Data type of model tensor inputs (determined by `prompt_embeds`)", - ), + OutputParam(name="batch_size", type_hint=int, description="The batch size of the prompt embeddings"), + OutputParam(name="dtype", type_hint=torch.dtype, description="The data type of the prompt embeddings"), + OutputParam.template("prompt_embeds", note="batch-expanded"), + OutputParam.template("prompt_embeds_mask", note="batch-expanded"), + OutputParam.template("negative_prompt_embeds", note="batch-expanded"), + OutputParam.template("negative_prompt_embeds_mask", note="batch-expanded"), ] @staticmethod @@ -221,41 +253,76 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) - return components, state -class QwenImageInputsDynamicStep(ModularPipelineBlocks): - model_name = "qwenimage" - - def __init__(self, image_latent_inputs: List[str] = ["image_latents"], additional_batch_inputs: List[str] = []): - """Initialize a configurable step that standardizes the inputs for the denoising step. It:\n" - - This step handles multiple common tasks to prepare inputs for the denoising step: - 1. For encoded image latents, use it update height/width if None, patchifies, and expands batch size - 2. For additional_batch_inputs: Only expands batch dimensions to match final batch size - - This is a dynamic block that allows you to configure which inputs to process. - - Args: - image_latent_inputs (List[str], optional): Names of image latent tensors to process. - These will be used to determine height/width, patchified, and batch-expanded. Can be a single string or - list of strings. Defaults to ["image_latents"]. Examples: ["image_latents"], ["control_image_latents"] - additional_batch_inputs (List[str], optional): - Names of additional conditional input tensors to expand batch size. These tensors will only have their - batch dimensions adjusted to match the final batch size. Can be a single string or list of strings. - Defaults to []. Examples: ["processed_mask_image"] +# auto_docstring +class QwenImageAdditionalInputsStep(ModularPipelineBlocks): + """ + Input processing step that: + 1. For image latent inputs: Updates height/width if None, patchifies, and expands batch size + 2. For additional batch inputs: Expands batch dimensions to match final batch size + + Configured inputs: + - Image latent inputs: ['image_latents'] + + This block should be placed after the encoder steps and the text input step. + + Components: + pachifier (`QwenImagePachifier`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + batch_size (`int`, *optional*, defaults to 1): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + + Outputs: + image_height (`int`): + The image height calculated from the image latents dimension + image_width (`int`): + The image width calculated from the image latents dimension + height (`int`): + if not provided, updated to image height + width (`int`): + if not provided, updated to image width + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. (patchified and + batch-expanded) + """ - Examples: - # Configure to process image_latents (default behavior) QwenImageInputsDynamicStep() + model_name = "qwenimage" - # Configure to process multiple image latent inputs - QwenImageInputsDynamicStep(image_latent_inputs=["image_latents", "control_image_latents"]) + def __init__( + self, + image_latent_inputs: list[InputParam] | None = None, + additional_batch_inputs: list[InputParam] | None = None, + ): + # by default, process `image_latents` + if image_latent_inputs is None: + image_latent_inputs = [InputParam.template("image_latents")] + if additional_batch_inputs is None: + additional_batch_inputs = [] - # Configure to process image latents and additional batch inputs QwenImageInputsDynamicStep( - image_latent_inputs=["image_latents"], additional_batch_inputs=["processed_mask_image"] - ) - """ if not isinstance(image_latent_inputs, list): - image_latent_inputs = [image_latent_inputs] + raise ValueError(f"image_latent_inputs must be a list, but got {type(image_latent_inputs)}") + else: + for input_param in image_latent_inputs: + if not isinstance(input_param, InputParam): + raise ValueError(f"image_latent_inputs must be a list of InputParam, but got {type(input_param)}") + if not isinstance(additional_batch_inputs, list): - additional_batch_inputs = [additional_batch_inputs] + raise ValueError(f"additional_batch_inputs must be a list, but got {type(additional_batch_inputs)}") + else: + for input_param in additional_batch_inputs: + if not isinstance(input_param, InputParam): + raise ValueError( + f"additional_batch_inputs must be a list of InputParam, but got {type(input_param)}" + ) self._image_latent_inputs = image_latent_inputs self._additional_batch_inputs = additional_batch_inputs @@ -263,69 +330,100 @@ def __init__(self, image_latent_inputs: List[str] = ["image_latents"], additiona @property def description(self) -> str: - # Functionality section summary_section = ( "Input processing step that:\n" - " 1. For image latent inputs: Updates height/width if None, patchifies latents, and expands batch size\n" + " 1. For image latent inputs: Updates height/width if None, patchifies, and expands batch size\n" " 2. For additional batch inputs: Expands batch dimensions to match final batch size" ) - # Inputs info inputs_info = "" if self._image_latent_inputs or self._additional_batch_inputs: inputs_info = "\n\nConfigured inputs:" if self._image_latent_inputs: - inputs_info += f"\n - Image latent inputs: {self._image_latent_inputs}" + inputs_info += f"\n - Image latent inputs: {[p.name for p in self._image_latent_inputs]}" if self._additional_batch_inputs: - inputs_info += f"\n - Additional batch inputs: {self._additional_batch_inputs}" + inputs_info += f"\n - Additional batch inputs: {[p.name for p in self._additional_batch_inputs]}" - # Placement guidance placement_section = "\n\nThis block should be placed after the encoder steps and the text input step." return summary_section + inputs_info + placement_section @property - def inputs(self) -> List[InputParam]: - inputs = [ - InputParam(name="num_images_per_prompt", default=1), - InputParam(name="batch_size", required=True), - InputParam(name="height"), - InputParam(name="width"), + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("pachifier", QwenImagePachifier, default_creation_method="from_config"), ] - # Add image latent inputs - for image_latent_input_name in self._image_latent_inputs: - inputs.append(InputParam(name=image_latent_input_name)) - - # Add additional batch inputs - for input_name in self._additional_batch_inputs: - inputs.append(InputParam(name=input_name)) + @property + def inputs(self) -> list[InputParam]: + inputs = [ + InputParam.template("num_images_per_prompt"), + InputParam.template("batch_size"), + InputParam.template("height"), + InputParam.template("width"), + ] + # default is `image_latents` + inputs += self._image_latent_inputs + self._additional_batch_inputs return inputs @property - def intermediate_outputs(self) -> List[OutputParam]: - return [ - OutputParam(name="image_height", type_hint=int, description="The height of the image latents"), - OutputParam(name="image_width", type_hint=int, description="The width of the image latents"), + def intermediate_outputs(self) -> list[OutputParam]: + outputs = [ + OutputParam( + name="image_height", + type_hint=int, + description="The image height calculated from the image latents dimension", + ), + OutputParam( + name="image_width", + type_hint=int, + description="The image width calculated from the image latents dimension", + ), ] - @property - def expected_components(self) -> List[ComponentSpec]: - return [ - ComponentSpec("pachifier", QwenImagePachifier, default_creation_method="from_config"), - ] + # `height`/`width` are not new outputs, but they will be updated if any image latent inputs are provided + if len(self._image_latent_inputs) > 0: + outputs.append( + OutputParam(name="height", type_hint=int, description="if not provided, updated to image height") + ) + outputs.append( + OutputParam(name="width", type_hint=int, description="if not provided, updated to image width") + ) + + # image latent inputs are modified in place (patchified and batch-expanded) + for input_param in self._image_latent_inputs: + outputs.append( + OutputParam( + name=input_param.name, + type_hint=input_param.type_hint, + description=input_param.description + " (patchified and batch-expanded)", + ) + ) + + # additional batch inputs (batch-expanded only) + for input_param in self._additional_batch_inputs: + outputs.append( + OutputParam( + name=input_param.name, + type_hint=input_param.type_hint, + description=input_param.description + " (batch-expanded)", + ) + ) + + return outputs def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - # Process image latent inputs (height/width calculation, patchify, and batch expansion) - for image_latent_input_name in self._image_latent_inputs: + # Process image latent inputs + for input_param in self._image_latent_inputs: + image_latent_input_name = input_param.name image_latent_tensor = getattr(block_state, image_latent_input_name) if image_latent_tensor is None: continue - # 1. Calculate height/width from latents + # 1. Calculate height/width from latents and update if not provided height, width = calculate_dimension_from_latents(image_latent_tensor, components.vae_scale_factor) block_state.height = block_state.height or height block_state.width = block_state.width or width @@ -335,7 +433,7 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) - if not hasattr(block_state, "image_width"): block_state.image_width = width - # 2. Patchify the image latent tensor + # 2. Patchify image_latent_tensor = components.pachifier.pack_latents(image_latent_tensor) # 3. Expand batch size @@ -349,12 +447,12 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) - setattr(block_state, image_latent_input_name, image_latent_tensor) # Process additional batch inputs (only batch expansion) - for input_name in self._additional_batch_inputs: + for input_param in self._additional_batch_inputs: + input_name = input_param.name input_tensor = getattr(block_state, input_name) if input_tensor is None: continue - # Only expand batch size input_tensor = repeat_tensor_to_batch_size( input_name=input_name, input_tensor=input_tensor, @@ -368,63 +466,436 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) - return components, state -class QwenImageEditPlusInputsDynamicStep(QwenImageInputsDynamicStep): +# auto_docstring +class QwenImageEditPlusAdditionalInputsStep(ModularPipelineBlocks): + """ + Input processing step for Edit Plus that: + 1. For image latent inputs (list): Collects heights/widths, patchifies each, concatenates, expands batch + 2. For additional batch inputs: Expands batch dimensions to match final batch size + Height/width defaults to last image in the list. + + Configured inputs: + - Image latent inputs: ['image_latents'] + + This block should be placed after the encoder steps and the text input step. + + Components: + pachifier (`QwenImagePachifier`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + batch_size (`int`, *optional*, defaults to 1): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + + Outputs: + image_height (`list`): + The image heights calculated from the image latents dimension + image_width (`list`): + The image widths calculated from the image latents dimension + height (`int`): + if not provided, updated to image height + width (`int`): + if not provided, updated to image width + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. (patchified, + concatenated, and batch-expanded) + """ + model_name = "qwenimage-edit-plus" + def __init__( + self, + image_latent_inputs: list[InputParam] | None = None, + additional_batch_inputs: list[InputParam] | None = None, + ): + if image_latent_inputs is None: + image_latent_inputs = [InputParam.template("image_latents")] + if additional_batch_inputs is None: + additional_batch_inputs = [] + + if not isinstance(image_latent_inputs, list): + raise ValueError(f"image_latent_inputs must be a list, but got {type(image_latent_inputs)}") + else: + for input_param in image_latent_inputs: + if not isinstance(input_param, InputParam): + raise ValueError(f"image_latent_inputs must be a list of InputParam, but got {type(input_param)}") + + if not isinstance(additional_batch_inputs, list): + raise ValueError(f"additional_batch_inputs must be a list, but got {type(additional_batch_inputs)}") + else: + for input_param in additional_batch_inputs: + if not isinstance(input_param, InputParam): + raise ValueError( + f"additional_batch_inputs must be a list of InputParam, but got {type(input_param)}" + ) + + self._image_latent_inputs = image_latent_inputs + self._additional_batch_inputs = additional_batch_inputs + super().__init__() + + @property + def description(self) -> str: + summary_section = ( + "Input processing step for Edit Plus that:\n" + " 1. For image latent inputs (list): Collects heights/widths, patchifies each, concatenates, expands batch\n" + " 2. For additional batch inputs: Expands batch dimensions to match final batch size\n" + " Height/width defaults to last image in the list." + ) + + inputs_info = "" + if self._image_latent_inputs or self._additional_batch_inputs: + inputs_info = "\n\nConfigured inputs:" + if self._image_latent_inputs: + inputs_info += f"\n - Image latent inputs: {[p.name for p in self._image_latent_inputs]}" + if self._additional_batch_inputs: + inputs_info += f"\n - Additional batch inputs: {[p.name for p in self._additional_batch_inputs]}" + + placement_section = "\n\nThis block should be placed after the encoder steps and the text input step." + + return summary_section + inputs_info + placement_section + @property - def intermediate_outputs(self) -> List[OutputParam]: + def expected_components(self) -> list[ComponentSpec]: return [ - OutputParam(name="image_height", type_hint=List[int], description="The height of the image latents"), - OutputParam(name="image_width", type_hint=List[int], description="The width of the image latents"), + ComponentSpec("pachifier", QwenImagePachifier, default_creation_method="from_config"), + ] + + @property + def inputs(self) -> list[InputParam]: + inputs = [ + InputParam.template("num_images_per_prompt"), + InputParam.template("batch_size"), + InputParam.template("height"), + InputParam.template("width"), + ] + + # default is `image_latents` + inputs += self._image_latent_inputs + self._additional_batch_inputs + + return inputs + + @property + def intermediate_outputs(self) -> list[OutputParam]: + outputs = [ + OutputParam( + name="image_height", + type_hint=list[int], + description="The image heights calculated from the image latents dimension", + ), + OutputParam( + name="image_width", + type_hint=list[int], + description="The image widths calculated from the image latents dimension", + ), ] + # `height`/`width` are updated if any image latent inputs are provided + if len(self._image_latent_inputs) > 0: + outputs.append( + OutputParam(name="height", type_hint=int, description="if not provided, updated to image height") + ) + outputs.append( + OutputParam(name="width", type_hint=int, description="if not provided, updated to image width") + ) + + # image latent inputs are modified in place (patchified, concatenated, and batch-expanded) + for input_param in self._image_latent_inputs: + outputs.append( + OutputParam( + name=input_param.name, + type_hint=input_param.type_hint, + description=input_param.description + " (patchified, concatenated, and batch-expanded)", + ) + ) + + # additional batch inputs (batch-expanded only) + for input_param in self._additional_batch_inputs: + outputs.append( + OutputParam( + name=input_param.name, + type_hint=input_param.type_hint, + description=input_param.description + " (batch-expanded)", + ) + ) + + return outputs + def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - # Process image latent inputs (height/width calculation, patchify, and batch expansion) - for image_latent_input_name in self._image_latent_inputs: + # Process image latent inputs + for input_param in self._image_latent_inputs: + image_latent_input_name = input_param.name image_latent_tensor = getattr(block_state, image_latent_input_name) if image_latent_tensor is None: continue - # Each image latent can have different size in QwenImage Edit Plus. + is_list = isinstance(image_latent_tensor, list) + if not is_list: + image_latent_tensor = [image_latent_tensor] + image_heights = [] image_widths = [] packed_image_latent_tensors = [] - for img_latent_tensor in image_latent_tensor: + for i, img_latent_tensor in enumerate(image_latent_tensor): # 1. Calculate height/width from latents height, width = calculate_dimension_from_latents(img_latent_tensor, components.vae_scale_factor) image_heights.append(height) image_widths.append(width) - # 2. Patchify the image latent tensor + # 2. Patchify img_latent_tensor = components.pachifier.pack_latents(img_latent_tensor) # 3. Expand batch size img_latent_tensor = repeat_tensor_to_batch_size( - input_name=image_latent_input_name, + input_name=f"{image_latent_input_name}[{i}]", input_tensor=img_latent_tensor, num_images_per_prompt=block_state.num_images_per_prompt, batch_size=block_state.batch_size, ) packed_image_latent_tensors.append(img_latent_tensor) + # Concatenate all packed latents along dim=1 packed_image_latent_tensors = torch.cat(packed_image_latent_tensors, dim=1) + + # Output lists of heights/widths block_state.image_height = image_heights block_state.image_width = image_widths - setattr(block_state, image_latent_input_name, packed_image_latent_tensors) + # Default height/width from last image block_state.height = block_state.height or image_heights[-1] block_state.width = block_state.width or image_widths[-1] + setattr(block_state, image_latent_input_name, packed_image_latent_tensors) + + # Process additional batch inputs (only batch expansion) + for input_param in self._additional_batch_inputs: + input_name = input_param.name + input_tensor = getattr(block_state, input_name) + if input_tensor is None: + continue + + input_tensor = repeat_tensor_to_batch_size( + input_name=input_name, + input_tensor=input_tensor, + num_images_per_prompt=block_state.num_images_per_prompt, + batch_size=block_state.batch_size, + ) + + setattr(block_state, input_name, input_tensor) + + self.set_block_state(state, block_state) + return components, state + + +# same as QwenImageAdditionalInputsStep, but with layered pachifier. + + +# auto_docstring +class QwenImageLayeredAdditionalInputsStep(ModularPipelineBlocks): + """ + Input processing step for Layered that: + 1. For image latent inputs: Updates height/width if None, patchifies with layered pachifier, and expands batch + size + 2. For additional batch inputs: Expands batch dimensions to match final batch size + + Configured inputs: + - Image latent inputs: ['image_latents'] + + This block should be placed after the encoder steps and the text input step. + + Components: + pachifier (`QwenImageLayeredPachifier`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + batch_size (`int`, *optional*, defaults to 1): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + + Outputs: + image_height (`int`): + The image height calculated from the image latents dimension + image_width (`int`): + The image width calculated from the image latents dimension + height (`int`): + if not provided, updated to image height + width (`int`): + if not provided, updated to image width + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. (patchified + with layered pachifier and batch-expanded) + """ + + model_name = "qwenimage-layered" + + def __init__( + self, + image_latent_inputs: list[InputParam] | None = None, + additional_batch_inputs: list[InputParam] | None = None, + ): + if image_latent_inputs is None: + image_latent_inputs = [InputParam.template("image_latents")] + if additional_batch_inputs is None: + additional_batch_inputs = [] + + if not isinstance(image_latent_inputs, list): + raise ValueError(f"image_latent_inputs must be a list, but got {type(image_latent_inputs)}") + else: + for input_param in image_latent_inputs: + if not isinstance(input_param, InputParam): + raise ValueError(f"image_latent_inputs must be a list of InputParam, but got {type(input_param)}") + + if not isinstance(additional_batch_inputs, list): + raise ValueError(f"additional_batch_inputs must be a list, but got {type(additional_batch_inputs)}") + else: + for input_param in additional_batch_inputs: + if not isinstance(input_param, InputParam): + raise ValueError( + f"additional_batch_inputs must be a list of InputParam, but got {type(input_param)}" + ) + + self._image_latent_inputs = image_latent_inputs + self._additional_batch_inputs = additional_batch_inputs + super().__init__() + + @property + def description(self) -> str: + summary_section = ( + "Input processing step for Layered that:\n" + " 1. For image latent inputs: Updates height/width if None, patchifies with layered pachifier, and expands batch size\n" + " 2. For additional batch inputs: Expands batch dimensions to match final batch size" + ) + + inputs_info = "" + if self._image_latent_inputs or self._additional_batch_inputs: + inputs_info = "\n\nConfigured inputs:" + if self._image_latent_inputs: + inputs_info += f"\n - Image latent inputs: {[p.name for p in self._image_latent_inputs]}" + if self._additional_batch_inputs: + inputs_info += f"\n - Additional batch inputs: {[p.name for p in self._additional_batch_inputs]}" + + placement_section = "\n\nThis block should be placed after the encoder steps and the text input step." + + return summary_section + inputs_info + placement_section + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("pachifier", QwenImageLayeredPachifier, default_creation_method="from_config"), + ] + + @property + def inputs(self) -> list[InputParam]: + inputs = [ + InputParam.template("num_images_per_prompt"), + InputParam.template("batch_size"), + ] + # default is `image_latents` + + inputs += self._image_latent_inputs + self._additional_batch_inputs + + return inputs + + @property + def intermediate_outputs(self) -> list[OutputParam]: + outputs = [ + OutputParam( + name="image_height", + type_hint=int, + description="The image height calculated from the image latents dimension", + ), + OutputParam( + name="image_width", + type_hint=int, + description="The image width calculated from the image latents dimension", + ), + ] + + if len(self._image_latent_inputs) > 0: + outputs.append( + OutputParam(name="height", type_hint=int, description="if not provided, updated to image height") + ) + outputs.append( + OutputParam(name="width", type_hint=int, description="if not provided, updated to image width") + ) + + # Add outputs for image latent inputs (patchified with layered pachifier and batch-expanded) + for input_param in self._image_latent_inputs: + outputs.append( + OutputParam( + name=input_param.name, + type_hint=input_param.type_hint, + description=input_param.description + " (patchified with layered pachifier and batch-expanded)", + ) + ) + + # Add outputs for additional batch inputs (batch-expanded only) + for input_param in self._additional_batch_inputs: + outputs.append( + OutputParam( + name=input_param.name, + type_hint=input_param.type_hint, + description=input_param.description + " (batch-expanded)", + ) + ) + + return outputs + + def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + # Process image latent inputs + for input_param in self._image_latent_inputs: + image_latent_input_name = input_param.name + image_latent_tensor = getattr(block_state, image_latent_input_name) + if image_latent_tensor is None: + continue + + # 1. Calculate height/width from latents and update if not provided + # Layered latents are (B, layers, C, H, W) + height = image_latent_tensor.shape[3] * components.vae_scale_factor + width = image_latent_tensor.shape[4] * components.vae_scale_factor + block_state.height = height + block_state.width = width + + if not hasattr(block_state, "image_height"): + block_state.image_height = height + if not hasattr(block_state, "image_width"): + block_state.image_width = width + + # 2. Patchify with layered pachifier + image_latent_tensor = components.pachifier.pack_latents(image_latent_tensor) + + # 3. Expand batch size + image_latent_tensor = repeat_tensor_to_batch_size( + input_name=image_latent_input_name, + input_tensor=image_latent_tensor, + num_images_per_prompt=block_state.num_images_per_prompt, + batch_size=block_state.batch_size, + ) + + setattr(block_state, image_latent_input_name, image_latent_tensor) + # Process additional batch inputs (only batch expansion) - for input_name in self._additional_batch_inputs: + for input_param in self._additional_batch_inputs: + input_name = input_param.name input_tensor = getattr(block_state, input_name) if input_tensor is None: continue - # Only expand batch size input_tensor = repeat_tensor_to_batch_size( input_name=input_name, input_tensor=input_tensor, @@ -438,7 +909,34 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) - return components, state +# auto_docstring class QwenImageControlNetInputsStep(ModularPipelineBlocks): + """ + prepare the `control_image_latents` for controlnet. Insert after all the other inputs steps. + + Inputs: + control_image_latents (`Tensor`): + The control image latents to use for the denoising process. Can be generated in controlnet vae encoder + step. + batch_size (`int`, *optional*, defaults to 1): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + + Outputs: + control_image_latents (`Tensor`): + The control image latents (patchified and batch-expanded). + height (`int`): + if not provided, updated to control image height + width (`int`): + if not provided, updated to control image width + """ + model_name = "qwenimage" @property @@ -446,13 +944,30 @@ def description(self) -> str: return "prepare the `control_image_latents` for controlnet. Insert after all the other inputs steps." @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ - InputParam(name="control_image_latents", required=True), - InputParam(name="batch_size", required=True), - InputParam(name="num_images_per_prompt", default=1), - InputParam(name="height"), - InputParam(name="width"), + InputParam( + name="control_image_latents", + required=True, + type_hint=torch.Tensor, + description="The control image latents to use for the denoising process. Can be generated in controlnet vae encoder step.", + ), + InputParam.template("batch_size"), + InputParam.template("num_images_per_prompt"), + InputParam.template("height"), + InputParam.template("width"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="control_image_latents", + type_hint=torch.Tensor, + description="The control image latents (patchified and batch-expanded).", + ), + OutputParam(name="height", type_hint=int, description="if not provided, updated to control image height"), + OutputParam(name="width", type_hint=int, description="if not provided, updated to control image width"), ] @torch.no_grad() diff --git a/src/diffusers/modular_pipelines/qwenimage/modular_blocks.py b/src/diffusers/modular_pipelines/qwenimage/modular_blocks.py deleted file mode 100644 index dcce0cab5dd1..000000000000 --- a/src/diffusers/modular_pipelines/qwenimage/modular_blocks.py +++ /dev/null @@ -1,1113 +0,0 @@ -# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from ...utils import logging -from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks -from ..modular_pipeline_utils import InsertableDict -from .before_denoise import ( - QwenImageControlNetBeforeDenoiserStep, - QwenImageCreateMaskLatentsStep, - QwenImageEditPlusRoPEInputsStep, - QwenImageEditRoPEInputsStep, - QwenImagePrepareLatentsStep, - QwenImagePrepareLatentsWithStrengthStep, - QwenImageRoPEInputsStep, - QwenImageSetTimestepsStep, - QwenImageSetTimestepsWithStrengthStep, -) -from .decoders import ( - QwenImageAfterDenoiseStep, - QwenImageDecoderStep, - QwenImageInpaintProcessImagesOutputStep, - QwenImageProcessImagesOutputStep, -) -from .denoise import ( - QwenImageControlNetDenoiseStep, - QwenImageDenoiseStep, - QwenImageEditDenoiseStep, - QwenImageEditInpaintDenoiseStep, - QwenImageInpaintControlNetDenoiseStep, - QwenImageInpaintDenoiseStep, - QwenImageLoopBeforeDenoiserControlNet, -) -from .encoders import ( - QwenImageControlNetVaeEncoderStep, - QwenImageEditPlusProcessImagesInputStep, - QwenImageEditPlusResizeDynamicStep, - QwenImageEditPlusTextEncoderStep, - QwenImageEditPlusVaeEncoderDynamicStep, - QwenImageEditResizeDynamicStep, - QwenImageEditTextEncoderStep, - QwenImageInpaintProcessImagesInputStep, - QwenImageProcessImagesInputStep, - QwenImageTextEncoderStep, - QwenImageVaeEncoderDynamicStep, -) -from .inputs import ( - QwenImageControlNetInputsStep, - QwenImageEditPlusInputsDynamicStep, - QwenImageInputsDynamicStep, - QwenImageTextInputsStep, -) - - -logger = logging.get_logger(__name__) - -# 1. QwenImage - -## 1.1 QwenImage/text2image - -#### QwenImage/decode -#### (standard decode step works for most tasks except for inpaint) -QwenImageDecodeBlocks = InsertableDict( - [ - ("decode", QwenImageDecoderStep()), - ("postprocess", QwenImageProcessImagesOutputStep()), - ] -) - - -class QwenImageDecodeStep(SequentialPipelineBlocks): - model_name = "qwenimage" - block_classes = QwenImageDecodeBlocks.values() - block_names = QwenImageDecodeBlocks.keys() - - @property - def description(self): - return "Decode step that decodes the latents to images and postprocess the generated image." - - -#### QwenImage/text2image presets -TEXT2IMAGE_BLOCKS = InsertableDict( - [ - ("text_encoder", QwenImageTextEncoderStep()), - ("input", QwenImageTextInputsStep()), - ("prepare_latents", QwenImagePrepareLatentsStep()), - ("set_timesteps", QwenImageSetTimestepsStep()), - ("prepare_rope_inputs", QwenImageRoPEInputsStep()), - ("denoise", QwenImageDenoiseStep()), - ("after_denoise", QwenImageAfterDenoiseStep()), - ("decode", QwenImageDecodeStep()), - ] -) - - -## 1.2 QwenImage/inpaint - -#### QwenImage/inpaint vae encoder -QwenImageInpaintVaeEncoderBlocks = InsertableDict( - [ - ( - "preprocess", - QwenImageInpaintProcessImagesInputStep, - ), # image, mask_image -> processed_image, processed_mask_image, mask_overlay_kwargs - ("encode", QwenImageVaeEncoderDynamicStep()), # processed_image -> image_latents - ] -) - - -class QwenImageInpaintVaeEncoderStep(SequentialPipelineBlocks): - model_name = "qwenimage" - block_classes = QwenImageInpaintVaeEncoderBlocks.values() - block_names = QwenImageInpaintVaeEncoderBlocks.keys() - - @property - def description(self) -> str: - return ( - "This step is used for processing image and mask inputs for inpainting tasks. It:\n" - " - Resizes the image to the target size, based on `height` and `width`.\n" - " - Processes and updates `image` and `mask_image`.\n" - " - Creates `image_latents`." - ) - - -#### QwenImage/inpaint inputs -QwenImageInpaintInputBlocks = InsertableDict( - [ - ("text_inputs", QwenImageTextInputsStep()), # default step to process text embeddings - ( - "additional_inputs", - QwenImageInputsDynamicStep( - image_latent_inputs=["image_latents"], additional_batch_inputs=["processed_mask_image"] - ), - ), - ] -) - - -class QwenImageInpaintInputStep(SequentialPipelineBlocks): - model_name = "qwenimage" - block_classes = QwenImageInpaintInputBlocks.values() - block_names = QwenImageInpaintInputBlocks.keys() - - @property - def description(self): - return "Input step that prepares the inputs for the inpainting denoising step. It:\n" - " - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents` and `processed_mask_image`).\n" - " - update height/width based `image_latents`, patchify `image_latents`." - - -# QwenImage/inpaint prepare latents -QwenImageInpaintPrepareLatentsBlocks = InsertableDict( - [ - ("add_noise_to_latents", QwenImagePrepareLatentsWithStrengthStep()), - ("create_mask_latents", QwenImageCreateMaskLatentsStep()), - ] -) - - -class QwenImageInpaintPrepareLatentsStep(SequentialPipelineBlocks): - model_name = "qwenimage" - block_classes = QwenImageInpaintPrepareLatentsBlocks.values() - block_names = QwenImageInpaintPrepareLatentsBlocks.keys() - - @property - def description(self) -> str: - return ( - "This step prepares the latents/image_latents and mask inputs for the inpainting denoising step. It:\n" - " - Add noise to the image latents to create the latents input for the denoiser.\n" - " - Create the pachified latents `mask` based on the processedmask image.\n" - ) - - -#### QwenImage/inpaint decode -QwenImageInpaintDecodeBlocks = InsertableDict( - [ - ("decode", QwenImageDecoderStep()), - ("postprocess", QwenImageInpaintProcessImagesOutputStep()), - ] -) - - -class QwenImageInpaintDecodeStep(SequentialPipelineBlocks): - model_name = "qwenimage" - block_classes = QwenImageInpaintDecodeBlocks.values() - block_names = QwenImageInpaintDecodeBlocks.keys() - - @property - def description(self): - return "Decode step that decodes the latents to images and postprocess the generated image, optional apply the mask overally to the original image." - - -#### QwenImage/inpaint presets -INPAINT_BLOCKS = InsertableDict( - [ - ("text_encoder", QwenImageTextEncoderStep()), - ("vae_encoder", QwenImageInpaintVaeEncoderStep()), - ("input", QwenImageInpaintInputStep()), - ("prepare_latents", QwenImagePrepareLatentsStep()), - ("set_timesteps", QwenImageSetTimestepsWithStrengthStep()), - ("prepare_inpaint_latents", QwenImageInpaintPrepareLatentsStep()), - ("prepare_rope_inputs", QwenImageRoPEInputsStep()), - ("denoise", QwenImageInpaintDenoiseStep()), - ("after_denoise", QwenImageAfterDenoiseStep()), - ("decode", QwenImageInpaintDecodeStep()), - ] -) - - -## 1.3 QwenImage/img2img - -#### QwenImage/img2img vae encoder -QwenImageImg2ImgVaeEncoderBlocks = InsertableDict( - [ - ("preprocess", QwenImageProcessImagesInputStep()), - ("encode", QwenImageVaeEncoderDynamicStep()), - ] -) - - -class QwenImageImg2ImgVaeEncoderStep(SequentialPipelineBlocks): - model_name = "qwenimage" - - block_classes = QwenImageImg2ImgVaeEncoderBlocks.values() - block_names = QwenImageImg2ImgVaeEncoderBlocks.keys() - - @property - def description(self) -> str: - return "Vae encoder step that preprocess andencode the image inputs into their latent representations." - - -#### QwenImage/img2img inputs -QwenImageImg2ImgInputBlocks = InsertableDict( - [ - ("text_inputs", QwenImageTextInputsStep()), # default step to process text embeddings - ("additional_inputs", QwenImageInputsDynamicStep(image_latent_inputs=["image_latents"])), - ] -) - - -class QwenImageImg2ImgInputStep(SequentialPipelineBlocks): - model_name = "qwenimage" - block_classes = QwenImageImg2ImgInputBlocks.values() - block_names = QwenImageImg2ImgInputBlocks.keys() - - @property - def description(self): - return "Input step that prepares the inputs for the img2img denoising step. It:\n" - " - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents`).\n" - " - update height/width based `image_latents`, patchify `image_latents`." - - -#### QwenImage/img2img presets -IMAGE2IMAGE_BLOCKS = InsertableDict( - [ - ("text_encoder", QwenImageTextEncoderStep()), - ("vae_encoder", QwenImageImg2ImgVaeEncoderStep()), - ("input", QwenImageImg2ImgInputStep()), - ("prepare_latents", QwenImagePrepareLatentsStep()), - ("set_timesteps", QwenImageSetTimestepsWithStrengthStep()), - ("prepare_img2img_latents", QwenImagePrepareLatentsWithStrengthStep()), - ("prepare_rope_inputs", QwenImageRoPEInputsStep()), - ("denoise", QwenImageDenoiseStep()), - ("after_denoise", QwenImageAfterDenoiseStep()), - ("decode", QwenImageDecodeStep()), - ] -) - - -## 1.4 QwenImage/controlnet - -#### QwenImage/controlnet presets -CONTROLNET_BLOCKS = InsertableDict( - [ - ("controlnet_vae_encoder", QwenImageControlNetVaeEncoderStep()), # vae encoder step for control_image - ("controlnet_inputs", QwenImageControlNetInputsStep()), # additional input step for controlnet - ( - "controlnet_before_denoise", - QwenImageControlNetBeforeDenoiserStep(), - ), # before denoise step (after set_timesteps step) - ( - "controlnet_denoise_loop_before", - QwenImageLoopBeforeDenoiserControlNet(), - ), # controlnet loop step (insert before the denoiseloop_denoiser) - ] -) - - -## 1.5 QwenImage/auto encoders - - -#### for inpaint and img2img tasks -class QwenImageAutoVaeEncoderStep(AutoPipelineBlocks): - block_classes = [QwenImageInpaintVaeEncoderStep, QwenImageImg2ImgVaeEncoderStep] - block_names = ["inpaint", "img2img"] - block_trigger_inputs = ["mask_image", "image"] - - @property - def description(self): - return ( - "Vae encoder step that encode the image inputs into their latent representations.\n" - + "This is an auto pipeline block.\n" - + " - `QwenImageInpaintVaeEncoderStep` (inpaint) is used when `mask_image` is provided.\n" - + " - `QwenImageImg2ImgVaeEncoderStep` (img2img) is used when `image` is provided.\n" - + " - if `mask_image` or `image` is not provided, step will be skipped." - ) - - -# for controlnet tasks -class QwenImageOptionalControlNetVaeEncoderStep(AutoPipelineBlocks): - block_classes = [QwenImageControlNetVaeEncoderStep] - block_names = ["controlnet"] - block_trigger_inputs = ["control_image"] - - @property - def description(self): - return ( - "Vae encoder step that encode the image inputs into their latent representations.\n" - + "This is an auto pipeline block.\n" - + " - `QwenImageControlNetVaeEncoderStep` (controlnet) is used when `control_image` is provided.\n" - + " - if `control_image` is not provided, step will be skipped." - ) - - -## 1.6 QwenImage/auto inputs - - -# text2image/inpaint/img2img -class QwenImageAutoInputStep(AutoPipelineBlocks): - block_classes = [QwenImageInpaintInputStep, QwenImageImg2ImgInputStep, QwenImageTextInputsStep] - block_names = ["inpaint", "img2img", "text2image"] - block_trigger_inputs = ["processed_mask_image", "image_latents", None] - - @property - def description(self): - return ( - "Input step that standardize the inputs for the denoising step, e.g. make sure inputs have consistent batch size, and patchified. \n" - " This is an auto pipeline block that works for text2image/inpaint/img2img tasks.\n" - + " - `QwenImageInpaintInputStep` (inpaint) is used when `processed_mask_image` is provided.\n" - + " - `QwenImageImg2ImgInputStep` (img2img) is used when `image_latents` is provided.\n" - + " - `QwenImageTextInputsStep` (text2image) is used when both `processed_mask_image` and `image_latents` are not provided.\n" - ) - - -# controlnet -class QwenImageOptionalControlNetInputStep(AutoPipelineBlocks): - block_classes = [QwenImageControlNetInputsStep] - block_names = ["controlnet"] - block_trigger_inputs = ["control_image_latents"] - - @property - def description(self): - return ( - "Controlnet input step that prepare the control_image_latents input.\n" - + "This is an auto pipeline block.\n" - + " - `QwenImageControlNetInputsStep` (controlnet) is used when `control_image_latents` is provided.\n" - + " - if `control_image_latents` is not provided, step will be skipped." - ) - - -## 1.7 QwenImage/auto before denoise step -# compose the steps into a BeforeDenoiseStep for text2image/img2img/inpaint tasks before combine into an auto step - -# QwenImage/text2image before denoise -QwenImageText2ImageBeforeDenoiseBlocks = InsertableDict( - [ - ("prepare_latents", QwenImagePrepareLatentsStep()), - ("set_timesteps", QwenImageSetTimestepsStep()), - ("prepare_rope_inputs", QwenImageRoPEInputsStep()), - ] -) - - -class QwenImageText2ImageBeforeDenoiseStep(SequentialPipelineBlocks): - model_name = "qwenimage" - block_classes = QwenImageText2ImageBeforeDenoiseBlocks.values() - block_names = QwenImageText2ImageBeforeDenoiseBlocks.keys() - - @property - def description(self): - return "Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for text2image task." - - -# QwenImage/inpaint before denoise -QwenImageInpaintBeforeDenoiseBlocks = InsertableDict( - [ - ("prepare_latents", QwenImagePrepareLatentsStep()), - ("set_timesteps", QwenImageSetTimestepsWithStrengthStep()), - ("prepare_inpaint_latents", QwenImageInpaintPrepareLatentsStep()), - ("prepare_rope_inputs", QwenImageRoPEInputsStep()), - ] -) - - -class QwenImageInpaintBeforeDenoiseStep(SequentialPipelineBlocks): - model_name = "qwenimage" - block_classes = QwenImageInpaintBeforeDenoiseBlocks.values() - block_names = QwenImageInpaintBeforeDenoiseBlocks.keys() - - @property - def description(self): - return "Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for inpaint task." - - -# QwenImage/img2img before denoise -QwenImageImg2ImgBeforeDenoiseBlocks = InsertableDict( - [ - ("prepare_latents", QwenImagePrepareLatentsStep()), - ("set_timesteps", QwenImageSetTimestepsWithStrengthStep()), - ("prepare_img2img_latents", QwenImagePrepareLatentsWithStrengthStep()), - ("prepare_rope_inputs", QwenImageRoPEInputsStep()), - ] -) - - -class QwenImageImg2ImgBeforeDenoiseStep(SequentialPipelineBlocks): - model_name = "qwenimage" - block_classes = QwenImageImg2ImgBeforeDenoiseBlocks.values() - block_names = QwenImageImg2ImgBeforeDenoiseBlocks.keys() - - @property - def description(self): - return "Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for img2img task." - - -# auto before_denoise step for text2image, inpaint, img2img tasks -class QwenImageAutoBeforeDenoiseStep(AutoPipelineBlocks): - block_classes = [ - QwenImageInpaintBeforeDenoiseStep, - QwenImageImg2ImgBeforeDenoiseStep, - QwenImageText2ImageBeforeDenoiseStep, - ] - block_names = ["inpaint", "img2img", "text2image"] - block_trigger_inputs = ["processed_mask_image", "image_latents", None] - - @property - def description(self): - return ( - "Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step.\n" - + "This is an auto pipeline block that works for text2img, inpainting, img2img tasks.\n" - + " - `QwenImageInpaintBeforeDenoiseStep` (inpaint) is used when `processed_mask_image` is provided.\n" - + " - `QwenImageImg2ImgBeforeDenoiseStep` (img2img) is used when `image_latents` is provided.\n" - + " - `QwenImageText2ImageBeforeDenoiseStep` (text2image) is used when both `processed_mask_image` and `image_latents` are not provided.\n" - ) - - -# auto before_denoise step for controlnet tasks -class QwenImageOptionalControlNetBeforeDenoiseStep(AutoPipelineBlocks): - block_classes = [QwenImageControlNetBeforeDenoiserStep] - block_names = ["controlnet"] - block_trigger_inputs = ["control_image_latents"] - - @property - def description(self): - return ( - "Controlnet before denoise step that prepare the controlnet input.\n" - + "This is an auto pipeline block.\n" - + " - `QwenImageControlNetBeforeDenoiserStep` (controlnet) is used when `control_image_latents` is provided.\n" - + " - if `control_image_latents` is not provided, step will be skipped." - ) - - -## 1.8 QwenImage/auto denoise - - -# auto denoise step for controlnet tasks: works for all tasks with controlnet -class QwenImageControlNetAutoDenoiseStep(AutoPipelineBlocks): - block_classes = [QwenImageInpaintControlNetDenoiseStep, QwenImageControlNetDenoiseStep] - block_names = ["inpaint_denoise", "denoise"] - block_trigger_inputs = ["mask", None] - - @property - def description(self): - return ( - "Controlnet step during the denoising process. \n" - " This is an auto pipeline block that works for inpaint and text2image/img2img tasks with controlnet.\n" - + " - `QwenImageInpaintControlNetDenoiseStep` (inpaint) is used when `mask` is provided.\n" - + " - `QwenImageControlNetDenoiseStep` (text2image/img2img) is used when `mask` is not provided.\n" - ) - - -# auto denoise step for everything: works for all tasks with or without controlnet -class QwenImageAutoDenoiseStep(AutoPipelineBlocks): - block_classes = [ - QwenImageControlNetAutoDenoiseStep, - QwenImageInpaintDenoiseStep, - QwenImageDenoiseStep, - ] - block_names = ["controlnet_denoise", "inpaint_denoise", "denoise"] - block_trigger_inputs = ["control_image_latents", "mask", None] - - @property - def description(self): - return ( - "Denoise step that iteratively denoise the latents. \n" - " This is an auto pipeline block that works for inpaint/text2image/img2img tasks. It also works with controlnet\n" - + " - `QwenImageControlNetAutoDenoiseStep` (controlnet) is used when `control_image_latents` is provided.\n" - + " - `QwenImageInpaintDenoiseStep` (inpaint) is used when `mask` is provided and `control_image_latents` is not provided.\n" - + " - `QwenImageDenoiseStep` (text2image/img2img) is used when `mask` is not provided and `control_image_latents` is not provided.\n" - ) - - -## 1.9 QwenImage/auto decode -# auto decode step for inpaint and text2image tasks - - -class QwenImageAutoDecodeStep(AutoPipelineBlocks): - block_classes = [QwenImageInpaintDecodeStep, QwenImageDecodeStep] - block_names = ["inpaint_decode", "decode"] - block_trigger_inputs = ["mask", None] - - @property - def description(self): - return ( - "Decode step that decode the latents into images. \n" - " This is an auto pipeline block that works for inpaint/text2image/img2img tasks, for both QwenImage and QwenImage-Edit.\n" - + " - `QwenImageInpaintDecodeStep` (inpaint) is used when `mask` is provided.\n" - + " - `QwenImageDecodeStep` (text2image/img2img) is used when `mask` is not provided.\n" - ) - - -class QwenImageCoreDenoiseStep(SequentialPipelineBlocks): - model_name = "qwenimage" - block_classes = [ - QwenImageAutoInputStep, - QwenImageOptionalControlNetInputStep, - QwenImageAutoBeforeDenoiseStep, - QwenImageOptionalControlNetBeforeDenoiseStep, - QwenImageAutoDenoiseStep, - QwenImageAfterDenoiseStep, - ] - block_names = [ - "input", - "controlnet_input", - "before_denoise", - "controlnet_before_denoise", - "denoise", - "after_denoise", - ] - - @property - def description(self): - return ( - "Core step that performs the denoising process. \n" - + " - `QwenImageAutoInputStep` (input) standardizes the inputs for the denoising step.\n" - + " - `QwenImageOptionalControlNetInputStep` (controlnet_input) prepares the controlnet input.\n" - + " - `QwenImageAutoBeforeDenoiseStep` (before_denoise) prepares the inputs for the denoising step.\n" - + " - `QwenImageOptionalControlNetBeforeDenoiseStep` (controlnet_before_denoise) prepares the controlnet input for the denoising step.\n" - + " - `QwenImageAutoDenoiseStep` (denoise) iteratively denoises the latents.\n" - + "This step support text-to-image, image-to-image, inpainting, and controlnet tasks for QwenImage:\n" - + " - for image-to-image generation, you need to provide `image_latents`\n" - + " - for inpainting, you need to provide `processed_mask_image` and `image_latents`\n" - + " - to run the controlnet workflow, you need to provide `control_image_latents`\n" - + " - for text-to-image generation, all you need to provide is prompt embeddings" - ) - - -## 1.10 QwenImage/auto block & presets -AUTO_BLOCKS = InsertableDict( - [ - ("text_encoder", QwenImageTextEncoderStep()), - ("vae_encoder", QwenImageAutoVaeEncoderStep()), - ("controlnet_vae_encoder", QwenImageOptionalControlNetVaeEncoderStep()), - ("denoise", QwenImageCoreDenoiseStep()), - ("decode", QwenImageAutoDecodeStep()), - ] -) - - -class QwenImageAutoBlocks(SequentialPipelineBlocks): - model_name = "qwenimage" - - block_classes = AUTO_BLOCKS.values() - block_names = AUTO_BLOCKS.keys() - - @property - def description(self): - return ( - "Auto Modular pipeline for text-to-image, image-to-image, inpainting, and controlnet tasks using QwenImage.\n" - + "- for image-to-image generation, you need to provide `image`\n" - + "- for inpainting, you need to provide `mask_image` and `image`, optionally you can provide `padding_mask_crop` \n" - + "- to run the controlnet workflow, you need to provide `control_image`\n" - + "- for text-to-image generation, all you need to provide is `prompt`" - ) - - -# 2. QwenImage-Edit - -## 2.1 QwenImage-Edit/edit - -#### QwenImage-Edit/edit vl encoder: take both image and text prompts -QwenImageEditVLEncoderBlocks = InsertableDict( - [ - ("resize", QwenImageEditResizeDynamicStep()), - ("encode", QwenImageEditTextEncoderStep()), - ] -) - - -class QwenImageEditVLEncoderStep(SequentialPipelineBlocks): - model_name = "qwenimage" - block_classes = QwenImageEditVLEncoderBlocks.values() - block_names = QwenImageEditVLEncoderBlocks.keys() - - @property - def description(self) -> str: - return "QwenImage-Edit VL encoder step that encode the image an text prompts together." - - -#### QwenImage-Edit/edit vae encoder -QwenImageEditVaeEncoderBlocks = InsertableDict( - [ - ("resize", QwenImageEditResizeDynamicStep()), # edit has a different resize step - ("preprocess", QwenImageProcessImagesInputStep()), # resized_image -> processed_image - ("encode", QwenImageVaeEncoderDynamicStep()), # processed_image -> image_latents - ] -) - - -class QwenImageEditVaeEncoderStep(SequentialPipelineBlocks): - model_name = "qwenimage" - block_classes = QwenImageEditVaeEncoderBlocks.values() - block_names = QwenImageEditVaeEncoderBlocks.keys() - - @property - def description(self) -> str: - return "Vae encoder step that encode the image inputs into their latent representations." - - -#### QwenImage-Edit/edit input -QwenImageEditInputBlocks = InsertableDict( - [ - ("text_inputs", QwenImageTextInputsStep()), # default step to process text embeddings - ("additional_inputs", QwenImageInputsDynamicStep(image_latent_inputs=["image_latents"])), - ] -) - - -class QwenImageEditInputStep(SequentialPipelineBlocks): - model_name = "qwenimage" - block_classes = QwenImageEditInputBlocks.values() - block_names = QwenImageEditInputBlocks.keys() - - @property - def description(self): - return "Input step that prepares the inputs for the edit denoising step. It:\n" - " - make sure the text embeddings have consistent batch size as well as the additional inputs: \n" - " - `image_latents`.\n" - " - update height/width based `image_latents`, patchify `image_latents`." - - -#### QwenImage/edit presets -EDIT_BLOCKS = InsertableDict( - [ - ("text_encoder", QwenImageEditVLEncoderStep()), - ("vae_encoder", QwenImageEditVaeEncoderStep()), - ("input", QwenImageEditInputStep()), - ("prepare_latents", QwenImagePrepareLatentsStep()), - ("set_timesteps", QwenImageSetTimestepsStep()), - ("prepare_rope_inputs", QwenImageEditRoPEInputsStep()), - ("denoise", QwenImageEditDenoiseStep()), - ("after_denoise", QwenImageAfterDenoiseStep()), - ("decode", QwenImageDecodeStep()), - ] -) - - -## 2.2 QwenImage-Edit/edit inpaint - -#### QwenImage-Edit/edit inpaint vae encoder: the difference from regular inpaint is the resize step -QwenImageEditInpaintVaeEncoderBlocks = InsertableDict( - [ - ("resize", QwenImageEditResizeDynamicStep()), # image -> resized_image - ( - "preprocess", - QwenImageInpaintProcessImagesInputStep, - ), # resized_image, mask_image -> processed_image, processed_mask_image, mask_overlay_kwargs - ( - "encode", - QwenImageVaeEncoderDynamicStep(input_name="processed_image", output_name="image_latents"), - ), # processed_image -> image_latents - ] -) - - -class QwenImageEditInpaintVaeEncoderStep(SequentialPipelineBlocks): - model_name = "qwenimage" - block_classes = QwenImageEditInpaintVaeEncoderBlocks.values() - block_names = QwenImageEditInpaintVaeEncoderBlocks.keys() - - @property - def description(self) -> str: - return ( - "This step is used for processing image and mask inputs for QwenImage-Edit inpaint tasks. It:\n" - " - resize the image for target area (1024 * 1024) while maintaining the aspect ratio.\n" - " - process the resized image and mask image.\n" - " - create image latents." - ) - - -#### QwenImage-Edit/edit inpaint presets -EDIT_INPAINT_BLOCKS = InsertableDict( - [ - ("text_encoder", QwenImageEditVLEncoderStep()), - ("vae_encoder", QwenImageEditInpaintVaeEncoderStep()), - ("input", QwenImageInpaintInputStep()), - ("prepare_latents", QwenImagePrepareLatentsStep()), - ("set_timesteps", QwenImageSetTimestepsWithStrengthStep()), - ("prepare_inpaint_latents", QwenImageInpaintPrepareLatentsStep()), - ("prepare_rope_inputs", QwenImageEditRoPEInputsStep()), - ("denoise", QwenImageEditInpaintDenoiseStep()), - ("after_denoise", QwenImageAfterDenoiseStep()), - ("decode", QwenImageInpaintDecodeStep()), - ] -) - - -## 2.3 QwenImage-Edit/auto encoders - - -class QwenImageEditAutoVaeEncoderStep(AutoPipelineBlocks): - block_classes = [ - QwenImageEditInpaintVaeEncoderStep, - QwenImageEditVaeEncoderStep, - ] - block_names = ["edit_inpaint", "edit"] - block_trigger_inputs = ["mask_image", "image"] - - @property - def description(self): - return ( - "Vae encoder step that encode the image inputs into their latent representations. \n" - " This is an auto pipeline block that works for edit and edit_inpaint tasks.\n" - + " - `QwenImageEditInpaintVaeEncoderStep` (edit_inpaint) is used when `mask_image` is provided.\n" - + " - `QwenImageEditVaeEncoderStep` (edit) is used when `image` is provided.\n" - + " - if `mask_image` or `image` is not provided, step will be skipped." - ) - - -## 2.4 QwenImage-Edit/auto inputs -class QwenImageEditAutoInputStep(AutoPipelineBlocks): - block_classes = [QwenImageInpaintInputStep, QwenImageEditInputStep] - block_names = ["edit_inpaint", "edit"] - block_trigger_inputs = ["processed_mask_image", "image_latents"] - - @property - def description(self): - return ( - "Input step that prepares the inputs for the edit denoising step.\n" - + " It is an auto pipeline block that works for edit and edit_inpaint tasks.\n" - + " - `QwenImageInpaintInputStep` (edit_inpaint) is used when `processed_mask_image` is provided.\n" - + " - `QwenImageEditInputStep` (edit) is used when `image_latents` is provided.\n" - + " - if `processed_mask_image` or `image_latents` is not provided, step will be skipped." - ) - - -## 2.5 QwenImage-Edit/auto before denoise -# compose the steps into a BeforeDenoiseStep for edit and edit_inpaint tasks before combine into an auto step - -#### QwenImage-Edit/edit before denoise -QwenImageEditBeforeDenoiseBlocks = InsertableDict( - [ - ("prepare_latents", QwenImagePrepareLatentsStep()), - ("set_timesteps", QwenImageSetTimestepsStep()), - ("prepare_rope_inputs", QwenImageEditRoPEInputsStep()), - ] -) - - -class QwenImageEditBeforeDenoiseStep(SequentialPipelineBlocks): - model_name = "qwenimage" - block_classes = QwenImageEditBeforeDenoiseBlocks.values() - block_names = QwenImageEditBeforeDenoiseBlocks.keys() - - @property - def description(self): - return "Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for edit task." - - -#### QwenImage-Edit/edit inpaint before denoise -QwenImageEditInpaintBeforeDenoiseBlocks = InsertableDict( - [ - ("prepare_latents", QwenImagePrepareLatentsStep()), - ("set_timesteps", QwenImageSetTimestepsWithStrengthStep()), - ("prepare_inpaint_latents", QwenImageInpaintPrepareLatentsStep()), - ("prepare_rope_inputs", QwenImageEditRoPEInputsStep()), - ] -) - - -class QwenImageEditInpaintBeforeDenoiseStep(SequentialPipelineBlocks): - model_name = "qwenimage" - block_classes = QwenImageEditInpaintBeforeDenoiseBlocks.values() - block_names = QwenImageEditInpaintBeforeDenoiseBlocks.keys() - - @property - def description(self): - return "Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for edit inpaint task." - - -# auto before_denoise step for edit and edit_inpaint tasks -class QwenImageEditAutoBeforeDenoiseStep(AutoPipelineBlocks): - model_name = "qwenimage-edit" - block_classes = [ - QwenImageEditInpaintBeforeDenoiseStep, - QwenImageEditBeforeDenoiseStep, - ] - block_names = ["edit_inpaint", "edit"] - block_trigger_inputs = ["processed_mask_image", "image_latents"] - - @property - def description(self): - return ( - "Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step.\n" - + "This is an auto pipeline block that works for edit (img2img) and edit inpaint tasks.\n" - + " - `QwenImageEditInpaintBeforeDenoiseStep` (edit_inpaint) is used when `processed_mask_image` is provided.\n" - + " - `QwenImageEditBeforeDenoiseStep` (edit) is used when `image_latents` is provided and `processed_mask_image` is not provided.\n" - + " - if `image_latents` or `processed_mask_image` is not provided, step will be skipped." - ) - - -## 2.6 QwenImage-Edit/auto denoise - - -class QwenImageEditAutoDenoiseStep(AutoPipelineBlocks): - model_name = "qwenimage-edit" - - block_classes = [QwenImageEditInpaintDenoiseStep, QwenImageEditDenoiseStep] - block_names = ["inpaint_denoise", "denoise"] - block_trigger_inputs = ["processed_mask_image", "image_latents"] - - @property - def description(self): - return ( - "Denoise step that iteratively denoise the latents. \n" - + "This block supports edit (img2img) and edit inpaint tasks for QwenImage Edit. \n" - + " - `QwenImageEditInpaintDenoiseStep` (inpaint) is used when `processed_mask_image` is provided.\n" - + " - `QwenImageEditDenoiseStep` (img2img) is used when `image_latents` is provided.\n" - + " - if `processed_mask_image` or `image_latents` is not provided, step will be skipped." - ) - - -## 2.7 QwenImage-Edit/auto blocks & presets - - -class QwenImageEditCoreDenoiseStep(SequentialPipelineBlocks): - model_name = "qwenimage-edit" - block_classes = [ - QwenImageEditAutoInputStep, - QwenImageEditAutoBeforeDenoiseStep, - QwenImageEditAutoDenoiseStep, - QwenImageAfterDenoiseStep, - ] - block_names = ["input", "before_denoise", "denoise", "after_denoise"] - - @property - def description(self): - return ( - "Core step that performs the denoising process. \n" - + " - `QwenImageEditAutoInputStep` (input) standardizes the inputs for the denoising step.\n" - + " - `QwenImageEditAutoBeforeDenoiseStep` (before_denoise) prepares the inputs for the denoising step.\n" - + " - `QwenImageEditAutoDenoiseStep` (denoise) iteratively denoises the latents.\n\n" - + "This step support edit (img2img) and edit inpainting workflow for QwenImage Edit:\n" - + " - When `processed_mask_image` is provided, it will be used for edit inpainting task.\n" - + " - When `image_latents` is provided, it will be used for edit (img2img) task.\n" - ) - - -EDIT_AUTO_BLOCKS = InsertableDict( - [ - ("text_encoder", QwenImageEditVLEncoderStep()), - ("vae_encoder", QwenImageEditAutoVaeEncoderStep()), - ("denoise", QwenImageEditCoreDenoiseStep()), - ("decode", QwenImageAutoDecodeStep()), - ] -) - - -class QwenImageEditAutoBlocks(SequentialPipelineBlocks): - model_name = "qwenimage-edit" - block_classes = EDIT_AUTO_BLOCKS.values() - block_names = EDIT_AUTO_BLOCKS.keys() - - @property - def description(self): - return ( - "Auto Modular pipeline for edit (img2img) and edit inpaint tasks using QwenImage-Edit.\n" - + "- for edit (img2img) generation, you need to provide `image`\n" - + "- for edit inpainting, you need to provide `mask_image` and `image`, optionally you can provide `padding_mask_crop` \n" - ) - - -#################### QwenImage Edit Plus ##################### - -# 3. QwenImage-Edit Plus - -## 3.1 QwenImage-Edit Plus / edit - -#### QwenImage-Edit Plus vl encoder: take both image and text prompts -QwenImageEditPlusVLEncoderBlocks = InsertableDict( - [ - ("resize", QwenImageEditPlusResizeDynamicStep()), - ("encode", QwenImageEditPlusTextEncoderStep()), - ] -) - - -class QwenImageEditPlusVLEncoderStep(SequentialPipelineBlocks): - model_name = "qwenimage" - block_classes = QwenImageEditPlusVLEncoderBlocks.values() - block_names = QwenImageEditPlusVLEncoderBlocks.keys() - - @property - def description(self) -> str: - return "QwenImage-Edit Plus VL encoder step that encode the image an text prompts together." - - -#### QwenImage-Edit Plus vae encoder -QwenImageEditPlusVaeEncoderBlocks = InsertableDict( - [ - ("resize", QwenImageEditPlusResizeDynamicStep()), # edit plus has a different resize step - ("preprocess", QwenImageEditPlusProcessImagesInputStep()), # vae_image -> processed_image - ("encode", QwenImageEditPlusVaeEncoderDynamicStep()), # processed_image -> image_latents - ] -) - - -class QwenImageEditPlusVaeEncoderStep(SequentialPipelineBlocks): - model_name = "qwenimage-edit-plus" - block_classes = QwenImageEditPlusVaeEncoderBlocks.values() - block_names = QwenImageEditPlusVaeEncoderBlocks.keys() - - @property - def description(self) -> str: - return "Vae encoder step that encode the image inputs into their latent representations." - - -#### QwenImage Edit Plus input blocks -QwenImageEditPlusInputBlocks = InsertableDict( - [ - ("text_inputs", QwenImageTextInputsStep()), # default step to process text embeddings - ( - "additional_inputs", - QwenImageEditPlusInputsDynamicStep(image_latent_inputs=["image_latents"]), - ), - ] -) - - -class QwenImageEditPlusInputStep(SequentialPipelineBlocks): - model_name = "qwenimage-edit-plus" - block_classes = QwenImageEditPlusInputBlocks.values() - block_names = QwenImageEditPlusInputBlocks.keys() - - -#### QwenImage Edit Plus presets -EDIT_PLUS_BLOCKS = InsertableDict( - [ - ("text_encoder", QwenImageEditPlusVLEncoderStep()), - ("vae_encoder", QwenImageEditPlusVaeEncoderStep()), - ("input", QwenImageEditPlusInputStep()), - ("prepare_latents", QwenImagePrepareLatentsStep()), - ("set_timesteps", QwenImageSetTimestepsStep()), - ("prepare_rope_inputs", QwenImageEditPlusRoPEInputsStep()), - ("denoise", QwenImageEditDenoiseStep()), - ("after_denoise", QwenImageAfterDenoiseStep()), - ("decode", QwenImageDecodeStep()), - ] -) - - -QwenImageEditPlusBeforeDenoiseBlocks = InsertableDict( - [ - ("prepare_latents", QwenImagePrepareLatentsStep()), - ("set_timesteps", QwenImageSetTimestepsStep()), - ("prepare_rope_inputs", QwenImageEditPlusRoPEInputsStep()), - ] -) - - -class QwenImageEditPlusBeforeDenoiseStep(SequentialPipelineBlocks): - model_name = "qwenimage-edit-plus" - block_classes = QwenImageEditPlusBeforeDenoiseBlocks.values() - block_names = QwenImageEditPlusBeforeDenoiseBlocks.keys() - - @property - def description(self): - return "Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for edit task." - - -# auto before_denoise step for edit tasks -class QwenImageEditPlusAutoBeforeDenoiseStep(AutoPipelineBlocks): - model_name = "qwenimage-edit-plus" - block_classes = [QwenImageEditPlusBeforeDenoiseStep] - block_names = ["edit"] - block_trigger_inputs = ["image_latents"] - - @property - def description(self): - return ( - "Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step.\n" - + "This is an auto pipeline block that works for edit (img2img) task.\n" - + " - `QwenImageEditPlusBeforeDenoiseStep` (edit) is used when `image_latents` is provided and `processed_mask_image` is not provided.\n" - + " - if `image_latents` is not provided, step will be skipped." - ) - - -## 3.2 QwenImage-Edit Plus/auto encoders - - -class QwenImageEditPlusAutoVaeEncoderStep(AutoPipelineBlocks): - block_classes = [QwenImageEditPlusVaeEncoderStep] - block_names = ["edit"] - block_trigger_inputs = ["image"] - - @property - def description(self): - return ( - "Vae encoder step that encode the image inputs into their latent representations. \n" - " This is an auto pipeline block that works for edit task.\n" - + " - `QwenImageEditPlusVaeEncoderStep` (edit) is used when `image` is provided.\n" - + " - if `image` is not provided, step will be skipped." - ) - - -## 3.3 QwenImage-Edit/auto blocks & presets - - -class QwenImageEditPlusAutoInputStep(AutoPipelineBlocks): - block_classes = [QwenImageEditPlusInputStep] - block_names = ["edit"] - block_trigger_inputs = ["image_latents"] - - @property - def description(self): - return ( - "Input step that prepares the inputs for the edit denoising step.\n" - + " It is an auto pipeline block that works for edit task.\n" - + " - `QwenImageEditPlusInputStep` (edit) is used when `image_latents` is provided.\n" - + " - if `image_latents` is not provided, step will be skipped." - ) - - -class QwenImageEditPlusCoreDenoiseStep(SequentialPipelineBlocks): - model_name = "qwenimage-edit-plus" - block_classes = [ - QwenImageEditPlusAutoInputStep, - QwenImageEditPlusAutoBeforeDenoiseStep, - QwenImageEditAutoDenoiseStep, - QwenImageAfterDenoiseStep, - ] - block_names = ["input", "before_denoise", "denoise", "after_denoise"] - - @property - def description(self): - return ( - "Core step that performs the denoising process. \n" - + " - `QwenImageEditAutoInputStep` (input) standardizes the inputs for the denoising step.\n" - + " - `QwenImageEditPlusAutoBeforeDenoiseStep` (before_denoise) prepares the inputs for the denoising step.\n" - + " - `QwenImageEditAutoDenoiseStep` (denoise) iteratively denoises the latents.\n\n" - + "This step support edit (img2img) workflow for QwenImage Edit Plus:\n" - + " - When `image_latents` is provided, it will be used for edit (img2img) task.\n" - ) - - -EDIT_PLUS_AUTO_BLOCKS = InsertableDict( - [ - ("text_encoder", QwenImageEditPlusVLEncoderStep()), - ("vae_encoder", QwenImageEditPlusAutoVaeEncoderStep()), - ("denoise", QwenImageEditPlusCoreDenoiseStep()), - ("decode", QwenImageAutoDecodeStep()), - ] -) - - -class QwenImageEditPlusAutoBlocks(SequentialPipelineBlocks): - model_name = "qwenimage-edit-plus" - block_classes = EDIT_PLUS_AUTO_BLOCKS.values() - block_names = EDIT_PLUS_AUTO_BLOCKS.keys() - - @property - def description(self): - return ( - "Auto Modular pipeline for edit (img2img) and edit tasks using QwenImage-Edit Plus.\n" - + "- for edit (img2img) generation, you need to provide `image`\n" - ) - - -# 3. all block presets supported in QwenImage, QwenImage-Edit, QwenImage-Edit Plus - - -ALL_BLOCKS = { - "text2image": TEXT2IMAGE_BLOCKS, - "img2img": IMAGE2IMAGE_BLOCKS, - "edit": EDIT_BLOCKS, - "edit_inpaint": EDIT_INPAINT_BLOCKS, - "edit_plus": EDIT_PLUS_BLOCKS, - "inpaint": INPAINT_BLOCKS, - "controlnet": CONTROLNET_BLOCKS, - "auto": AUTO_BLOCKS, - "edit_auto": EDIT_AUTO_BLOCKS, - "edit_plus_auto": EDIT_PLUS_AUTO_BLOCKS, -} diff --git a/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py b/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py new file mode 100644 index 000000000000..bf87028b2f90 --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py @@ -0,0 +1,1224 @@ +# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from ...utils import logging +from ..modular_pipeline import AutoPipelineBlocks, ConditionalPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import InputParam, InsertableDict, OutputParam +from .before_denoise import ( + QwenImageControlNetBeforeDenoiserStep, + QwenImageCreateMaskLatentsStep, + QwenImagePrepareLatentsStep, + QwenImagePrepareLatentsWithStrengthStep, + QwenImageRoPEInputsStep, + QwenImageSetTimestepsStep, + QwenImageSetTimestepsWithStrengthStep, +) +from .decoders import ( + QwenImageAfterDenoiseStep, + QwenImageDecoderStep, + QwenImageInpaintProcessImagesOutputStep, + QwenImageProcessImagesOutputStep, +) +from .denoise import ( + QwenImageControlNetDenoiseStep, + QwenImageDenoiseStep, + QwenImageInpaintControlNetDenoiseStep, + QwenImageInpaintDenoiseStep, +) +from .encoders import ( + QwenImageControlNetVaeEncoderStep, + QwenImageInpaintProcessImagesInputStep, + QwenImageProcessImagesInputStep, + QwenImageTextEncoderStep, + QwenImageVaeEncoderStep, +) +from .inputs import ( + QwenImageAdditionalInputsStep, + QwenImageControlNetInputsStep, + QwenImageTextInputsStep, +) + + +logger = logging.get_logger(__name__) + + +# ==================== +# 1. TEXT ENCODER +# ==================== + + +# auto_docstring +class QwenImageAutoTextEncoderStep(AutoPipelineBlocks): + """ + Text encoder step that encodes the text prompt into a text embedding. This is an auto pipeline block. + + Components: + text_encoder (`Qwen2_5_VLForConditionalGeneration`): The text encoder to use tokenizer (`Qwen2Tokenizer`): + The tokenizer to use guider (`ClassifierFreeGuidance`) + + Inputs: + prompt (`str`, *optional*): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + max_sequence_length (`int`, *optional*, defaults to 1024): + Maximum sequence length for prompt encoding. + + Outputs: + prompt_embeds (`Tensor`): + The prompt embeddings. + prompt_embeds_mask (`Tensor`): + The encoder attention mask. + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. + """ + + model_name = "qwenimage" + block_classes = [QwenImageTextEncoderStep()] + block_names = ["text_encoder"] + block_trigger_inputs = ["prompt"] + + @property + def description(self) -> str: + return "Text encoder step that encodes the text prompt into a text embedding. This is an auto pipeline block." + " - `QwenImageTextEncoderStep` (text_encoder) is used when `prompt` is provided." + " - if `prompt` is not provided, step will be skipped." + + +# ==================== +# 2. VAE ENCODER +# ==================== + + +# auto_docstring +class QwenImageInpaintVaeEncoderStep(SequentialPipelineBlocks): + """ + This step is used for processing image and mask inputs for inpainting tasks. It: + - Resizes the image to the target size, based on `height` and `width`. + - Processes and updates `image` and `mask_image`. + - Creates `image_latents`. + + Components: + image_mask_processor (`InpaintProcessor`) vae (`AutoencoderKLQwenImage`) + + Inputs: + mask_image (`Image`): + Mask image for inpainting. + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + padding_mask_crop (`int`, *optional*): + Padding for mask cropping in inpainting. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + processed_image (`Tensor`): + The processed image + processed_mask_image (`Tensor`): + The processed mask image + mask_overlay_kwargs (`dict`): + The kwargs for the postprocess step to apply the mask overlay + image_latents (`Tensor`): + The latent representation of the input image. + """ + + model_name = "qwenimage" + block_classes = [QwenImageInpaintProcessImagesInputStep(), QwenImageVaeEncoderStep()] + block_names = ["preprocess", "encode"] + + @property + def description(self) -> str: + return ( + "This step is used for processing image and mask inputs for inpainting tasks. It:\n" + " - Resizes the image to the target size, based on `height` and `width`.\n" + " - Processes and updates `image` and `mask_image`.\n" + " - Creates `image_latents`." + ) + + +# auto_docstring +class QwenImageImg2ImgVaeEncoderStep(SequentialPipelineBlocks): + """ + Vae encoder step that preprocess andencode the image inputs into their latent representations. + + Components: + image_processor (`VaeImageProcessor`) vae (`AutoencoderKLQwenImage`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + processed_image (`Tensor`): + The processed image + image_latents (`Tensor`): + The latent representation of the input image. + """ + + model_name = "qwenimage" + + block_classes = [QwenImageProcessImagesInputStep(), QwenImageVaeEncoderStep()] + block_names = ["preprocess", "encode"] + + @property + def description(self) -> str: + return "Vae encoder step that preprocess andencode the image inputs into their latent representations." + + +class QwenImageAutoVaeEncoderStep(AutoPipelineBlocks): + block_classes = [QwenImageInpaintVaeEncoderStep, QwenImageImg2ImgVaeEncoderStep] + block_names = ["inpaint", "img2img"] + block_trigger_inputs = ["mask_image", "image"] + + @property + def description(self): + return ( + "Vae encoder step that encode the image inputs into their latent representations.\n" + + "This is an auto pipeline block.\n" + + " - `QwenImageInpaintVaeEncoderStep` (inpaint) is used when `mask_image` is provided.\n" + + " - `QwenImageImg2ImgVaeEncoderStep` (img2img) is used when `image` is provided.\n" + + " - if `mask_image` or `image` is not provided, step will be skipped." + ) + + +# optional controlnet vae encoder +# auto_docstring +class QwenImageOptionalControlNetVaeEncoderStep(AutoPipelineBlocks): + """ + Vae encoder step that encode the image inputs into their latent representations. + This is an auto pipeline block. + - `QwenImageControlNetVaeEncoderStep` (controlnet) is used when `control_image` is provided. + - if `control_image` is not provided, step will be skipped. + + Components: + vae (`AutoencoderKLQwenImage`) controlnet (`QwenImageControlNetModel`) control_image_processor + (`VaeImageProcessor`) + + Inputs: + control_image (`Image`, *optional*): + Control image for ControlNet conditioning. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + control_image_latents (`Tensor`): + The latents representing the control image + """ + + block_classes = [QwenImageControlNetVaeEncoderStep] + block_names = ["controlnet"] + block_trigger_inputs = ["control_image"] + + @property + def description(self): + return ( + "Vae encoder step that encode the image inputs into their latent representations.\n" + + "This is an auto pipeline block.\n" + + " - `QwenImageControlNetVaeEncoderStep` (controlnet) is used when `control_image` is provided.\n" + + " - if `control_image` is not provided, step will be skipped." + ) + + +# ==================== +# 3. DENOISE (input -> prepare_latents -> set_timesteps -> prepare_rope_inputs -> denoise -> after_denoise) +# ==================== + + +# assemble input steps +# auto_docstring +class QwenImageImg2ImgInputStep(SequentialPipelineBlocks): + """ + Input step that prepares the inputs for the img2img denoising step. It: + + Components: + pachifier (`QwenImagePachifier`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + + Outputs: + batch_size (`int`): + The batch size of the prompt embeddings + dtype (`dtype`): + The data type of the prompt embeddings + prompt_embeds (`Tensor`): + The prompt embeddings. (batch-expanded) + prompt_embeds_mask (`Tensor`): + The encoder attention mask. (batch-expanded) + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. (batch-expanded) + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. (batch-expanded) + image_height (`int`): + The image height calculated from the image latents dimension + image_width (`int`): + The image width calculated from the image latents dimension + height (`int`): + if not provided, updated to image height + width (`int`): + if not provided, updated to image width + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. (patchified and + batch-expanded) + """ + + model_name = "qwenimage" + block_classes = [QwenImageTextInputsStep(), QwenImageAdditionalInputsStep()] + block_names = ["text_inputs", "additional_inputs"] + + @property + def description(self): + return "Input step that prepares the inputs for the img2img denoising step. It:\n" + " - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents`).\n" + " - update height/width based `image_latents`, patchify `image_latents`." + + +# auto_docstring +class QwenImageInpaintInputStep(SequentialPipelineBlocks): + """ + Input step that prepares the inputs for the inpainting denoising step. It: + + Components: + pachifier (`QwenImagePachifier`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + image_latents (`Tensor`, *optional*): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`, *optional*): + The processed mask image + + Outputs: + batch_size (`int`): + The batch size of the prompt embeddings + dtype (`dtype`): + The data type of the prompt embeddings + prompt_embeds (`Tensor`): + The prompt embeddings. (batch-expanded) + prompt_embeds_mask (`Tensor`): + The encoder attention mask. (batch-expanded) + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. (batch-expanded) + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. (batch-expanded) + image_height (`int`): + The image height calculated from the image latents dimension + image_width (`int`): + The image width calculated from the image latents dimension + height (`int`): + if not provided, updated to image height + width (`int`): + if not provided, updated to image width + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. (patchified and + batch-expanded) + processed_mask_image (`Tensor`): + The processed mask image (batch-expanded) + """ + + model_name = "qwenimage" + block_classes = [ + QwenImageTextInputsStep(), + QwenImageAdditionalInputsStep( + additional_batch_inputs=[ + InputParam(name="processed_mask_image", type_hint=torch.Tensor, description="The processed mask image") + ] + ), + ] + block_names = ["text_inputs", "additional_inputs"] + + @property + def description(self): + return "Input step that prepares the inputs for the inpainting denoising step. It:\n" + " - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents` and `processed_mask_image`).\n" + " - update height/width based `image_latents`, patchify `image_latents`." + + +# assemble prepare latents steps +# auto_docstring +class QwenImageInpaintPrepareLatentsStep(SequentialPipelineBlocks): + """ + This step prepares the latents/image_latents and mask inputs for the inpainting denoising step. It: + - Add noise to the image latents to create the latents input for the denoiser. + - Create the pachified latents `mask` based on the processedmask image. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) pachifier (`QwenImagePachifier`) + + Inputs: + latents (`Tensor`): + The initial random noised, can be generated in prepare latent step. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. (Can be + generated from vae encoder and updated in input step.) + timesteps (`Tensor`): + The timesteps to use for the denoising process. Can be generated in set_timesteps step. + processed_mask_image (`Tensor`): + The processed mask to use for the inpainting process. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + dtype (`dtype`, *optional*, defaults to torch.float32): + The dtype of the model inputs, can be generated in input step. + + Outputs: + initial_noise (`Tensor`): + The initial random noised used for inpainting denoising. + latents (`Tensor`): + The scaled noisy latents to use for inpainting/image-to-image denoising. + mask (`Tensor`): + The mask to use for the inpainting process. + """ + + model_name = "qwenimage" + block_classes = [QwenImagePrepareLatentsWithStrengthStep(), QwenImageCreateMaskLatentsStep()] + block_names = ["add_noise_to_latents", "create_mask_latents"] + + @property + def description(self) -> str: + return ( + "This step prepares the latents/image_latents and mask inputs for the inpainting denoising step. It:\n" + " - Add noise to the image latents to create the latents input for the denoiser.\n" + " - Create the pachified latents `mask` based on the processedmask image.\n" + ) + + +# assemble denoising steps + + +# Qwen Image (text2image) +# auto_docstring +class QwenImageCoreDenoiseStep(SequentialPipelineBlocks): + """ + step that denoise noise into image for text2image task. It includes the denoise loop, as well as prepare the inputs + (timesteps, latents, rope inputs etc.). + + Components: + pachifier (`QwenImagePachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage" + block_classes = [ + QwenImageTextInputsStep(), + QwenImagePrepareLatentsStep(), + QwenImageSetTimestepsStep(), + QwenImageRoPEInputsStep(), + QwenImageDenoiseStep(), + QwenImageAfterDenoiseStep(), + ] + block_names = [ + "input", + "prepare_latents", + "set_timesteps", + "prepare_rope_inputs", + "denoise", + "after_denoise", + ] + + @property + def description(self): + return "step that denoise noise into image for text2image task. It includes the denoise loop, as well as prepare the inputs (timesteps, latents, rope inputs etc.)." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +# Qwen Image (inpainting) +# auto_docstring +class QwenImageInpaintCoreDenoiseStep(SequentialPipelineBlocks): + """ + Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for inpaint + task. + + Components: + pachifier (`QwenImagePachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + image_latents (`Tensor`, *optional*): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`, *optional*): + The processed mask image + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage" + block_classes = [ + QwenImageInpaintInputStep(), + QwenImagePrepareLatentsStep(), + QwenImageSetTimestepsWithStrengthStep(), + QwenImageInpaintPrepareLatentsStep(), + QwenImageRoPEInputsStep(), + QwenImageInpaintDenoiseStep(), + QwenImageAfterDenoiseStep(), + ] + block_names = [ + "input", + "prepare_latents", + "set_timesteps", + "prepare_inpaint_latents", + "prepare_rope_inputs", + "denoise", + "after_denoise", + ] + + @property + def description(self): + return "Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for inpaint task." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +# Qwen Image (image2image) +# auto_docstring +class QwenImageImg2ImgCoreDenoiseStep(SequentialPipelineBlocks): + """ + Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for img2img + task. + + Components: + pachifier (`QwenImagePachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage" + block_classes = [ + QwenImageImg2ImgInputStep(), + QwenImagePrepareLatentsStep(), + QwenImageSetTimestepsWithStrengthStep(), + QwenImagePrepareLatentsWithStrengthStep(), + QwenImageRoPEInputsStep(), + QwenImageDenoiseStep(), + QwenImageAfterDenoiseStep(), + ] + block_names = [ + "input", + "prepare_latents", + "set_timesteps", + "prepare_img2img_latents", + "prepare_rope_inputs", + "denoise", + "after_denoise", + ] + + @property + def description(self): + return "Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for img2img task." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +# Qwen Image (text2image) with controlnet +# auto_docstring +class QwenImageControlNetCoreDenoiseStep(SequentialPipelineBlocks): + """ + step that denoise noise into image for text2image task. It includes the denoise loop, as well as prepare the inputs + (timesteps, latents, rope inputs etc.). + + Components: + pachifier (`QwenImagePachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) controlnet + (`QwenImageControlNetModel`) guider (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + control_image_latents (`Tensor`): + The control image latents to use for the denoising process. Can be generated in controlnet vae encoder + step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + control_guidance_start (`float`, *optional*, defaults to 0.0): + When to start applying ControlNet. + control_guidance_end (`float`, *optional*, defaults to 1.0): + When to stop applying ControlNet. + controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0): + Scale for ControlNet conditioning. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage" + block_classes = [ + QwenImageTextInputsStep(), + QwenImageControlNetInputsStep(), + QwenImagePrepareLatentsStep(), + QwenImageSetTimestepsStep(), + QwenImageRoPEInputsStep(), + QwenImageControlNetBeforeDenoiserStep(), + QwenImageControlNetDenoiseStep(), + QwenImageAfterDenoiseStep(), + ] + block_names = [ + "input", + "controlnet_input", + "prepare_latents", + "set_timesteps", + "prepare_rope_inputs", + "controlnet_before_denoise", + "controlnet_denoise", + "after_denoise", + ] + + @property + def description(self): + return "step that denoise noise into image for text2image task. It includes the denoise loop, as well as prepare the inputs (timesteps, latents, rope inputs etc.)." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +# Qwen Image (inpainting) with controlnet +# auto_docstring +class QwenImageControlNetInpaintCoreDenoiseStep(SequentialPipelineBlocks): + """ + Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for inpaint + task. + + Components: + pachifier (`QwenImagePachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) controlnet + (`QwenImageControlNetModel`) guider (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + image_latents (`Tensor`, *optional*): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`, *optional*): + The processed mask image + control_image_latents (`Tensor`): + The control image latents to use for the denoising process. Can be generated in controlnet vae encoder + step. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. + control_guidance_start (`float`, *optional*, defaults to 0.0): + When to start applying ControlNet. + control_guidance_end (`float`, *optional*, defaults to 1.0): + When to stop applying ControlNet. + controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0): + Scale for ControlNet conditioning. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage" + block_classes = [ + QwenImageInpaintInputStep(), + QwenImageControlNetInputsStep(), + QwenImagePrepareLatentsStep(), + QwenImageSetTimestepsWithStrengthStep(), + QwenImageInpaintPrepareLatentsStep(), + QwenImageRoPEInputsStep(), + QwenImageControlNetBeforeDenoiserStep(), + QwenImageInpaintControlNetDenoiseStep(), + QwenImageAfterDenoiseStep(), + ] + block_names = [ + "input", + "controlnet_input", + "prepare_latents", + "set_timesteps", + "prepare_inpaint_latents", + "prepare_rope_inputs", + "controlnet_before_denoise", + "controlnet_denoise", + "after_denoise", + ] + + @property + def description(self): + return "Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for inpaint task." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +# Qwen Image (image2image) with controlnet +# auto_docstring +class QwenImageControlNetImg2ImgCoreDenoiseStep(SequentialPipelineBlocks): + """ + Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for img2img + task. + + Components: + pachifier (`QwenImagePachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) controlnet + (`QwenImageControlNetModel`) guider (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + control_image_latents (`Tensor`): + The control image latents to use for the denoising process. Can be generated in controlnet vae encoder + step. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. + control_guidance_start (`float`, *optional*, defaults to 0.0): + When to start applying ControlNet. + control_guidance_end (`float`, *optional*, defaults to 1.0): + When to stop applying ControlNet. + controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0): + Scale for ControlNet conditioning. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage" + block_classes = [ + QwenImageImg2ImgInputStep(), + QwenImageControlNetInputsStep(), + QwenImagePrepareLatentsStep(), + QwenImageSetTimestepsWithStrengthStep(), + QwenImagePrepareLatentsWithStrengthStep(), + QwenImageRoPEInputsStep(), + QwenImageControlNetBeforeDenoiserStep(), + QwenImageControlNetDenoiseStep(), + QwenImageAfterDenoiseStep(), + ] + block_names = [ + "input", + "controlnet_input", + "prepare_latents", + "set_timesteps", + "prepare_img2img_latents", + "prepare_rope_inputs", + "controlnet_before_denoise", + "controlnet_denoise", + "after_denoise", + ] + + @property + def description(self): + return "Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for img2img task." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +# Auto denoise step for QwenImage +class QwenImageAutoCoreDenoiseStep(ConditionalPipelineBlocks): + block_classes = [ + QwenImageCoreDenoiseStep, + QwenImageInpaintCoreDenoiseStep, + QwenImageImg2ImgCoreDenoiseStep, + QwenImageControlNetCoreDenoiseStep, + QwenImageControlNetInpaintCoreDenoiseStep, + QwenImageControlNetImg2ImgCoreDenoiseStep, + ] + block_names = [ + "text2image", + "inpaint", + "img2img", + "controlnet_text2image", + "controlnet_inpaint", + "controlnet_img2img", + ] + block_trigger_inputs = ["control_image_latents", "processed_mask_image", "image_latents"] + default_block_name = "text2image" + + def select_block(self, control_image_latents=None, processed_mask_image=None, image_latents=None): + if control_image_latents is not None: + if processed_mask_image is not None: + return "controlnet_inpaint" + elif image_latents is not None: + return "controlnet_img2img" + else: + return "controlnet_text2image" + else: + if processed_mask_image is not None: + return "inpaint" + elif image_latents is not None: + return "img2img" + else: + return "text2image" + + @property + def description(self): + return ( + "Core step that performs the denoising process. \n" + + " - `QwenImageCoreDenoiseStep` (text2image) for text2image tasks.\n" + + " - `QwenImageInpaintCoreDenoiseStep` (inpaint) for inpaint tasks.\n" + + " - `QwenImageImg2ImgCoreDenoiseStep` (img2img) for img2img tasks.\n" + + " - `QwenImageControlNetCoreDenoiseStep` (controlnet_text2image) for text2image tasks with controlnet.\n" + + " - `QwenImageControlNetInpaintCoreDenoiseStep` (controlnet_inpaint) for inpaint tasks with controlnet.\n" + + " - `QwenImageControlNetImg2ImgCoreDenoiseStep` (controlnet_img2img) for img2img tasks with controlnet.\n" + + "This step support text-to-image, image-to-image, inpainting, and controlnet tasks for QwenImage:\n" + + " - for image-to-image generation, you need to provide `image_latents`\n" + + " - for inpainting, you need to provide `processed_mask_image` and `image_latents`\n" + + " - to run the controlnet workflow, you need to provide `control_image_latents`\n" + + " - for text-to-image generation, all you need to provide is prompt embeddings" + ) + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +# ==================== +# 4. DECODE +# ==================== + + +# standard decode step works for most tasks except for inpaint +# auto_docstring +class QwenImageDecodeStep(SequentialPipelineBlocks): + """ + Decode step that decodes the latents to images and postprocess the generated image. + + Components: + vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`) + + Inputs: + latents (`Tensor`): + The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise + step. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + + Outputs: + images (`list`): + Generated images. (tensor output of the vae decoder.) + """ + + model_name = "qwenimage" + block_classes = [QwenImageDecoderStep(), QwenImageProcessImagesOutputStep()] + block_names = ["decode", "postprocess"] + + @property + def description(self): + return "Decode step that decodes the latents to images and postprocess the generated image." + + +# Inpaint decode step +# auto_docstring +class QwenImageInpaintDecodeStep(SequentialPipelineBlocks): + """ + Decode step that decodes the latents to images and postprocess the generated image, optional apply the mask + overally to the original image. + + Components: + vae (`AutoencoderKLQwenImage`) image_mask_processor (`InpaintProcessor`) + + Inputs: + latents (`Tensor`): + The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise + step. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + mask_overlay_kwargs (`dict`, *optional*): + The kwargs for the postprocess step to apply the mask overlay. generated in + InpaintProcessImagesInputStep. + + Outputs: + images (`list`): + Generated images. (tensor output of the vae decoder.) + """ + + model_name = "qwenimage" + block_classes = [QwenImageDecoderStep(), QwenImageInpaintProcessImagesOutputStep()] + block_names = ["decode", "postprocess"] + + @property + def description(self): + return "Decode step that decodes the latents to images and postprocess the generated image, optional apply the mask overally to the original image." + + +# Auto decode step for QwenImage +class QwenImageAutoDecodeStep(AutoPipelineBlocks): + block_classes = [QwenImageInpaintDecodeStep, QwenImageDecodeStep] + block_names = ["inpaint_decode", "decode"] + block_trigger_inputs = ["mask", None] + + @property + def description(self): + return ( + "Decode step that decode the latents into images. \n" + " This is an auto pipeline block that works for inpaint/text2image/img2img tasks, for both QwenImage and QwenImage-Edit.\n" + + " - `QwenImageInpaintDecodeStep` (inpaint) is used when `mask` is provided.\n" + + " - `QwenImageDecodeStep` (text2image/img2img) is used when `mask` is not provided.\n" + ) + + +# ==================== +# 5. AUTO BLOCKS & PRESETS +# ==================== +AUTO_BLOCKS = InsertableDict( + [ + ("text_encoder", QwenImageAutoTextEncoderStep()), + ("vae_encoder", QwenImageAutoVaeEncoderStep()), + ("controlnet_vae_encoder", QwenImageOptionalControlNetVaeEncoderStep()), + ("denoise", QwenImageAutoCoreDenoiseStep()), + ("decode", QwenImageAutoDecodeStep()), + ] +) + + +# auto_docstring +class QwenImageAutoBlocks(SequentialPipelineBlocks): + """ + Auto Modular pipeline for text-to-image, image-to-image, inpainting, and controlnet tasks using QwenImage. + + Supported workflows: + - `text2image`: requires `prompt` + - `image2image`: requires `prompt`, `image` + - `inpainting`: requires `prompt`, `mask_image`, `image` + - `controlnet_text2image`: requires `prompt`, `control_image` + - `controlnet_image2image`: requires `prompt`, `image`, `control_image` + - `controlnet_inpainting`: requires `prompt`, `mask_image`, `image`, `control_image` + + Components: + text_encoder (`Qwen2_5_VLForConditionalGeneration`): The text encoder to use tokenizer (`Qwen2Tokenizer`): + The tokenizer to use guider (`ClassifierFreeGuidance`) image_mask_processor (`InpaintProcessor`) vae + (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`) controlnet (`QwenImageControlNetModel`) + control_image_processor (`VaeImageProcessor`) pachifier (`QwenImagePachifier`) scheduler + (`FlowMatchEulerDiscreteScheduler`) transformer (`QwenImageTransformer2DModel`) + + Inputs: + prompt (`str`, *optional*): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + max_sequence_length (`int`, *optional*, defaults to 1024): + Maximum sequence length for prompt encoding. + mask_image (`Image`, *optional*): + Mask image for inpainting. + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + padding_mask_crop (`int`, *optional*): + Padding for mask cropping in inpainting. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + control_image (`Image`, *optional*): + Control image for ControlNet conditioning. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + latents (`Tensor`): + Pre-generated noisy latents for image generation. + num_inference_steps (`int`): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + image_latents (`Tensor`, *optional*): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`, *optional*): + The processed mask image + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. + control_image_latents (`Tensor`, *optional*): + The control image latents to use for the denoising process. Can be generated in controlnet vae encoder + step. + control_guidance_start (`float`, *optional*, defaults to 0.0): + When to start applying ControlNet. + control_guidance_end (`float`, *optional*, defaults to 1.0): + When to stop applying ControlNet. + controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0): + Scale for ControlNet conditioning. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + mask_overlay_kwargs (`dict`, *optional*): + The kwargs for the postprocess step to apply the mask overlay. generated in + InpaintProcessImagesInputStep. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "qwenimage" + + block_classes = AUTO_BLOCKS.values() + block_names = AUTO_BLOCKS.keys() + + # Workflow map defines the trigger conditions for each workflow. + # How to define: + # - Only include required inputs and trigger inputs (inputs that determine which blocks run) + # - currently, only supports `True` means the workflow triggers when the input is not None + + _workflow_map = { + "text2image": {"prompt": True}, + "image2image": {"prompt": True, "image": True}, + "inpainting": {"prompt": True, "mask_image": True, "image": True}, + "controlnet_text2image": {"prompt": True, "control_image": True}, + "controlnet_image2image": {"prompt": True, "image": True, "control_image": True}, + "controlnet_inpainting": {"prompt": True, "mask_image": True, "image": True, "control_image": True}, + } + + @property + def description(self): + return "Auto Modular pipeline for text-to-image, image-to-image, inpainting, and controlnet tasks using QwenImage." + + @property + def outputs(self): + return [OutputParam.template("images")] diff --git a/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit.py b/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit.py new file mode 100644 index 000000000000..37b80b69ec7e --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit.py @@ -0,0 +1,796 @@ +# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch + +from ...utils import logging +from ..modular_pipeline import AutoPipelineBlocks, ConditionalPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import InputParam, InsertableDict, OutputParam +from .before_denoise import ( + QwenImageCreateMaskLatentsStep, + QwenImageEditRoPEInputsStep, + QwenImagePrepareLatentsStep, + QwenImagePrepareLatentsWithStrengthStep, + QwenImageSetTimestepsStep, + QwenImageSetTimestepsWithStrengthStep, +) +from .decoders import ( + QwenImageAfterDenoiseStep, + QwenImageDecoderStep, + QwenImageInpaintProcessImagesOutputStep, + QwenImageProcessImagesOutputStep, +) +from .denoise import ( + QwenImageEditDenoiseStep, + QwenImageEditInpaintDenoiseStep, +) +from .encoders import ( + QwenImageEditInpaintProcessImagesInputStep, + QwenImageEditProcessImagesInputStep, + QwenImageEditResizeStep, + QwenImageEditTextEncoderStep, + QwenImageVaeEncoderStep, +) +from .inputs import ( + QwenImageAdditionalInputsStep, + QwenImageTextInputsStep, +) + + +logger = logging.get_logger(__name__) + + +# ==================== +# 1. TEXT ENCODER +# ==================== + + +# auto_docstring +class QwenImageEditVLEncoderStep(SequentialPipelineBlocks): + """ + QwenImage-Edit VL encoder step that encode the image and text prompts together. + + Components: + image_resize_processor (`VaeImageProcessor`) text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor + (`Qwen2VLProcessor`) guider (`ClassifierFreeGuidance`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + + Outputs: + resized_image (`list`): + The resized images + prompt_embeds (`Tensor`): + The prompt embeddings. + prompt_embeds_mask (`Tensor`): + The encoder attention mask. + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. + """ + + model_name = "qwenimage-edit" + block_classes = [ + QwenImageEditResizeStep(), + QwenImageEditTextEncoderStep(), + ] + block_names = ["resize", "encode"] + + @property + def description(self) -> str: + return "QwenImage-Edit VL encoder step that encode the image and text prompts together." + + +# ==================== +# 2. VAE ENCODER +# ==================== + + +# Edit VAE encoder +# auto_docstring +class QwenImageEditVaeEncoderStep(SequentialPipelineBlocks): + """ + Vae encoder step that encode the image inputs into their latent representations. + + Components: + image_resize_processor (`VaeImageProcessor`) image_processor (`VaeImageProcessor`) vae + (`AutoencoderKLQwenImage`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + resized_image (`list`): + The resized images + processed_image (`Tensor`): + The processed image + image_latents (`Tensor`): + The latent representation of the input image. + """ + + model_name = "qwenimage-edit" + block_classes = [ + QwenImageEditResizeStep(), + QwenImageEditProcessImagesInputStep(), + QwenImageVaeEncoderStep(), + ] + block_names = ["resize", "preprocess", "encode"] + + @property + def description(self) -> str: + return "Vae encoder step that encode the image inputs into their latent representations." + + +# Edit Inpaint VAE encoder +# auto_docstring +class QwenImageEditInpaintVaeEncoderStep(SequentialPipelineBlocks): + """ + This step is used for processing image and mask inputs for QwenImage-Edit inpaint tasks. It: + - resize the image for target area (1024 * 1024) while maintaining the aspect ratio. + - process the resized image and mask image. + - create image latents. + + Components: + image_resize_processor (`VaeImageProcessor`) image_mask_processor (`InpaintProcessor`) vae + (`AutoencoderKLQwenImage`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + mask_image (`Image`): + Mask image for inpainting. + padding_mask_crop (`int`, *optional*): + Padding for mask cropping in inpainting. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + resized_image (`list`): + The resized images + processed_image (`Tensor`): + The processed image + processed_mask_image (`Tensor`): + The processed mask image + mask_overlay_kwargs (`dict`): + The kwargs for the postprocess step to apply the mask overlay + image_latents (`Tensor`): + The latent representation of the input image. + """ + + model_name = "qwenimage-edit" + block_classes = [ + QwenImageEditResizeStep(), + QwenImageEditInpaintProcessImagesInputStep(), + QwenImageVaeEncoderStep(), + ] + block_names = ["resize", "preprocess", "encode"] + + @property + def description(self) -> str: + return ( + "This step is used for processing image and mask inputs for QwenImage-Edit inpaint tasks. It:\n" + " - resize the image for target area (1024 * 1024) while maintaining the aspect ratio.\n" + " - process the resized image and mask image.\n" + " - create image latents." + ) + + +# Auto VAE encoder +class QwenImageEditAutoVaeEncoderStep(AutoPipelineBlocks): + block_classes = [QwenImageEditInpaintVaeEncoderStep, QwenImageEditVaeEncoderStep] + block_names = ["edit_inpaint", "edit"] + block_trigger_inputs = ["mask_image", "image"] + + @property + def description(self): + return ( + "Vae encoder step that encode the image inputs into their latent representations.\n" + "This is an auto pipeline block.\n" + " - `QwenImageEditInpaintVaeEncoderStep` (edit_inpaint) is used when `mask_image` is provided.\n" + " - `QwenImageEditVaeEncoderStep` (edit) is used when `image` is provided.\n" + " - if `mask_image` or `image` is not provided, step will be skipped." + ) + + +# ==================== +# 3. DENOISE (input -> prepare_latents -> set_timesteps -> prepare_rope_inputs -> denoise -> after_denoise) +# ==================== + + +# assemble input steps +# auto_docstring +class QwenImageEditInputStep(SequentialPipelineBlocks): + """ + Input step that prepares the inputs for the edit denoising step. It: + - make sure the text embeddings have consistent batch size as well as the additional inputs. + - update height/width based `image_latents`, patchify `image_latents`. + + Components: + pachifier (`QwenImagePachifier`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + + Outputs: + batch_size (`int`): + The batch size of the prompt embeddings + dtype (`dtype`): + The data type of the prompt embeddings + prompt_embeds (`Tensor`): + The prompt embeddings. (batch-expanded) + prompt_embeds_mask (`Tensor`): + The encoder attention mask. (batch-expanded) + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. (batch-expanded) + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. (batch-expanded) + image_height (`int`): + The image height calculated from the image latents dimension + image_width (`int`): + The image width calculated from the image latents dimension + height (`int`): + if not provided, updated to image height + width (`int`): + if not provided, updated to image width + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. (patchified and + batch-expanded) + """ + + model_name = "qwenimage-edit" + block_classes = [ + QwenImageTextInputsStep(), + QwenImageAdditionalInputsStep(), + ] + block_names = ["text_inputs", "additional_inputs"] + + @property + def description(self): + return ( + "Input step that prepares the inputs for the edit denoising step. It:\n" + " - make sure the text embeddings have consistent batch size as well as the additional inputs.\n" + " - update height/width based `image_latents`, patchify `image_latents`." + ) + + +# auto_docstring +class QwenImageEditInpaintInputStep(SequentialPipelineBlocks): + """ + Input step that prepares the inputs for the edit inpaint denoising step. It: + - make sure the text embeddings have consistent batch size as well as the additional inputs. + - update height/width based `image_latents`, patchify `image_latents`. + + Components: + pachifier (`QwenImagePachifier`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`, *optional*): + The processed mask image + + Outputs: + batch_size (`int`): + The batch size of the prompt embeddings + dtype (`dtype`): + The data type of the prompt embeddings + prompt_embeds (`Tensor`): + The prompt embeddings. (batch-expanded) + prompt_embeds_mask (`Tensor`): + The encoder attention mask. (batch-expanded) + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. (batch-expanded) + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. (batch-expanded) + image_height (`int`): + The image height calculated from the image latents dimension + image_width (`int`): + The image width calculated from the image latents dimension + height (`int`): + if not provided, updated to image height + width (`int`): + if not provided, updated to image width + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. (patchified and + batch-expanded) + processed_mask_image (`Tensor`): + The processed mask image (batch-expanded) + """ + + model_name = "qwenimage-edit" + block_classes = [ + QwenImageTextInputsStep(), + QwenImageAdditionalInputsStep( + additional_batch_inputs=[ + InputParam(name="processed_mask_image", type_hint=torch.Tensor, description="The processed mask image") + ] + ), + ] + block_names = ["text_inputs", "additional_inputs"] + + @property + def description(self): + return ( + "Input step that prepares the inputs for the edit inpaint denoising step. It:\n" + " - make sure the text embeddings have consistent batch size as well as the additional inputs.\n" + " - update height/width based `image_latents`, patchify `image_latents`." + ) + + +# assemble prepare latents steps +# auto_docstring +class QwenImageEditInpaintPrepareLatentsStep(SequentialPipelineBlocks): + """ + This step prepares the latents/image_latents and mask inputs for the edit inpainting denoising step. It: + - Add noise to the image latents to create the latents input for the denoiser. + - Create the patchified latents `mask` based on the processed mask image. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) pachifier (`QwenImagePachifier`) + + Inputs: + latents (`Tensor`): + The initial random noised, can be generated in prepare latent step. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. (Can be + generated from vae encoder and updated in input step.) + timesteps (`Tensor`): + The timesteps to use for the denoising process. Can be generated in set_timesteps step. + processed_mask_image (`Tensor`): + The processed mask to use for the inpainting process. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + dtype (`dtype`, *optional*, defaults to torch.float32): + The dtype of the model inputs, can be generated in input step. + + Outputs: + initial_noise (`Tensor`): + The initial random noised used for inpainting denoising. + latents (`Tensor`): + The scaled noisy latents to use for inpainting/image-to-image denoising. + mask (`Tensor`): + The mask to use for the inpainting process. + """ + + model_name = "qwenimage-edit" + block_classes = [QwenImagePrepareLatentsWithStrengthStep(), QwenImageCreateMaskLatentsStep()] + block_names = ["add_noise_to_latents", "create_mask_latents"] + + @property + def description(self) -> str: + return ( + "This step prepares the latents/image_latents and mask inputs for the edit inpainting denoising step. It:\n" + " - Add noise to the image latents to create the latents input for the denoiser.\n" + " - Create the patchified latents `mask` based on the processed mask image.\n" + ) + + +# Qwen Image Edit (image2image) core denoise step +# auto_docstring +class QwenImageEditCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoising workflow for QwenImage-Edit edit (img2img) task. + + Components: + pachifier (`QwenImagePachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage-edit" + block_classes = [ + QwenImageEditInputStep(), + QwenImagePrepareLatentsStep(), + QwenImageSetTimestepsStep(), + QwenImageEditRoPEInputsStep(), + QwenImageEditDenoiseStep(), + QwenImageAfterDenoiseStep(), + ] + block_names = [ + "input", + "prepare_latents", + "set_timesteps", + "prepare_rope_inputs", + "denoise", + "after_denoise", + ] + + @property + def description(self): + return "Core denoising workflow for QwenImage-Edit edit (img2img) task." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +# Qwen Image Edit (inpainting) core denoise step +# auto_docstring +class QwenImageEditInpaintCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoising workflow for QwenImage-Edit edit inpaint task. + + Components: + pachifier (`QwenImagePachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`, *optional*): + The processed mask image + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage-edit" + block_classes = [ + QwenImageEditInpaintInputStep(), + QwenImagePrepareLatentsStep(), + QwenImageSetTimestepsWithStrengthStep(), + QwenImageEditInpaintPrepareLatentsStep(), + QwenImageEditRoPEInputsStep(), + QwenImageEditInpaintDenoiseStep(), + QwenImageAfterDenoiseStep(), + ] + block_names = [ + "input", + "prepare_latents", + "set_timesteps", + "prepare_inpaint_latents", + "prepare_rope_inputs", + "denoise", + "after_denoise", + ] + + @property + def description(self): + return "Core denoising workflow for QwenImage-Edit edit inpaint task." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +# Auto core denoise step for QwenImage Edit +class QwenImageEditAutoCoreDenoiseStep(ConditionalPipelineBlocks): + model_name = "qwenimage-edit" + block_classes = [ + QwenImageEditInpaintCoreDenoiseStep, + QwenImageEditCoreDenoiseStep, + ] + block_names = ["edit_inpaint", "edit"] + block_trigger_inputs = ["processed_mask_image", "image_latents"] + default_block_name = "edit" + + def select_block(self, processed_mask_image=None, image_latents=None) -> str | None: + if processed_mask_image is not None: + return "edit_inpaint" + elif image_latents is not None: + return "edit" + return None + + @property + def description(self): + return ( + "Auto core denoising step that selects the appropriate workflow based on inputs.\n" + " - `QwenImageEditInpaintCoreDenoiseStep` when `processed_mask_image` is provided\n" + " - `QwenImageEditCoreDenoiseStep` when `image_latents` is provided\n" + "Supports edit (img2img) and edit inpainting tasks for QwenImage-Edit." + ) + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +# ==================== +# 4. DECODE +# ==================== + + +# Decode step (standard) +# auto_docstring +class QwenImageEditDecodeStep(SequentialPipelineBlocks): + """ + Decode step that decodes the latents to images and postprocess the generated image. + + Components: + vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`) + + Inputs: + latents (`Tensor`): + The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise + step. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + + Outputs: + images (`list`): + Generated images. (tensor output of the vae decoder.) + """ + + model_name = "qwenimage-edit" + block_classes = [QwenImageDecoderStep(), QwenImageProcessImagesOutputStep()] + block_names = ["decode", "postprocess"] + + @property + def description(self): + return "Decode step that decodes the latents to images and postprocess the generated image." + + +# Inpaint decode step +# auto_docstring +class QwenImageEditInpaintDecodeStep(SequentialPipelineBlocks): + """ + Decode step that decodes the latents to images and postprocess the generated image, optionally apply the mask + overlay to the original image. + + Components: + vae (`AutoencoderKLQwenImage`) image_mask_processor (`InpaintProcessor`) + + Inputs: + latents (`Tensor`): + The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise + step. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + mask_overlay_kwargs (`dict`, *optional*): + The kwargs for the postprocess step to apply the mask overlay. generated in + InpaintProcessImagesInputStep. + + Outputs: + images (`list`): + Generated images. (tensor output of the vae decoder.) + """ + + model_name = "qwenimage-edit" + block_classes = [QwenImageDecoderStep(), QwenImageInpaintProcessImagesOutputStep()] + block_names = ["decode", "postprocess"] + + @property + def description(self): + return "Decode step that decodes the latents to images and postprocess the generated image, optionally apply the mask overlay to the original image." + + +# Auto decode step +class QwenImageEditAutoDecodeStep(AutoPipelineBlocks): + block_classes = [QwenImageEditInpaintDecodeStep, QwenImageEditDecodeStep] + block_names = ["inpaint_decode", "decode"] + block_trigger_inputs = ["mask", None] + + @property + def description(self): + return ( + "Decode step that decode the latents into images.\n" + "This is an auto pipeline block.\n" + " - `QwenImageEditInpaintDecodeStep` (inpaint) is used when `mask` is provided.\n" + " - `QwenImageEditDecodeStep` (edit) is used when `mask` is not provided.\n" + ) + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +# ==================== +# 5. AUTO BLOCKS & PRESETS +# ==================== + +EDIT_AUTO_BLOCKS = InsertableDict( + [ + ("text_encoder", QwenImageEditVLEncoderStep()), + ("vae_encoder", QwenImageEditAutoVaeEncoderStep()), + ("denoise", QwenImageEditAutoCoreDenoiseStep()), + ("decode", QwenImageEditAutoDecodeStep()), + ] +) + + +# auto_docstring +class QwenImageEditAutoBlocks(SequentialPipelineBlocks): + """ + Auto Modular pipeline for edit (img2img) and edit inpaint tasks using QwenImage-Edit. + - for edit (img2img) generation, you need to provide `image` + - for edit inpainting, you need to provide `mask_image` and `image`, optionally you can provide + `padding_mask_crop` + + + Supported workflows: + - `image_conditioned`: requires `prompt`, `image` + - `image_conditioned_inpainting`: requires `prompt`, `mask_image`, `image` + + Components: + image_resize_processor (`VaeImageProcessor`) text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor + (`Qwen2VLProcessor`) guider (`ClassifierFreeGuidance`) image_mask_processor (`InpaintProcessor`) vae + (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`) pachifier (`QwenImagePachifier`) scheduler + (`FlowMatchEulerDiscreteScheduler`) transformer (`QwenImageTransformer2DModel`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + mask_image (`Image`, *optional*): + Mask image for inpainting. + padding_mask_crop (`int`, *optional*): + Padding for mask cropping in inpainting. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`, *optional*): + The processed mask image + latents (`Tensor`): + Pre-generated noisy latents for image generation. + num_inference_steps (`int`): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + mask_overlay_kwargs (`dict`, *optional*): + The kwargs for the postprocess step to apply the mask overlay. generated in + InpaintProcessImagesInputStep. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "qwenimage-edit" + block_classes = EDIT_AUTO_BLOCKS.values() + block_names = EDIT_AUTO_BLOCKS.keys() + _workflow_map = { + "image_conditioned": {"prompt": True, "image": True}, + "image_conditioned_inpainting": {"prompt": True, "mask_image": True, "image": True}, + } + + @property + def description(self): + return ( + "Auto Modular pipeline for edit (img2img) and edit inpaint tasks using QwenImage-Edit.\n" + "- for edit (img2img) generation, you need to provide `image`\n" + "- for edit inpainting, you need to provide `mask_image` and `image`, optionally you can provide `padding_mask_crop`\n" + ) + + @property + def outputs(self): + return [OutputParam.template("images")] diff --git a/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit_plus.py b/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit_plus.py new file mode 100644 index 000000000000..4a1f418d7b45 --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit_plus.py @@ -0,0 +1,407 @@ +# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...utils import logging +from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import ( + QwenImageEditPlusRoPEInputsStep, + QwenImagePrepareLatentsStep, + QwenImageSetTimestepsStep, +) +from .decoders import ( + QwenImageAfterDenoiseStep, + QwenImageDecoderStep, + QwenImageProcessImagesOutputStep, +) +from .denoise import ( + QwenImageEditDenoiseStep, +) +from .encoders import ( + QwenImageEditPlusProcessImagesInputStep, + QwenImageEditPlusResizeStep, + QwenImageEditPlusTextEncoderStep, + QwenImageVaeEncoderStep, +) +from .inputs import ( + QwenImageEditPlusAdditionalInputsStep, + QwenImageTextInputsStep, +) + + +logger = logging.get_logger(__name__) + + +# ==================== +# 1. TEXT ENCODER +# ==================== + + +# auto_docstring +class QwenImageEditPlusVLEncoderStep(SequentialPipelineBlocks): + """ + QwenImage-Edit Plus VL encoder step that encodes the image and text prompts together. + + Components: + image_resize_processor (`VaeImageProcessor`) text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor + (`Qwen2VLProcessor`) guider (`ClassifierFreeGuidance`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + + Outputs: + resized_image (`list`): + Images resized to 1024x1024 target area for VAE encoding + resized_cond_image (`list`): + Images resized to 384x384 target area for VL text encoding + prompt_embeds (`Tensor`): + The prompt embeddings. + prompt_embeds_mask (`Tensor`): + The encoder attention mask. + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. + """ + + model_name = "qwenimage-edit-plus" + block_classes = [ + QwenImageEditPlusResizeStep(), + QwenImageEditPlusTextEncoderStep(), + ] + block_names = ["resize", "encode"] + + @property + def description(self) -> str: + return "QwenImage-Edit Plus VL encoder step that encodes the image and text prompts together." + + +# ==================== +# 2. VAE ENCODER +# ==================== + + +# auto_docstring +class QwenImageEditPlusVaeEncoderStep(SequentialPipelineBlocks): + """ + VAE encoder step that encodes image inputs into latent representations. + Each image is resized independently based on its own aspect ratio to 1024x1024 target area. + + Components: + image_resize_processor (`VaeImageProcessor`) image_processor (`VaeImageProcessor`) vae + (`AutoencoderKLQwenImage`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + resized_image (`list`): + Images resized to 1024x1024 target area for VAE encoding + resized_cond_image (`list`): + Images resized to 384x384 target area for VL text encoding + processed_image (`Tensor`): + The processed image + image_latents (`Tensor`): + The latent representation of the input image. + """ + + model_name = "qwenimage-edit-plus" + block_classes = [ + QwenImageEditPlusResizeStep(), + QwenImageEditPlusProcessImagesInputStep(), + QwenImageVaeEncoderStep(), + ] + block_names = ["resize", "preprocess", "encode"] + + @property + def description(self) -> str: + return ( + "VAE encoder step that encodes image inputs into latent representations.\n" + "Each image is resized independently based on its own aspect ratio to 1024x1024 target area." + ) + + +# ==================== +# 3. DENOISE (input -> prepare_latents -> set_timesteps -> prepare_rope_inputs -> denoise -> after_denoise) +# ==================== + + +# assemble input steps +# auto_docstring +class QwenImageEditPlusInputStep(SequentialPipelineBlocks): + """ + Input step that prepares the inputs for the Edit Plus denoising step. It: + - Standardizes text embeddings batch size. + - Processes list of image latents: patchifies, concatenates along dim=1, expands batch. + - Outputs lists of image_height/image_width for RoPE calculation. + - Defaults height/width from last image in the list. + + Components: + pachifier (`QwenImagePachifier`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + + Outputs: + batch_size (`int`): + The batch size of the prompt embeddings + dtype (`dtype`): + The data type of the prompt embeddings + prompt_embeds (`Tensor`): + The prompt embeddings. (batch-expanded) + prompt_embeds_mask (`Tensor`): + The encoder attention mask. (batch-expanded) + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. (batch-expanded) + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. (batch-expanded) + image_height (`list`): + The image heights calculated from the image latents dimension + image_width (`list`): + The image widths calculated from the image latents dimension + height (`int`): + if not provided, updated to image height + width (`int`): + if not provided, updated to image width + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. (patchified, + concatenated, and batch-expanded) + """ + + model_name = "qwenimage-edit-plus" + block_classes = [ + QwenImageTextInputsStep(), + QwenImageEditPlusAdditionalInputsStep(), + ] + block_names = ["text_inputs", "additional_inputs"] + + @property + def description(self): + return ( + "Input step that prepares the inputs for the Edit Plus denoising step. It:\n" + " - Standardizes text embeddings batch size.\n" + " - Processes list of image latents: patchifies, concatenates along dim=1, expands batch.\n" + " - Outputs lists of image_height/image_width for RoPE calculation.\n" + " - Defaults height/width from last image in the list." + ) + + +# Qwen Image Edit Plus (image2image) core denoise step +# auto_docstring +class QwenImageEditPlusCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoising workflow for QwenImage-Edit Plus edit (img2img) task. + + Components: + pachifier (`QwenImagePachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage-edit-plus" + block_classes = [ + QwenImageEditPlusInputStep(), + QwenImagePrepareLatentsStep(), + QwenImageSetTimestepsStep(), + QwenImageEditPlusRoPEInputsStep(), + QwenImageEditDenoiseStep(), + QwenImageAfterDenoiseStep(), + ] + block_names = [ + "input", + "prepare_latents", + "set_timesteps", + "prepare_rope_inputs", + "denoise", + "after_denoise", + ] + + @property + def description(self): + return "Core denoising workflow for QwenImage-Edit Plus edit (img2img) task." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +# ==================== +# 4. DECODE +# ==================== + + +# auto_docstring +class QwenImageEditPlusDecodeStep(SequentialPipelineBlocks): + """ + Decode step that decodes the latents to images and postprocesses the generated image. + + Components: + vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`) + + Inputs: + latents (`Tensor`): + The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise + step. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + + Outputs: + images (`list`): + Generated images. (tensor output of the vae decoder.) + """ + + model_name = "qwenimage-edit-plus" + block_classes = [QwenImageDecoderStep(), QwenImageProcessImagesOutputStep()] + block_names = ["decode", "postprocess"] + + @property + def description(self): + return "Decode step that decodes the latents to images and postprocesses the generated image." + + +# ==================== +# 5. AUTO BLOCKS & PRESETS +# ==================== + +EDIT_PLUS_AUTO_BLOCKS = InsertableDict( + [ + ("text_encoder", QwenImageEditPlusVLEncoderStep()), + ("vae_encoder", QwenImageEditPlusVaeEncoderStep()), + ("denoise", QwenImageEditPlusCoreDenoiseStep()), + ("decode", QwenImageEditPlusDecodeStep()), + ] +) + + +# auto_docstring +class QwenImageEditPlusAutoBlocks(SequentialPipelineBlocks): + """ + Auto Modular pipeline for edit (img2img) tasks using QwenImage-Edit Plus. + - `image` is required input (can be single image or list of images). + - Each image is resized independently based on its own aspect ratio. + - VL encoder uses 384x384 target area, VAE encoder uses 1024x1024 target area. + + Components: + image_resize_processor (`VaeImageProcessor`) text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor + (`Qwen2VLProcessor`) guider (`ClassifierFreeGuidance`) image_processor (`VaeImageProcessor`) vae + (`AutoencoderKLQwenImage`) pachifier (`QwenImagePachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) + transformer (`QwenImageTransformer2DModel`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "qwenimage-edit-plus" + block_classes = EDIT_PLUS_AUTO_BLOCKS.values() + block_names = EDIT_PLUS_AUTO_BLOCKS.keys() + + @property + def description(self): + return ( + "Auto Modular pipeline for edit (img2img) tasks using QwenImage-Edit Plus.\n" + "- `image` is required input (can be single image or list of images).\n" + "- Each image is resized independently based on its own aspect ratio.\n" + "- VL encoder uses 384x384 target area, VAE encoder uses 1024x1024 target area." + ) + + @property + def outputs(self): + return [OutputParam.template("images")] diff --git a/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_layered.py b/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_layered.py new file mode 100644 index 000000000000..a10454f1fb0c --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_layered.py @@ -0,0 +1,366 @@ +# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...utils import logging +from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import ( + QwenImageLayeredPrepareLatentsStep, + QwenImageLayeredRoPEInputsStep, + QwenImageLayeredSetTimestepsStep, +) +from .decoders import ( + QwenImageLayeredAfterDenoiseStep, + QwenImageLayeredDecoderStep, +) +from .denoise import ( + QwenImageLayeredDenoiseStep, +) +from .encoders import ( + QwenImageEditProcessImagesInputStep, + QwenImageLayeredGetImagePromptStep, + QwenImageLayeredPermuteLatentsStep, + QwenImageLayeredResizeStep, + QwenImageTextEncoderStep, + QwenImageVaeEncoderStep, +) +from .inputs import ( + QwenImageLayeredAdditionalInputsStep, + QwenImageTextInputsStep, +) + + +logger = logging.get_logger(__name__) + + +# ==================== +# 1. TEXT ENCODER +# ==================== + + +# auto_docstring +class QwenImageLayeredTextEncoderStep(SequentialPipelineBlocks): + """ + QwenImage-Layered Text encoder step that encode the text prompt, will generate a prompt based on image if not + provided. + + Components: + image_resize_processor (`VaeImageProcessor`) text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor + (`Qwen2VLProcessor`) tokenizer (`Qwen2Tokenizer`): The tokenizer to use guider (`ClassifierFreeGuidance`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + resolution (`int`, *optional*, defaults to 640): + The target area to resize the image to, can be 1024 or 640 + prompt (`str`, *optional*): + The prompt or prompts to guide image generation. + use_en_prompt (`bool`, *optional*, defaults to False): + Whether to use English prompt template + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + max_sequence_length (`int`, *optional*, defaults to 1024): + Maximum sequence length for prompt encoding. + + Outputs: + resized_image (`list`): + The resized images + prompt (`str`): + The prompt or prompts to guide image generation. If not provided, updated using image caption + prompt_embeds (`Tensor`): + The prompt embeddings. + prompt_embeds_mask (`Tensor`): + The encoder attention mask. + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. + """ + + model_name = "qwenimage-layered" + block_classes = [ + QwenImageLayeredResizeStep(), + QwenImageLayeredGetImagePromptStep(), + QwenImageTextEncoderStep(), + ] + block_names = ["resize", "get_image_prompt", "encode"] + + @property + def description(self) -> str: + return "QwenImage-Layered Text encoder step that encode the text prompt, will generate a prompt based on image if not provided." + + +# ==================== +# 2. VAE ENCODER +# ==================== + + +# Edit VAE encoder +# auto_docstring +class QwenImageLayeredVaeEncoderStep(SequentialPipelineBlocks): + """ + Vae encoder step that encode the image inputs into their latent representations. + + Components: + image_resize_processor (`VaeImageProcessor`) image_processor (`VaeImageProcessor`) vae + (`AutoencoderKLQwenImage`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + resolution (`int`, *optional*, defaults to 640): + The target area to resize the image to, can be 1024 or 640 + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + resized_image (`list`): + The resized images + processed_image (`Tensor`): + The processed image + image_latents (`Tensor`): + The latent representation of the input image. + """ + + model_name = "qwenimage-layered" + block_classes = [ + QwenImageLayeredResizeStep(), + QwenImageEditProcessImagesInputStep(), + QwenImageVaeEncoderStep(), + QwenImageLayeredPermuteLatentsStep(), + ] + block_names = ["resize", "preprocess", "encode", "permute"] + + @property + def description(self) -> str: + return "Vae encoder step that encode the image inputs into their latent representations." + + +# ==================== +# 3. DENOISE (input -> prepare_latents -> set_timesteps -> prepare_rope_inputs -> denoise -> after_denoise) +# ==================== + + +# assemble input steps +# auto_docstring +class QwenImageLayeredInputStep(SequentialPipelineBlocks): + """ + Input step that prepares the inputs for the layered denoising step. It: + - make sure the text embeddings have consistent batch size as well as the additional inputs. + - update height/width based `image_latents`, patchify `image_latents`. + + Components: + pachifier (`QwenImageLayeredPachifier`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + + Outputs: + batch_size (`int`): + The batch size of the prompt embeddings + dtype (`dtype`): + The data type of the prompt embeddings + prompt_embeds (`Tensor`): + The prompt embeddings. (batch-expanded) + prompt_embeds_mask (`Tensor`): + The encoder attention mask. (batch-expanded) + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. (batch-expanded) + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. (batch-expanded) + image_height (`int`): + The image height calculated from the image latents dimension + image_width (`int`): + The image width calculated from the image latents dimension + height (`int`): + if not provided, updated to image height + width (`int`): + if not provided, updated to image width + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. (patchified + with layered pachifier and batch-expanded) + """ + + model_name = "qwenimage-layered" + block_classes = [ + QwenImageTextInputsStep(), + QwenImageLayeredAdditionalInputsStep(), + ] + block_names = ["text_inputs", "additional_inputs"] + + @property + def description(self): + return ( + "Input step that prepares the inputs for the layered denoising step. It:\n" + " - make sure the text embeddings have consistent batch size as well as the additional inputs.\n" + " - update height/width based `image_latents`, patchify `image_latents`." + ) + + +# Qwen Image Layered (image2image) core denoise step +# auto_docstring +class QwenImageLayeredCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoising workflow for QwenImage-Layered img2img task. + + Components: + pachifier (`QwenImageLayeredPachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_embeds_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + layers (`int`, *optional*, defaults to 4): + Number of layers to extract from the image + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "qwenimage-layered" + block_classes = [ + QwenImageLayeredInputStep(), + QwenImageLayeredPrepareLatentsStep(), + QwenImageLayeredSetTimestepsStep(), + QwenImageLayeredRoPEInputsStep(), + QwenImageLayeredDenoiseStep(), + QwenImageLayeredAfterDenoiseStep(), + ] + block_names = [ + "input", + "prepare_latents", + "set_timesteps", + "prepare_rope_inputs", + "denoise", + "after_denoise", + ] + + @property + def description(self): + return "Core denoising workflow for QwenImage-Layered img2img task." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +# ==================== +# 4. AUTO BLOCKS & PRESETS +# ==================== + +LAYERED_AUTO_BLOCKS = InsertableDict( + [ + ("text_encoder", QwenImageLayeredTextEncoderStep()), + ("vae_encoder", QwenImageLayeredVaeEncoderStep()), + ("denoise", QwenImageLayeredCoreDenoiseStep()), + ("decode", QwenImageLayeredDecoderStep()), + ] +) + + +# auto_docstring +class QwenImageLayeredAutoBlocks(SequentialPipelineBlocks): + """ + Auto Modular pipeline for layered denoising tasks using QwenImage-Layered. + + Components: + image_resize_processor (`VaeImageProcessor`) text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor + (`Qwen2VLProcessor`) tokenizer (`Qwen2Tokenizer`): The tokenizer to use guider (`ClassifierFreeGuidance`) + image_processor (`VaeImageProcessor`) vae (`AutoencoderKLQwenImage`) pachifier (`QwenImageLayeredPachifier`) + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`QwenImageTransformer2DModel`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + resolution (`int`, *optional*, defaults to 640): + The target area to resize the image to, can be 1024 or 640 + prompt (`str`, *optional*): + The prompt or prompts to guide image generation. + use_en_prompt (`bool`, *optional*, defaults to False): + Whether to use English prompt template + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + max_sequence_length (`int`, *optional*, defaults to 1024): + Maximum sequence length for prompt encoding. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + layers (`int`, *optional*, defaults to 4): + Number of layers to extract from the image + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "qwenimage-layered" + block_classes = LAYERED_AUTO_BLOCKS.values() + block_names = LAYERED_AUTO_BLOCKS.keys() + + @property + def description(self): + return "Auto Modular pipeline for layered denoising tasks using QwenImage-Layered." + + @property + def outputs(self): + return [OutputParam.template("images")] diff --git a/src/diffusers/modular_pipelines/qwenimage/modular_pipeline.py b/src/diffusers/modular_pipelines/qwenimage/modular_pipeline.py index 59e1a13a5db2..892435989d00 100644 --- a/src/diffusers/modular_pipelines/qwenimage/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/qwenimage/modular_pipeline.py @@ -90,6 +90,88 @@ def unpack_latents(self, latents, height, width, vae_scale_factor=8): return latents +class QwenImageLayeredPachifier(ConfigMixin): + """ + A class to pack and unpack latents for QwenImage Layered. + + Unlike QwenImagePachifier, this handles 5D latents with shape (B, layers+1, C, H, W). + """ + + config_name = "config.json" + + @register_to_config + def __init__(self, patch_size: int = 2): + super().__init__() + + def pack_latents(self, latents): + """ + Pack latents from (B, layers, C, H, W) to (B, layers * H/2 * W/2, C*4). + """ + + if latents.ndim != 5: + raise ValueError(f"Latents must have 5 dimensions (B, layers, C, H, W), but got {latents.ndim}") + + batch_size, layers, num_channels_latents, latent_height, latent_width = latents.shape + patch_size = self.config.patch_size + + if latent_height % patch_size != 0 or latent_width % patch_size != 0: + raise ValueError( + f"Latent height and width must be divisible by {patch_size}, but got {latent_height} and {latent_width}" + ) + + latents = latents.view( + batch_size, + layers, + num_channels_latents, + latent_height // patch_size, + patch_size, + latent_width // patch_size, + patch_size, + ) + latents = latents.permute(0, 1, 3, 5, 2, 4, 6) + latents = latents.reshape( + batch_size, + layers * (latent_height // patch_size) * (latent_width // patch_size), + num_channels_latents * patch_size * patch_size, + ) + return latents + + def unpack_latents(self, latents, height, width, layers, vae_scale_factor=8): + """ + Unpack latents from (B, seq, C*4) to (B, C, layers+1, H, W). + """ + + if latents.ndim != 3: + raise ValueError(f"Latents must have 3 dimensions, but got {latents.ndim}") + + batch_size, _, channels = latents.shape + patch_size = self.config.patch_size + + height = patch_size * (int(height) // (vae_scale_factor * patch_size)) + width = patch_size * (int(width) // (vae_scale_factor * patch_size)) + + latents = latents.view( + batch_size, + layers + 1, + height // patch_size, + width // patch_size, + channels // (patch_size * patch_size), + patch_size, + patch_size, + ) + latents = latents.permute(0, 1, 4, 2, 5, 3, 6) + latents = latents.reshape( + batch_size, + layers + 1, + channels // (patch_size * patch_size), + height, + width, + ) + latents = latents.permute(0, 2, 1, 3, 4) # (b, c, f, h, w) + + return latents + + class QwenImageModularPipeline(ModularPipeline, QwenImageLoraLoaderMixin): """ A ModularPipeline for QwenImage. @@ -203,3 +285,13 @@ class QwenImageEditPlusModularPipeline(QwenImageEditModularPipeline): """ default_blocks_name = "QwenImageEditPlusAutoBlocks" + + +class QwenImageLayeredModularPipeline(QwenImageModularPipeline): + """ + A ModularPipeline for QwenImage-Layered. + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "QwenImageLayeredAutoBlocks" diff --git a/src/diffusers/modular_pipelines/qwenimage/prompt_templates.py b/src/diffusers/modular_pipelines/qwenimage/prompt_templates.py new file mode 100644 index 000000000000..8e7beb555760 --- /dev/null +++ b/src/diffusers/modular_pipelines/qwenimage/prompt_templates.py @@ -0,0 +1,121 @@ +# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Prompt templates for QwenImage pipelines. + +This module centralizes all prompt templates used across different QwenImage pipeline variants: +- QwenImage (base): Text-only encoding for text-to-image generation +- QwenImage Edit: VL encoding with single image for image editing +- QwenImage Edit Plus: VL encoding with multiple images for multi-reference editing +- QwenImage Layered: Auto-captioning for image decomposition +""" + +# ============================================ +# QwenImage Base (text-only encoding) +# ============================================ +# Used for text-to-image generation where only text prompt is encoded + +QWENIMAGE_PROMPT_TEMPLATE = ( + "<|im_start|>system\n" + "Describe the image by detailing the color, shape, size, texture, quantity, text, " + "spatial relationships of the objects and background:<|im_end|>\n" + "<|im_start|>user\n{}<|im_end|>\n" + "<|im_start|>assistant\n" +) +QWENIMAGE_PROMPT_TEMPLATE_START_IDX = 34 + + +# ============================================ +# QwenImage Edit (VL encoding with single image) +# ============================================ +# Used for single-image editing where both image and text are encoded together + +QWENIMAGE_EDIT_PROMPT_TEMPLATE = ( + "<|im_start|>system\n" + "Describe the key features of the input image (color, shape, size, texture, objects, background), " + "then explain how the user's text instruction should alter or modify the image. " + "Generate a new image that meets the user's requirements while maintaining consistency " + "with the original input where appropriate.<|im_end|>\n" + "<|im_start|>user\n" + "<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n" + "<|im_start|>assistant\n" +) +QWENIMAGE_EDIT_PROMPT_TEMPLATE_START_IDX = 64 + + +# ============================================ +# QwenImage Edit Plus (VL encoding with multiple images) +# ============================================ +# Used for multi-reference editing where multiple images and text are encoded together +# The img_template is used to format each image in the prompt + +QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE = ( + "<|im_start|>system\n" + "Describe the key features of the input image (color, shape, size, texture, objects, background), " + "then explain how the user's text instruction should alter or modify the image. " + "Generate a new image that meets the user's requirements while maintaining consistency " + "with the original input where appropriate.<|im_end|>\n" + "<|im_start|>user\n{}<|im_end|>\n" + "<|im_start|>assistant\n" +) +QWENIMAGE_EDIT_PLUS_IMG_TEMPLATE = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>" +QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE_START_IDX = 64 + + +# ============================================ +# QwenImage Layered (auto-captioning) +# ============================================ +# Used for image decomposition where the VL model generates a caption from the input image +# if no prompt is provided. These prompts instruct the model to describe the image in detail. + +QWENIMAGE_LAYERED_CAPTION_PROMPT_EN = ( + "<|im_start|>system\n" + "You are a helpful assistant.<|im_end|>\n" + "<|im_start|>user\n" + "# Image Annotator\n" + "You are a professional image annotator. Please write an image caption based on the input image:\n" + "1. Write the caption using natural, descriptive language without structured formats or rich text.\n" + "2. Enrich caption details by including:\n" + " - Object attributes, such as quantity, color, shape, size, material, state, position, actions, and so on\n" + " - Vision Relations between objects, such as spatial relations, functional relations, possessive relations, " + "attachment relations, action relations, comparative relations, causal relations, and so on\n" + " - Environmental details, such as weather, lighting, colors, textures, atmosphere, and so on\n" + " - Identify the text clearly visible in the image, without translation or explanation, " + "and highlight it in the caption with quotation marks\n" + "3. Maintain authenticity and accuracy:\n" + " - Avoid generalizations\n" + " - Describe all visible information in the image, while do not add information not explicitly shown in the image\n" + "<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n" + "<|im_start|>assistant\n" +) + +QWENIMAGE_LAYERED_CAPTION_PROMPT_CN = ( + "<|im_start|>system\n" + "You are a helpful assistant.<|im_end|>\n" + "<|im_start|>user\n" + "# 图像标注器\n" + "你是一个专业的图像标注器。请基于输入图像,撰写图注:\n" + "1. 使用自然、描述性的语言撰写图注,不要使用结构化形式或富文本形式。\n" + "2. 通过加入以下内容,丰富图注细节:\n" + " - 对象的属性:如数量、颜色、形状、大小、位置、材质、状态、动作等\n" + " - 对象间的视觉关系:如空间关系、功能关系、动作关系、从属关系、比较关系、因果关系等\n" + " - 环境细节:例如天气、光照、颜色、纹理、气氛等\n" + " - 文字内容:识别图像中清晰可见的文字,不做翻译和解释,用引号在图注中强调\n" + "3. 保持真实性与准确性:\n" + " - 不要使用笼统的描述\n" + " - 描述图像中所有可见的信息,但不要加入没有在图像中出现的内容\n" + "<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n" + "<|im_start|>assistant\n" +) diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/__init__.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/__init__.py index 59ec46dc6d36..44f1c555cef3 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_xl/__init__.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/__init__.py @@ -21,21 +21,7 @@ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) else: - _import_structure["encoders"] = ["StableDiffusionXLTextEncoderStep"] - _import_structure["modular_blocks"] = [ - "ALL_BLOCKS", - "AUTO_BLOCKS", - "CONTROLNET_BLOCKS", - "IMAGE2IMAGE_BLOCKS", - "INPAINT_BLOCKS", - "IP_ADAPTER_BLOCKS", - "TEXT2IMAGE_BLOCKS", - "StableDiffusionXLAutoBlocks", - "StableDiffusionXLAutoControlnetStep", - "StableDiffusionXLAutoDecodeStep", - "StableDiffusionXLAutoIPAdapterStep", - "StableDiffusionXLAutoVaeEncoderStep", - ] + _import_structure["modular_blocks_stable_diffusion_xl"] = ["StableDiffusionXLAutoBlocks"] _import_structure["modular_pipeline"] = ["StableDiffusionXLModularPipeline"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: @@ -45,23 +31,7 @@ except OptionalDependencyNotAvailable: from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 else: - from .encoders import ( - StableDiffusionXLTextEncoderStep, - ) - from .modular_blocks import ( - ALL_BLOCKS, - AUTO_BLOCKS, - CONTROLNET_BLOCKS, - IMAGE2IMAGE_BLOCKS, - INPAINT_BLOCKS, - IP_ADAPTER_BLOCKS, - TEXT2IMAGE_BLOCKS, - StableDiffusionXLAutoBlocks, - StableDiffusionXLAutoControlnetStep, - StableDiffusionXLAutoDecodeStep, - StableDiffusionXLAutoIPAdapterStep, - StableDiffusionXLAutoVaeEncoderStep, - ) + from .modular_blocks_stable_diffusion_xl import StableDiffusionXLAutoBlocks from .modular_pipeline import StableDiffusionXLModularPipeline else: import sys diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/before_denoise.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/before_denoise.py index 70cbf0c1c78d..4a393e7ce296 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_xl/before_denoise.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/before_denoise.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, List, Optional, Tuple, Union +from typing import Any import PIL import torch @@ -46,10 +46,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -64,15 +64,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -105,7 +105,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -210,7 +210,7 @@ def description(self) -> str: ) @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("num_images_per_prompt", default=1), InputParam( @@ -236,18 +236,18 @@ def inputs(self) -> List[InputParam]: ), InputParam( "ip_adapter_embeds", - type_hint=List[torch.Tensor], + type_hint=list[torch.Tensor], description="Pre-generated image embeddings for IP-Adapter. Can be generated from ip_adapter step.", ), InputParam( "negative_ip_adapter_embeds", - type_hint=List[torch.Tensor], + type_hint=list[torch.Tensor], description="Pre-generated negative image embeddings for IP-Adapter. Can be generated from ip_adapter step.", ), ] @property - def intermediate_outputs(self) -> List[str]: + def intermediate_outputs(self) -> list[str]: return [ OutputParam( "batch_size", @@ -285,13 +285,13 @@ def intermediate_outputs(self) -> List[str]: ), OutputParam( "ip_adapter_embeds", - type_hint=List[torch.Tensor], + type_hint=list[torch.Tensor], kwargs_type="denoiser_input_fields", # already in intermedites state but declare here again for denoiser_input_fields description="image embeddings for IP-Adapter", ), OutputParam( "negative_ip_adapter_embeds", - type_hint=List[torch.Tensor], + type_hint=list[torch.Tensor], kwargs_type="denoiser_input_fields", # already in intermedites state but declare here again for denoiser_input_fields description="negative image embeddings for IP-Adapter", ), @@ -393,7 +393,7 @@ class StableDiffusionXLImg2ImgSetTimestepsStep(ModularPipelineBlocks): model_name = "stable-diffusion-xl" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", EulerDiscreteScheduler), ] @@ -406,7 +406,7 @@ def description(self) -> str: ) @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("num_inference_steps", default=50), InputParam("timesteps"), @@ -425,7 +425,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[str]: + def intermediate_outputs(self) -> list[str]: return [ OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference"), OutputParam( @@ -537,7 +537,7 @@ class StableDiffusionXLSetTimestepsStep(ModularPipelineBlocks): model_name = "stable-diffusion-xl" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", EulerDiscreteScheduler), ] @@ -547,7 +547,7 @@ def description(self) -> str: return "Step that sets the scheduler's timesteps for inference" @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("num_inference_steps", default=50), InputParam("timesteps"), @@ -556,7 +556,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference"), OutputParam( @@ -605,7 +605,7 @@ class StableDiffusionXLInpaintPrepareLatentsStep(ModularPipelineBlocks): model_name = "stable-diffusion-xl" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", EulerDiscreteScheduler), ] @@ -615,7 +615,7 @@ def description(self) -> str: return "Step that prepares the latents for the inpainting process" @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [ InputParam("latents"), InputParam("num_images_per_prompt", default=1), @@ -664,7 +664,7 @@ def inputs(self) -> List[Tuple[str, Any]]: ] @property - def intermediate_outputs(self) -> List[str]: + def intermediate_outputs(self) -> list[str]: return [ OutputParam( "latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process" @@ -879,7 +879,7 @@ class StableDiffusionXLImg2ImgPrepareLatentsStep(ModularPipelineBlocks): model_name = "stable-diffusion-xl" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("vae", AutoencoderKL), ComponentSpec("scheduler", EulerDiscreteScheduler), @@ -890,7 +890,7 @@ def description(self) -> str: return "Step that prepares the latents for the image-to-image generation process" @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [ InputParam("latents"), InputParam("num_images_per_prompt", default=1), @@ -918,7 +918,7 @@ def inputs(self) -> List[Tuple[str, Any]]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process" @@ -955,7 +955,7 @@ class StableDiffusionXLPrepareLatentsStep(ModularPipelineBlocks): model_name = "stable-diffusion-xl" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", EulerDiscreteScheduler), ComponentSpec("vae", AutoencoderKL), @@ -966,7 +966,7 @@ def description(self) -> str: return "Prepare latents step that prepares the latents for the text-to-image generation process" @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("height"), InputParam("width"), @@ -983,7 +983,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process" @@ -1061,13 +1061,13 @@ class StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep(ModularPipelineB model_name = "stable-diffusion-xl" @property - def expected_configs(self) -> List[ConfigSpec]: + def expected_configs(self) -> list[ConfigSpec]: return [ ConfigSpec("requires_aesthetics_score", False), ] @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("unet", UNet2DConditionModel), ComponentSpec( @@ -1083,7 +1083,7 @@ def description(self) -> str: return "Step that prepares the additional conditioning for the image-to-image/inpainting generation process" @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [ InputParam("original_size"), InputParam("target_size"), @@ -1115,7 +1115,7 @@ def inputs(self) -> List[Tuple[str, Any]]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "add_time_ids", @@ -1284,7 +1284,7 @@ def description(self) -> str: return "Step that prepares the additional conditioning for the text-to-image generation process" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("unet", UNet2DConditionModel), ComponentSpec( @@ -1296,7 +1296,7 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [ InputParam("original_size"), InputParam("target_size"), @@ -1326,7 +1326,7 @@ def inputs(self) -> List[Tuple[str, Any]]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "add_time_ids", @@ -1458,7 +1458,7 @@ class StableDiffusionXLControlNetInputStep(ModularPipelineBlocks): model_name = "stable-diffusion-xl" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("controlnet", ControlNetModel), ComponentSpec( @@ -1474,7 +1474,7 @@ def description(self) -> str: return "step that prepare inputs for controlnet" @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [ InputParam("control_image", required=True), InputParam("control_guidance_start", default=0.0), @@ -1502,26 +1502,26 @@ def inputs(self) -> List[Tuple[str, Any]]: ), InputParam( "crops_coords", - type_hint=Optional[Tuple[int]], + type_hint=tuple[int] | None, description="The crop coordinates to use for preprocess/postprocess the image and mask, for inpainting task only. Can be generated in vae_encode step.", ), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam("controlnet_cond", type_hint=torch.Tensor, description="The processed control image"), OutputParam( - "control_guidance_start", type_hint=List[float], description="The controlnet guidance start values" + "control_guidance_start", type_hint=list[float], description="The controlnet guidance start values" ), OutputParam( - "control_guidance_end", type_hint=List[float], description="The controlnet guidance end values" + "control_guidance_end", type_hint=list[float], description="The controlnet guidance end values" ), OutputParam( - "conditioning_scale", type_hint=List[float], description="The controlnet conditioning scale values" + "conditioning_scale", type_hint=list[float], description="The controlnet conditioning scale values" ), OutputParam("guess_mode", type_hint=bool, description="Whether guess mode is used"), - OutputParam("controlnet_keep", type_hint=List[float], description="The controlnet keep values"), + OutputParam("controlnet_keep", type_hint=list[float], description="The controlnet keep values"), ] # Modified from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl.StableDiffusionXLControlNetPipeline.prepare_image @@ -1672,7 +1672,7 @@ class StableDiffusionXLControlNetUnionInputStep(ModularPipelineBlocks): model_name = "stable-diffusion-xl" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("controlnet", ControlNetUnionModel), ComponentSpec( @@ -1688,7 +1688,7 @@ def description(self) -> str: return "step that prepares inputs for the ControlNetUnion model" @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [ InputParam("control_image", required=True), InputParam("control_mode", required=True), @@ -1723,18 +1723,18 @@ def inputs(self) -> List[Tuple[str, Any]]: ), InputParam( "crops_coords", - type_hint=Optional[Tuple[int]], + type_hint=tuple[int] | None, description="The crop coordinates to use for preprocess/postprocess the image and mask, for inpainting task only. Can be generated in vae_encode step.", ), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ - OutputParam("controlnet_cond", type_hint=List[torch.Tensor], description="The processed control images"), + OutputParam("controlnet_cond", type_hint=list[torch.Tensor], description="The processed control images"), OutputParam( "control_type_idx", - type_hint=List[int], + type_hint=list[int], description="The control mode indices", kwargs_type="controlnet_kwargs", ), @@ -1747,10 +1747,10 @@ def intermediate_outputs(self) -> List[OutputParam]: OutputParam("control_guidance_start", type_hint=float, description="The controlnet guidance start value"), OutputParam("control_guidance_end", type_hint=float, description="The controlnet guidance end value"), OutputParam( - "conditioning_scale", type_hint=List[float], description="The controlnet conditioning scale values" + "conditioning_scale", type_hint=list[float], description="The controlnet conditioning scale values" ), OutputParam("guess_mode", type_hint=bool, description="Whether guess mode is used"), - OutputParam("controlnet_keep", type_hint=List[float], description="The controlnet keep values"), + OutputParam("controlnet_keep", type_hint=list[float], description="The controlnet keep values"), ] # Modified from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl.StableDiffusionXLControlNetPipeline.prepare_image diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py index 6e0307260d1d..7e505559f685 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, List, Tuple, Union +from typing import Any import numpy as np import PIL @@ -36,7 +36,7 @@ class StableDiffusionXLDecodeStep(ModularPipelineBlocks): model_name = "stable-diffusion-xl" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("vae", AutoencoderKL), ComponentSpec( @@ -52,7 +52,7 @@ def description(self) -> str: return "Step that decodes the denoised latents into images" @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [ InputParam("output_type", default="pil"), InputParam( @@ -64,11 +64,11 @@ def inputs(self) -> List[Tuple[str, Any]]: ] @property - def intermediate_outputs(self) -> List[str]: + def intermediate_outputs(self) -> list[str]: return [ OutputParam( "images", - type_hint=Union[List[PIL.Image.Image], List[torch.Tensor], List[np.array]], + type_hint=list[PIL.Image.Image] | list[torch.Tensor] | list[np.array], description="The generated images, can be a PIL.Image.Image, torch.Tensor or a numpy array", ) ] @@ -153,7 +153,7 @@ def description(self) -> str: ) @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec( "image_processor", @@ -164,19 +164,19 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [ InputParam("image"), InputParam("mask_image"), InputParam("padding_mask_crop"), InputParam( "images", - type_hint=Union[List[PIL.Image.Image], List[torch.Tensor], List[np.array]], + type_hint=list[PIL.Image.Image] | list[torch.Tensor] | list[np.array], description="The generated images from the decode step", ), InputParam( "crops_coords", - type_hint=Tuple[int, int], + type_hint=tuple[int, int], description="The crop coordinates to use for preprocess/postprocess the image and mask, for inpainting task only. Can be generated in vae_encode step.", ), ] diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py index 862315e59169..0190bc3ea62f 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, List, Optional, Tuple +from typing import Any import torch @@ -41,7 +41,7 @@ class StableDiffusionXLLoopBeforeDenoiser(ModularPipelineBlocks): model_name = "stable-diffusion-xl" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", EulerDiscreteScheduler), ] @@ -55,7 +55,7 @@ def description(self) -> str: ) @property - def inputs(self) -> List[str]: + def inputs(self) -> list[str]: return [ InputParam( "latents", @@ -77,7 +77,7 @@ class StableDiffusionXLInpaintLoopBeforeDenoiser(ModularPipelineBlocks): model_name = "stable-diffusion-xl" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", EulerDiscreteScheduler), ComponentSpec("unet", UNet2DConditionModel), @@ -91,7 +91,7 @@ def description(self) -> str: ) @property - def inputs(self) -> List[str]: + def inputs(self) -> list[str]: return [ InputParam( "latents", @@ -101,12 +101,12 @@ def inputs(self) -> List[str]: ), InputParam( "mask", - type_hint=Optional[torch.Tensor], + type_hint=torch.Tensor | None, description="The mask to use for the denoising process, for inpainting task only. Can be generated in vae_encode or prepare_latent step.", ), InputParam( "masked_image_latents", - type_hint=Optional[torch.Tensor], + type_hint=torch.Tensor | None, description="The masked image latents to use for the denoising process, for inpainting task only. Can be generated in vae_encode or prepare_latent step.", ), ] @@ -148,7 +148,7 @@ class StableDiffusionXLLoopDenoiser(ModularPipelineBlocks): model_name = "stable-diffusion-xl" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec( "guider", @@ -168,7 +168,7 @@ def description(self) -> str: ) @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [ InputParam("cross_attention_kwargs"), InputParam( @@ -179,7 +179,7 @@ def inputs(self) -> List[Tuple[str, Any]]: ), InputParam( "timestep_cond", - type_hint=Optional[torch.Tensor], + type_hint=torch.Tensor | None, description="The guidance scale embedding to use for Latent Consistency Models(LCMs). Can be generated in prepare_additional_conditioning step.", ), InputParam( @@ -262,7 +262,7 @@ class StableDiffusionXLControlNetLoopDenoiser(ModularPipelineBlocks): model_name = "stable-diffusion-xl" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec( "guider", @@ -283,7 +283,7 @@ def description(self) -> str: ) @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [ InputParam("cross_attention_kwargs"), InputParam( @@ -306,12 +306,12 @@ def inputs(self) -> List[Tuple[str, Any]]: InputParam( "controlnet_keep", required=True, - type_hint=List[float], + type_hint=list[float], description="The controlnet keep values to use for the denoising process. Can be generated in prepare_controlnet_inputs step.", ), InputParam( "timestep_cond", - type_hint=Optional[torch.Tensor], + type_hint=torch.Tensor | None, description="The guidance scale embedding to use for Latent Consistency Models(LCMs), can be generated by prepare_additional_conditioning step", ), InputParam( @@ -472,7 +472,7 @@ class StableDiffusionXLLoopAfterDenoiser(ModularPipelineBlocks): model_name = "stable-diffusion-xl" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", EulerDiscreteScheduler), ] @@ -486,14 +486,14 @@ def description(self) -> str: ) @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [ InputParam("eta", default=0.0), InputParam("generator"), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")] # YiYi TODO: move this out of here @@ -537,7 +537,7 @@ class StableDiffusionXLInpaintLoopAfterDenoiser(ModularPipelineBlocks): model_name = "stable-diffusion-xl" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", EulerDiscreteScheduler), ComponentSpec("unet", UNet2DConditionModel), @@ -552,7 +552,7 @@ def description(self) -> str: ) @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [ InputParam("eta", default=0.0), InputParam("generator"), @@ -564,23 +564,23 @@ def inputs(self) -> List[Tuple[str, Any]]: ), InputParam( "mask", - type_hint=Optional[torch.Tensor], + type_hint=torch.Tensor | None, description="The mask to use for the denoising process, for inpainting task only. Can be generated in vae_encode or prepare_latent step.", ), InputParam( "noise", - type_hint=Optional[torch.Tensor], + type_hint=torch.Tensor | None, description="The noise added to the image latents, for inpainting task only. Can be generated in prepare_latent step.", ), InputParam( "image_latents", - type_hint=Optional[torch.Tensor], + type_hint=torch.Tensor | None, description="The image latents to use for the denoising process, for inpainting/image-to-image task only. Can be generated in vae_encode or prepare_latent step.", ), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")] @staticmethod @@ -654,7 +654,7 @@ def description(self) -> str: ) @property - def loop_expected_components(self) -> List[ComponentSpec]: + def loop_expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec( "guider", @@ -667,7 +667,7 @@ def loop_expected_components(self) -> List[ComponentSpec]: ] @property - def loop_inputs(self) -> List[InputParam]: + def loop_inputs(self) -> list[InputParam]: return [ InputParam( "timesteps", diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/encoders.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/encoders.py index 90b254b6f5d4..8387ae7bd6b6 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_xl/encoders.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/encoders.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Optional, Tuple - import torch from transformers import ( CLIPImageProcessor, @@ -45,7 +43,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -71,7 +69,7 @@ def description(self) -> str: ) @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("image_encoder", CLIPVisionModelWithProjection), ComponentSpec( @@ -90,7 +88,7 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam( "ip_adapter_image", @@ -101,7 +99,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam("ip_adapter_embeds", type_hint=torch.Tensor, description="IP adapter image embeddings"), OutputParam( @@ -223,7 +221,7 @@ def description(self) -> str: return "Text Encoder step that generate text_embeddings to guide the image generation" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("text_encoder", CLIPTextModel), ComponentSpec("text_encoder_2", CLIPTextModelWithProjection), @@ -238,11 +236,11 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def expected_configs(self) -> List[ConfigSpec]: + def expected_configs(self) -> list[ConfigSpec]: return [ConfigSpec("force_zeros_for_empty_prompt", True)] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("prompt"), InputParam("prompt_2"), @@ -253,7 +251,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "prompt_embeds", @@ -296,26 +294,26 @@ def check_inputs(block_state): def encode_prompt( components, prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, + prompt_2: str | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, prepare_unconditional_embeds: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders device: (`torch.device`): @@ -324,11 +322,11 @@ def encode_prompt( number of images that should be generated per prompt prepare_unconditional_embeds (`bool`): whether to use prepare unconditional embeddings or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -450,7 +448,7 @@ def encode_prompt( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) - uncond_tokens: List[str] + uncond_tokens: list[str] if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -584,7 +582,7 @@ def description(self) -> str: return "Vae Encoder step that encode the input image into a latent representation" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("vae", AutoencoderKL), ComponentSpec( @@ -596,7 +594,7 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("image", required=True), InputParam("height"), @@ -605,13 +603,13 @@ def inputs(self) -> List[InputParam]: InputParam("dtype", type_hint=torch.dtype, description="Data type of model tensor inputs"), InputParam( "preprocess_kwargs", - type_hint=Optional[dict], + type_hint=dict | None, description="A kwargs dictionary that if specified is passed along to the `ImageProcessor` as defined under `self.image_processor` in [diffusers.image_processor.VaeImageProcessor]", ), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "image_latents", @@ -687,7 +685,7 @@ class StableDiffusionXLInpaintVaeEncoderStep(ModularPipelineBlocks): model_name = "stable-diffusion-xl" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("vae", AutoencoderKL), ComponentSpec( @@ -711,7 +709,7 @@ def description(self) -> str: return "Vae encoder step that prepares the image and mask for the inpainting process" @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("height"), InputParam("width"), @@ -723,7 +721,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "image_latents", type_hint=torch.Tensor, description="The latents representation of the input image" @@ -736,7 +734,7 @@ def intermediate_outputs(self) -> List[OutputParam]: ), OutputParam( "crops_coords", - type_hint=Optional[Tuple[int, int]], + type_hint=tuple[int, int] | None, description="The crop coordinates to use for the preprocess/postprocess of the image and mask", ), ] diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks_stable_diffusion_xl.py similarity index 55% rename from src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks.py rename to src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks_stable_diffusion_xl.py index 68b5e33755b5..a7a18e514777 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks_stable_diffusion_xl.py @@ -14,7 +14,7 @@ from ...utils import logging from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks -from ..modular_pipeline_utils import InsertableDict +from ..modular_pipeline_utils import OutputParam from .before_denoise import ( StableDiffusionXLControlNetInputStep, StableDiffusionXLControlNetUnionInputStep, @@ -277,7 +277,161 @@ def description(self): # ip-adapter, controlnet, text2img, img2img, inpainting +# auto_docstring class StableDiffusionXLAutoBlocks(SequentialPipelineBlocks): + """ + Auto Modular pipeline for text-to-image, image-to-image, inpainting, and controlnet tasks using Stable Diffusion + XL. + + Supported workflows: + - `text2image`: requires `prompt` + - `image2image`: requires `image`, `prompt` + - `inpainting`: requires `mask_image`, `image`, `prompt` + - `controlnet_text2image`: requires `control_image`, `prompt` + - `controlnet_image2image`: requires `control_image`, `image`, `prompt` + - `controlnet_inpainting`: requires `control_image`, `mask_image`, `image`, `prompt` + - `controlnet_union_text2image`: requires `control_image`, `control_mode`, `prompt` + - `controlnet_union_image2image`: requires `control_image`, `control_mode`, `image`, `prompt` + - `controlnet_union_inpainting`: requires `control_image`, `control_mode`, `mask_image`, `image`, `prompt` + - `ip_adapter_text2image`: requires `ip_adapter_image`, `prompt` + - `ip_adapter_image2image`: requires `ip_adapter_image`, `image`, `prompt` + - `ip_adapter_inpainting`: requires `ip_adapter_image`, `mask_image`, `image`, `prompt` + - `ip_adapter_controlnet_text2image`: requires `ip_adapter_image`, `control_image`, `prompt` + - `ip_adapter_controlnet_image2image`: requires `ip_adapter_image`, `control_image`, `image`, `prompt` + - `ip_adapter_controlnet_inpainting`: requires `ip_adapter_image`, `control_image`, `mask_image`, `image`, + `prompt` + - `ip_adapter_controlnet_union_text2image`: requires `ip_adapter_image`, `control_image`, `control_mode`, + `prompt` + - `ip_adapter_controlnet_union_image2image`: requires `ip_adapter_image`, `control_image`, `control_mode`, + `image`, `prompt` + - `ip_adapter_controlnet_union_inpainting`: requires `ip_adapter_image`, `control_image`, `control_mode`, + `mask_image`, `image`, `prompt` + + Components: + text_encoder (`CLIPTextModel`) text_encoder_2 (`CLIPTextModelWithProjection`) tokenizer (`CLIPTokenizer`) + tokenizer_2 (`CLIPTokenizer`) guider (`ClassifierFreeGuidance`) image_encoder + (`CLIPVisionModelWithProjection`) feature_extractor (`CLIPImageProcessor`) unet (`UNet2DConditionModel`) vae + (`AutoencoderKL`) image_processor (`VaeImageProcessor`) mask_processor (`VaeImageProcessor`) scheduler + (`EulerDiscreteScheduler`) controlnet (`ControlNetUnionModel`) control_image_processor (`VaeImageProcessor`) + + Configs: + force_zeros_for_empty_prompt (default: True) requires_aesthetics_score (default: False) + + Inputs: + prompt (`None`, *optional*): + TODO: Add description. + prompt_2 (`None`, *optional*): + TODO: Add description. + negative_prompt (`None`, *optional*): + TODO: Add description. + negative_prompt_2 (`None`, *optional*): + TODO: Add description. + cross_attention_kwargs (`None`, *optional*): + TODO: Add description. + clip_skip (`None`, *optional*): + TODO: Add description. + ip_adapter_image (`Image | ndarray | Tensor | list | list | list`, *optional*): + The image(s) to be used as ip adapter + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + image (`None`, *optional*): + TODO: Add description. + mask_image (`None`, *optional*): + TODO: Add description. + padding_mask_crop (`None`, *optional*): + TODO: Add description. + dtype (`dtype`, *optional*): + The dtype of the model inputs + generator (`None`, *optional*): + TODO: Add description. + preprocess_kwargs (`dict | NoneType`, *optional*): + A kwargs dictionary that if specified is passed along to the `ImageProcessor` as defined under + `self.image_processor` in [diffusers.image_processor.VaeImageProcessor] + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + ip_adapter_embeds (`list`, *optional*): + Pre-generated image embeddings for IP-Adapter. Can be generated from ip_adapter step. + negative_ip_adapter_embeds (`list`, *optional*): + Pre-generated negative image embeddings for IP-Adapter. Can be generated from ip_adapter step. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + denoising_end (`None`, *optional*): + TODO: Add description. + strength (`None`, *optional*, defaults to 0.3): + TODO: Add description. + denoising_start (`None`, *optional*): + TODO: Add description. + latents (`None`): + TODO: Add description. + image_latents (`Tensor`, *optional*): + The latents representing the reference image for image-to-image/inpainting generation. Can be generated + in vae_encode step. + mask (`Tensor`, *optional*): + The mask for the inpainting generation. Can be generated in vae_encode step. + masked_image_latents (`Tensor`, *optional*): + The masked image latents for the inpainting generation (only for inpainting-specific unet). Can be + generated in vae_encode step. + original_size (`None`, *optional*): + TODO: Add description. + target_size (`None`, *optional*): + TODO: Add description. + negative_original_size (`None`, *optional*): + TODO: Add description. + negative_target_size (`None`, *optional*): + TODO: Add description. + crops_coords_top_left (`None`, *optional*, defaults to (0, 0)): + TODO: Add description. + negative_crops_coords_top_left (`None`, *optional*, defaults to (0, 0)): + TODO: Add description. + aesthetic_score (`None`, *optional*, defaults to 6.0): + TODO: Add description. + negative_aesthetic_score (`None`, *optional*, defaults to 2.0): + TODO: Add description. + control_image (`None`, *optional*): + TODO: Add description. + control_mode (`None`, *optional*): + TODO: Add description. + control_guidance_start (`None`, *optional*, defaults to 0.0): + TODO: Add description. + control_guidance_end (`None`, *optional*, defaults to 1.0): + TODO: Add description. + controlnet_conditioning_scale (`None`, *optional*, defaults to 1.0): + TODO: Add description. + guess_mode (`None`, *optional*, defaults to False): + TODO: Add description. + crops_coords (`tuple | NoneType`, *optional*): + The crop coordinates to use for preprocess/postprocess the image and mask, for inpainting task only. Can + be generated in vae_encode step. + controlnet_cond (`Tensor`, *optional*): + The control image to use for the denoising process. Can be generated in prepare_controlnet_inputs step. + conditioning_scale (`float`, *optional*): + The controlnet conditioning scale value to use for the denoising process. Can be generated in + prepare_controlnet_inputs step. + controlnet_keep (`list`, *optional*): + The controlnet keep values to use for the denoising process. Can be generated in + prepare_controlnet_inputs step. + **denoiser_input_fields (`None`, *optional*): + All conditional model inputs that need to be prepared with guider. It should contain + prompt_embeds/negative_prompt_embeds, add_time_ids/negative_add_time_ids, + pooled_prompt_embeds/negative_pooled_prompt_embeds, and ip_adapter_embeds/negative_ip_adapter_embeds + (optional).please add `kwargs_type=denoiser_input_fields` to their parameter spec (`OutputParam`) when + they are created and added to the pipeline state + eta (`None`, *optional*, defaults to 0.0): + TODO: Add description. + output_type (`None`, *optional*, defaults to pil): + TODO: Add description. + + Outputs: + images (`list`): + Generated images. + """ + block_classes = [ StableDiffusionXLTextEncoderStep, StableDiffusionXLAutoIPAdapterStep, @@ -293,103 +447,66 @@ class StableDiffusionXLAutoBlocks(SequentialPipelineBlocks): "decode", ] - @property - def description(self): - return ( - "Auto Modular pipeline for text-to-image, image-to-image, inpainting, and controlnet tasks using Stable Diffusion XL.\n" - + "- for image-to-image generation, you need to provide either `image` or `image_latents`\n" - + "- for inpainting, you need to provide `mask_image` and `image`, optionally you can provide `padding_mask_crop` \n" - + "- to run the controlnet workflow, you need to provide `control_image`\n" - + "- to run the controlnet_union workflow, you need to provide `control_image` and `control_mode`\n" - + "- to run the ip_adapter workflow, you need to provide `ip_adapter_image`\n" - + "- for text-to-image generation, all you need to provide is `prompt`" - ) - - -# controlnet (input + denoise step) -class StableDiffusionXLAutoControlnetStep(SequentialPipelineBlocks): - block_classes = [ - StableDiffusionXLAutoControlNetInputStep, - StableDiffusionXLAutoControlNetDenoiseStep, - ] - block_names = ["controlnet_input", "controlnet_denoise"] + _workflow_map = { + "text2image": {"prompt": True}, + "image2image": {"image": True, "prompt": True}, + "inpainting": {"mask_image": True, "image": True, "prompt": True}, + "controlnet_text2image": {"control_image": True, "prompt": True}, + "controlnet_image2image": {"control_image": True, "image": True, "prompt": True}, + "controlnet_inpainting": {"control_image": True, "mask_image": True, "image": True, "prompt": True}, + "controlnet_union_text2image": {"control_image": True, "control_mode": True, "prompt": True}, + "controlnet_union_image2image": {"control_image": True, "control_mode": True, "image": True, "prompt": True}, + "controlnet_union_inpainting": { + "control_image": True, + "control_mode": True, + "mask_image": True, + "image": True, + "prompt": True, + }, + "ip_adapter_text2image": {"ip_adapter_image": True, "prompt": True}, + "ip_adapter_image2image": {"ip_adapter_image": True, "image": True, "prompt": True}, + "ip_adapter_inpainting": {"ip_adapter_image": True, "mask_image": True, "image": True, "prompt": True}, + "ip_adapter_controlnet_text2image": {"ip_adapter_image": True, "control_image": True, "prompt": True}, + "ip_adapter_controlnet_image2image": { + "ip_adapter_image": True, + "control_image": True, + "image": True, + "prompt": True, + }, + "ip_adapter_controlnet_inpainting": { + "ip_adapter_image": True, + "control_image": True, + "mask_image": True, + "image": True, + "prompt": True, + }, + "ip_adapter_controlnet_union_text2image": { + "ip_adapter_image": True, + "control_image": True, + "control_mode": True, + "prompt": True, + }, + "ip_adapter_controlnet_union_image2image": { + "ip_adapter_image": True, + "control_image": True, + "control_mode": True, + "image": True, + "prompt": True, + }, + "ip_adapter_controlnet_union_inpainting": { + "ip_adapter_image": True, + "control_image": True, + "control_mode": True, + "mask_image": True, + "image": True, + "prompt": True, + }, + } @property def description(self): - return ( - "Controlnet auto step that prepare the controlnet input and denoise the latents. " - + "It works for both controlnet and controlnet_union and supports text2img, img2img and inpainting tasks." - + " (it should be replace at 'denoise' step)" - ) - - -TEXT2IMAGE_BLOCKS = InsertableDict( - [ - ("text_encoder", StableDiffusionXLTextEncoderStep), - ("input", StableDiffusionXLInputStep), - ("set_timesteps", StableDiffusionXLSetTimestepsStep), - ("prepare_latents", StableDiffusionXLPrepareLatentsStep), - ("prepare_add_cond", StableDiffusionXLPrepareAdditionalConditioningStep), - ("denoise", StableDiffusionXLDenoiseStep), - ("decode", StableDiffusionXLDecodeStep), - ] -) + return "Auto Modular pipeline for text-to-image, image-to-image, inpainting, and controlnet tasks using Stable Diffusion XL." -IMAGE2IMAGE_BLOCKS = InsertableDict( - [ - ("text_encoder", StableDiffusionXLTextEncoderStep), - ("vae_encoder", StableDiffusionXLVaeEncoderStep), - ("input", StableDiffusionXLInputStep), - ("set_timesteps", StableDiffusionXLImg2ImgSetTimestepsStep), - ("prepare_latents", StableDiffusionXLImg2ImgPrepareLatentsStep), - ("prepare_add_cond", StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep), - ("denoise", StableDiffusionXLDenoiseStep), - ("decode", StableDiffusionXLDecodeStep), - ] -) - -INPAINT_BLOCKS = InsertableDict( - [ - ("text_encoder", StableDiffusionXLTextEncoderStep), - ("vae_encoder", StableDiffusionXLInpaintVaeEncoderStep), - ("input", StableDiffusionXLInputStep), - ("set_timesteps", StableDiffusionXLImg2ImgSetTimestepsStep), - ("prepare_latents", StableDiffusionXLInpaintPrepareLatentsStep), - ("prepare_add_cond", StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep), - ("denoise", StableDiffusionXLInpaintDenoiseStep), - ("decode", StableDiffusionXLInpaintDecodeStep), - ] -) - -CONTROLNET_BLOCKS = InsertableDict( - [ - ("denoise", StableDiffusionXLAutoControlnetStep), - ] -) - - -IP_ADAPTER_BLOCKS = InsertableDict( - [ - ("ip_adapter", StableDiffusionXLAutoIPAdapterStep), - ] -) - -AUTO_BLOCKS = InsertableDict( - [ - ("text_encoder", StableDiffusionXLTextEncoderStep), - ("ip_adapter", StableDiffusionXLAutoIPAdapterStep), - ("vae_encoder", StableDiffusionXLAutoVaeEncoderStep), - ("denoise", StableDiffusionXLCoreDenoiseStep), - ("decode", StableDiffusionXLAutoDecodeStep), - ] -) - - -ALL_BLOCKS = { - "text2img": TEXT2IMAGE_BLOCKS, - "img2img": IMAGE2IMAGE_BLOCKS, - "inpaint": INPAINT_BLOCKS, - "controlnet": CONTROLNET_BLOCKS, - "ip_adapter": IP_ADAPTER_BLOCKS, - "auto": AUTO_BLOCKS, -} + @property + def outputs(self): + return [OutputParam.template("images")] diff --git a/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_pipeline.py b/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_pipeline.py index f2a4c96073ea..209e2b11814f 100644 --- a/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_pipeline.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Optional, Tuple, Union - import numpy as np import PIL import torch @@ -94,30 +92,30 @@ def num_channels_latents(self): # auto_docstring SDXL_INPUTS_SCHEMA = { "prompt": InputParam( - "prompt", type_hint=Union[str, List[str]], description="The prompt or prompts to guide the image generation" + "prompt", type_hint=str | list[str], description="The prompt or prompts to guide the image generation" ), "prompt_2": InputParam( "prompt_2", - type_hint=Union[str, List[str]], + type_hint=str | list[str], description="The prompt or prompts to be sent to the tokenizer_2 and text_encoder_2", ), "negative_prompt": InputParam( "negative_prompt", - type_hint=Union[str, List[str]], + type_hint=str | list[str], description="The prompt or prompts not to guide the image generation", ), "negative_prompt_2": InputParam( "negative_prompt_2", - type_hint=Union[str, List[str]], + type_hint=str | list[str], description="The negative prompt or prompts for text_encoder_2", ), "cross_attention_kwargs": InputParam( "cross_attention_kwargs", - type_hint=Optional[dict], + type_hint=dict | None, description="Kwargs dictionary passed to the AttentionProcessor", ), "clip_skip": InputParam( - "clip_skip", type_hint=Optional[int], description="Number of layers to skip in CLIP text encoder" + "clip_skip", type_hint=int | None, description="Number of layers to skip in CLIP text encoder" ), "image": InputParam( "image", @@ -133,11 +131,11 @@ def num_channels_latents(self): ), "generator": InputParam( "generator", - type_hint=Optional[Union[torch.Generator, List[torch.Generator]]], + type_hint=torch.Generator | list[torch.Generator] | None, description="Generator(s) for deterministic generation", ), - "height": InputParam("height", type_hint=Optional[int], description="Height in pixels of the generated image"), - "width": InputParam("width", type_hint=Optional[int], description="Width in pixels of the generated image"), + "height": InputParam("height", type_hint=int | None, description="Height in pixels of the generated image"), + "width": InputParam("width", type_hint=int | None, description="Width in pixels of the generated image"), "num_images_per_prompt": InputParam( "num_images_per_prompt", type_hint=int, default=1, description="Number of images to generate per prompt" ), @@ -145,14 +143,14 @@ def num_channels_latents(self): "num_inference_steps", type_hint=int, default=50, description="Number of denoising steps" ), "timesteps": InputParam( - "timesteps", type_hint=Optional[torch.Tensor], description="Custom timesteps for the denoising process" + "timesteps", type_hint=torch.Tensor | None, description="Custom timesteps for the denoising process" ), "sigmas": InputParam( - "sigmas", type_hint=Optional[torch.Tensor], description="Custom sigmas for the denoising process" + "sigmas", type_hint=torch.Tensor | None, description="Custom sigmas for the denoising process" ), "denoising_end": InputParam( "denoising_end", - type_hint=Optional[float], + type_hint=float | None, description="Fraction of denoising process to complete before termination", ), # YiYi Notes: img2img defaults to 0.3, inpainting defaults to 0.9999 @@ -160,43 +158,43 @@ def num_channels_latents(self): "strength", type_hint=float, default=0.3, description="How much to transform the reference image" ), "denoising_start": InputParam( - "denoising_start", type_hint=Optional[float], description="Starting point of the denoising process" + "denoising_start", type_hint=float | None, description="Starting point of the denoising process" ), "latents": InputParam( - "latents", type_hint=Optional[torch.Tensor], description="Pre-generated noisy latents for image generation" + "latents", type_hint=torch.Tensor | None, description="Pre-generated noisy latents for image generation" ), "padding_mask_crop": InputParam( "padding_mask_crop", - type_hint=Optional[Tuple[int, int]], + type_hint=tuple[int, int] | None, description="Size of margin in crop for image and mask", ), "original_size": InputParam( "original_size", - type_hint=Optional[Tuple[int, int]], + type_hint=tuple[int, int] | None, description="Original size of the image for SDXL's micro-conditioning", ), "target_size": InputParam( - "target_size", type_hint=Optional[Tuple[int, int]], description="Target size for SDXL's micro-conditioning" + "target_size", type_hint=tuple[int, int] | None, description="Target size for SDXL's micro-conditioning" ), "negative_original_size": InputParam( "negative_original_size", - type_hint=Optional[Tuple[int, int]], + type_hint=tuple[int, int] | None, description="Negative conditioning based on image resolution", ), "negative_target_size": InputParam( "negative_target_size", - type_hint=Optional[Tuple[int, int]], + type_hint=tuple[int, int] | None, description="Negative conditioning based on target resolution", ), "crops_coords_top_left": InputParam( "crops_coords_top_left", - type_hint=Tuple[int, int], + type_hint=tuple[int, int], default=(0, 0), description="Top-left coordinates for SDXL's micro-conditioning", ), "negative_crops_coords_top_left": InputParam( "negative_crops_coords_top_left", - type_hint=Tuple[int, int], + type_hint=tuple[int, int], default=(0, 0), description="Negative conditioning crop coordinates", ), @@ -221,19 +219,19 @@ def num_channels_latents(self): ), "control_guidance_start": InputParam( "control_guidance_start", - type_hint=Union[float, List[float]], + type_hint=float | list[float], default=0.0, description="When ControlNet starts applying", ), "control_guidance_end": InputParam( "control_guidance_end", - type_hint=Union[float, List[float]], + type_hint=float | list[float], default=1.0, description="When ControlNet stops applying", ), "controlnet_conditioning_scale": InputParam( "controlnet_conditioning_scale", - type_hint=Union[float, List[float]], + type_hint=float | list[float], default=1.0, description="Scale factor for ControlNet outputs", ), @@ -244,7 +242,7 @@ def num_channels_latents(self): description="Enables ControlNet encoder to recognize input without prompts", ), "control_mode": InputParam( - "control_mode", type_hint=List[int], required=True, description="Control mode for union controlnet" + "control_mode", type_hint=list[int], required=True, description="Control mode for union controlnet" ), "prompt_embeds": InputParam( "prompt_embeds", @@ -264,7 +262,7 @@ def num_channels_latents(self): "batch_size": InputParam("batch_size", type_hint=int, required=True, description="Number of prompts"), "dtype": InputParam("dtype", type_hint=torch.dtype, description="Data type of model tensor inputs"), "preprocess_kwargs": InputParam( - "preprocess_kwargs", type_hint=Optional[dict], description="Kwargs for ImageProcessor" + "preprocess_kwargs", type_hint=dict | None, description="Kwargs for ImageProcessor" ), "latent_timestep": InputParam( "latent_timestep", type_hint=torch.Tensor, required=True, description="Initial noise level timestep" @@ -284,18 +282,18 @@ def num_channels_latents(self): ), "timestep_cond": InputParam("timestep_cond", type_hint=torch.Tensor, description="Timestep conditioning for LCM"), "noise": InputParam("noise", type_hint=torch.Tensor, description="Noise added to image latents"), - "crops_coords": InputParam("crops_coords", type_hint=Optional[Tuple[int]], description="Crop coordinates"), + "crops_coords": InputParam("crops_coords", type_hint=tuple[int] | None, description="Crop coordinates"), "ip_adapter_embeds": InputParam( - "ip_adapter_embeds", type_hint=List[torch.Tensor], description="Image embeddings for IP-Adapter" + "ip_adapter_embeds", type_hint=list[torch.Tensor], description="Image embeddings for IP-Adapter" ), "negative_ip_adapter_embeds": InputParam( "negative_ip_adapter_embeds", - type_hint=List[torch.Tensor], + type_hint=list[torch.Tensor], description="Negative image embeddings for IP-Adapter", ), "images": InputParam( "images", - type_hint=Union[List[PIL.Image.Image], List[torch.Tensor], List[np.array]], + type_hint=list[PIL.Image.Image] | list[torch.Tensor] | list[np.array], required=True, description="Generated images", ), @@ -324,7 +322,7 @@ def num_channels_latents(self): "masked_image_latents": OutputParam( "masked_image_latents", type_hint=torch.Tensor, description="Masked image latents for inpainting" ), - "crops_coords": OutputParam("crops_coords", type_hint=Optional[Tuple[int]], description="Crop coordinates"), + "crops_coords": OutputParam("crops_coords", type_hint=tuple[int] | None, description="Crop coordinates"), "timesteps": OutputParam("timesteps", type_hint=torch.Tensor, description="Timesteps for inference"), "num_inference_steps": OutputParam("num_inference_steps", type_hint=int, description="Number of denoising steps"), "latent_timestep": OutputParam( @@ -338,16 +336,16 @@ def num_channels_latents(self): "latents": OutputParam("latents", type_hint=torch.Tensor, description="Denoised latents"), "noise": OutputParam("noise", type_hint=torch.Tensor, description="Noise added to image latents"), "ip_adapter_embeds": OutputParam( - "ip_adapter_embeds", type_hint=List[torch.Tensor], description="Image embeddings for IP-Adapter" + "ip_adapter_embeds", type_hint=list[torch.Tensor], description="Image embeddings for IP-Adapter" ), "negative_ip_adapter_embeds": OutputParam( "negative_ip_adapter_embeds", - type_hint=List[torch.Tensor], + type_hint=list[torch.Tensor], description="Negative image embeddings for IP-Adapter", ), "images": OutputParam( "images", - type_hint=Union[List[PIL.Image.Image], List[torch.Tensor], List[np.array]], + type_hint=list[PIL.Image.Image] | list[torch.Tensor] | list[np.array], description="Generated images", ), } @@ -356,9 +354,7 @@ def num_channels_latents(self): SDXL_OUTPUTS_SCHEMA = { "images": OutputParam( "images", - type_hint=Union[ - Tuple[Union[List[PIL.Image.Image], List[torch.Tensor], List[np.array]]], StableDiffusionXLPipelineOutput - ], + type_hint=tuple[list[PIL.Image.Image] | list[torch.Tensor] | list[np.array]] | StableDiffusionXLPipelineOutput, description="The final generated images", ) } diff --git a/src/diffusers/modular_pipelines/wan/__init__.py b/src/diffusers/modular_pipelines/wan/__init__.py index 73f67c9afed2..284b6c9fa436 100644 --- a/src/diffusers/modular_pipelines/wan/__init__.py +++ b/src/diffusers/modular_pipelines/wan/__init__.py @@ -21,16 +21,16 @@ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) else: - _import_structure["decoders"] = ["WanImageVaeDecoderStep"] - _import_structure["encoders"] = ["WanTextEncoderStep"] - _import_structure["modular_blocks"] = [ - "ALL_BLOCKS", - "Wan22AutoBlocks", - "WanAutoBlocks", - "WanAutoImageEncoderStep", - "WanAutoVaeImageEncoderStep", + _import_structure["modular_blocks_wan"] = ["WanBlocks"] + _import_structure["modular_blocks_wan22"] = ["Wan22Blocks"] + _import_structure["modular_blocks_wan22_i2v"] = ["Wan22Image2VideoBlocks"] + _import_structure["modular_blocks_wan_i2v"] = ["WanImage2VideoAutoBlocks"] + _import_structure["modular_pipeline"] = [ + "Wan22Image2VideoModularPipeline", + "Wan22ModularPipeline", + "WanImage2VideoModularPipeline", + "WanModularPipeline", ] - _import_structure["modular_pipeline"] = ["WanModularPipeline"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: @@ -39,16 +39,16 @@ except OptionalDependencyNotAvailable: from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 else: - from .decoders import WanImageVaeDecoderStep - from .encoders import WanTextEncoderStep - from .modular_blocks import ( - ALL_BLOCKS, - Wan22AutoBlocks, - WanAutoBlocks, - WanAutoImageEncoderStep, - WanAutoVaeImageEncoderStep, + from .modular_blocks_wan import WanBlocks + from .modular_blocks_wan22 import Wan22Blocks + from .modular_blocks_wan22_i2v import Wan22Image2VideoBlocks + from .modular_blocks_wan_i2v import WanImage2VideoAutoBlocks + from .modular_pipeline import ( + Wan22Image2VideoModularPipeline, + Wan22ModularPipeline, + WanImage2VideoModularPipeline, + WanModularPipeline, ) - from .modular_pipeline import WanModularPipeline else: import sys diff --git a/src/diffusers/modular_pipelines/wan/before_denoise.py b/src/diffusers/modular_pipelines/wan/before_denoise.py index e2f8d3e7d88b..398b9665522c 100644 --- a/src/diffusers/modular_pipelines/wan/before_denoise.py +++ b/src/diffusers/modular_pipelines/wan/before_denoise.py @@ -13,7 +13,6 @@ # limitations under the License. import inspect -from typing import List, Optional, Tuple, Union import torch @@ -93,7 +92,7 @@ def repeat_tensor_to_batch_size( def calculate_dimension_from_latents( latents: torch.Tensor, vae_scale_factor_temporal: int, vae_scale_factor_spatial: int -) -> Tuple[int, int]: +) -> tuple[int, int]: """Calculate image dimensions from latent tensor dimensions. This function converts latent temporal and spatial dimensions to image temporal and spatial dimensions by @@ -108,7 +107,7 @@ def calculate_dimension_from_latents( Typically 8 for most VAEs (image is 8x larger than latents in each dimension) Returns: - Tuple[int, int]: The calculated image dimensions as (height, width) + tuple[int, int]: The calculated image dimensions as (height, width) Raises: ValueError: If latents tensor doesn't have 4 or 5 dimensions @@ -129,10 +128,10 @@ def calculate_dimension_from_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -147,15 +146,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -201,13 +200,13 @@ def description(self) -> str: ) @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("transformer", WanTransformer3DModel), ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("num_videos_per_prompt", default=1), InputParam( @@ -224,7 +223,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[str]: + def intermediate_outputs(self) -> list[str]: return [ OutputParam( "batch_size", @@ -280,8 +279,8 @@ class WanAdditionalInputsStep(ModularPipelineBlocks): def __init__( self, - image_latent_inputs: List[str] = ["first_frame_latents"], - additional_batch_inputs: List[str] = [], + image_latent_inputs: list[str] = ["image_condition_latents"], + additional_batch_inputs: list[str] = [], ): """Initialize a configurable step that standardizes the inputs for the denoising step. It:\n" @@ -292,22 +291,18 @@ def __init__( This is a dynamic block that allows you to configure which inputs to process. Args: - image_latent_inputs (List[str], optional): Names of image latent tensors to process. + image_latent_inputs (list[str], optional): Names of image latent tensors to process. In additional to adjust batch size of these inputs, they will be used to determine height/width. Can be - a single string or list of strings. Defaults to ["first_frame_latents"]. + a single string or list of strings. Defaults to ["image_condition_latents"]. additional_batch_inputs (List[str], optional): Names of additional conditional input tensors to expand batch size. These tensors will only have their batch dimensions adjusted to match the final batch size. Can be a single string or list of strings. Defaults to []. Examples: - # Configure to process first_frame_latents (default behavior) WanAdditionalInputsStep() - - # Configure to process multiple image latent inputs - WanAdditionalInputsStep(image_latent_inputs=["first_frame_latents", "last_frame_latents"]) - - # Configure to process image latents and additional batch inputs WanAdditionalInputsStep( - image_latent_inputs=["first_frame_latents"], additional_batch_inputs=["image_embeds"] + # Configure to process image_condition_latents (default behavior) WanAdditionalInputsStep() # Configure to + process image latents and additional batch inputs WanAdditionalInputsStep( + image_latent_inputs=["image_condition_latents"], additional_batch_inputs=["image_embeds"] ) """ if not isinstance(image_latent_inputs, list): @@ -343,7 +338,7 @@ def description(self) -> str: return summary_section + inputs_info + placement_section @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: inputs = [ InputParam(name="num_videos_per_prompt", default=1), InputParam(name="batch_size", required=True), @@ -413,7 +408,7 @@ class WanSetTimestepsStep(ModularPipelineBlocks): model_name = "wan" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", UniPCMultistepScheduler), ] @@ -423,7 +418,7 @@ def description(self) -> str: return "Step that sets the scheduler's timesteps for inference" @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("num_inference_steps", default=50), InputParam("timesteps"), @@ -455,12 +450,12 @@ def description(self) -> str: return "Prepare latents step that prepares the latents for the text-to-video generation process" @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("height", type_hint=int), InputParam("width", type_hint=int), InputParam("num_frames", type_hint=int), - InputParam("latents", type_hint=Optional[torch.Tensor]), + InputParam("latents", type_hint=torch.Tensor | None), InputParam("num_videos_per_prompt", type_hint=int, default=1), InputParam("generator"), InputParam( @@ -473,7 +468,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process" @@ -504,10 +499,10 @@ def prepare_latents( height: int = 480, width: int = 832, num_frames: int = 81, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: if latents is not None: return latents.to(device=device, dtype=dtype) @@ -557,81 +552,3 @@ def __call__(self, components: WanModularPipeline, state: PipelineState) -> Pipe self.set_block_state(state, block_state) return components, state - - -class WanPrepareFirstFrameLatentsStep(ModularPipelineBlocks): - model_name = "wan" - - @property - def description(self) -> str: - return "step that prepares the masked first frame latents and add it to the latent condition" - - @property - def inputs(self) -> List[InputParam]: - return [ - InputParam("first_frame_latents", type_hint=Optional[torch.Tensor]), - InputParam("num_frames", type_hint=int), - ] - - def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState: - block_state = self.get_block_state(state) - - batch_size, _, _, latent_height, latent_width = block_state.first_frame_latents.shape - - mask_lat_size = torch.ones(batch_size, 1, block_state.num_frames, latent_height, latent_width) - mask_lat_size[:, :, list(range(1, block_state.num_frames))] = 0 - - first_frame_mask = mask_lat_size[:, :, 0:1] - first_frame_mask = torch.repeat_interleave( - first_frame_mask, dim=2, repeats=components.vae_scale_factor_temporal - ) - mask_lat_size = torch.concat([first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2) - mask_lat_size = mask_lat_size.view( - batch_size, -1, components.vae_scale_factor_temporal, latent_height, latent_width - ) - mask_lat_size = mask_lat_size.transpose(1, 2) - mask_lat_size = mask_lat_size.to(block_state.first_frame_latents.device) - block_state.first_frame_latents = torch.concat([mask_lat_size, block_state.first_frame_latents], dim=1) - - self.set_block_state(state, block_state) - return components, state - - -class WanPrepareFirstLastFrameLatentsStep(ModularPipelineBlocks): - model_name = "wan" - - @property - def description(self) -> str: - return "step that prepares the masked latents with first and last frames and add it to the latent condition" - - @property - def inputs(self) -> List[InputParam]: - return [ - InputParam("first_last_frame_latents", type_hint=Optional[torch.Tensor]), - InputParam("num_frames", type_hint=int), - ] - - def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState: - block_state = self.get_block_state(state) - - batch_size, _, _, latent_height, latent_width = block_state.first_last_frame_latents.shape - - mask_lat_size = torch.ones(batch_size, 1, block_state.num_frames, latent_height, latent_width) - mask_lat_size[:, :, list(range(1, block_state.num_frames - 1))] = 0 - - first_frame_mask = mask_lat_size[:, :, 0:1] - first_frame_mask = torch.repeat_interleave( - first_frame_mask, dim=2, repeats=components.vae_scale_factor_temporal - ) - mask_lat_size = torch.concat([first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2) - mask_lat_size = mask_lat_size.view( - batch_size, -1, components.vae_scale_factor_temporal, latent_height, latent_width - ) - mask_lat_size = mask_lat_size.transpose(1, 2) - mask_lat_size = mask_lat_size.to(block_state.first_last_frame_latents.device) - block_state.first_last_frame_latents = torch.concat( - [mask_lat_size, block_state.first_last_frame_latents], dim=1 - ) - - self.set_block_state(state, block_state) - return components, state diff --git a/src/diffusers/modular_pipelines/wan/decoders.py b/src/diffusers/modular_pipelines/wan/decoders.py index 7cec318c1706..9d1a4cf4f348 100644 --- a/src/diffusers/modular_pipelines/wan/decoders.py +++ b/src/diffusers/modular_pipelines/wan/decoders.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, List, Tuple, Union +from typing import Any import numpy as np import PIL @@ -29,11 +29,11 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name -class WanImageVaeDecoderStep(ModularPipelineBlocks): +class WanVaeDecoderStep(ModularPipelineBlocks): model_name = "wan" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("vae", AutoencoderKLWan), ComponentSpec( @@ -49,22 +49,25 @@ def description(self) -> str: return "Step that decodes the denoised latents into images" @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [ InputParam( "latents", required=True, type_hint=torch.Tensor, description="The denoised latents from the denoising step", - ) + ), + InputParam( + "output_type", default="np", type_hint=str, description="The output type of the decoded videos" + ), ] @property - def intermediate_outputs(self) -> List[str]: + def intermediate_outputs(self) -> list[str]: return [ OutputParam( "videos", - type_hint=Union[List[List[PIL.Image.Image]], List[torch.Tensor], List[np.ndarray]], + type_hint=list[list[PIL.Image.Image]] | list[torch.Tensor] | list[np.ndarray], description="The generated videos, can be a PIL.Image.Image, torch.Tensor or a numpy array", ) ] @@ -87,7 +90,8 @@ def __call__(self, components, state: PipelineState) -> PipelineState: latents = latents.to(vae_dtype) block_state.videos = components.vae.decode(latents, return_dict=False)[0] - block_state.videos = components.video_processor.postprocess_video(block_state.videos, output_type="np") + output_type = getattr(block_state, "output_type", "np") + block_state.videos = components.video_processor.postprocess_video(block_state.videos, output_type=output_type) self.set_block_state(state, block_state) diff --git a/src/diffusers/modular_pipelines/wan/denoise.py b/src/diffusers/modular_pipelines/wan/denoise.py index 2da36f52da87..2f51f353012e 100644 --- a/src/diffusers/modular_pipelines/wan/denoise.py +++ b/src/diffusers/modular_pipelines/wan/denoise.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, List, Tuple +from typing import Any import torch @@ -46,7 +46,7 @@ def description(self) -> str: ) @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam( "latents", @@ -80,7 +80,7 @@ def description(self) -> str: ) @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam( "latents", @@ -89,52 +89,10 @@ def inputs(self) -> List[InputParam]: description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.", ), InputParam( - "first_frame_latents", + "image_condition_latents", required=True, type_hint=torch.Tensor, - description="The first frame latents to use for the denoising process. Can be generated in prepare_first_frame_latents step.", - ), - InputParam( - "dtype", - required=True, - type_hint=torch.dtype, - description="The dtype of the model inputs. Can be generated in input step.", - ), - ] - - @torch.no_grad() - def __call__(self, components: WanModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): - block_state.latent_model_input = torch.cat([block_state.latents, block_state.first_frame_latents], dim=1).to( - block_state.dtype - ) - return components, block_state - - -class WanFLF2VLoopBeforeDenoiser(ModularPipelineBlocks): - model_name = "wan" - - @property - def description(self) -> str: - return ( - "step within the denoising loop that prepares the latent input for the denoiser. " - "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " - "object (e.g. `WanDenoiseLoopWrapper`)" - ) - - @property - def inputs(self) -> List[InputParam]: - return [ - InputParam( - "latents", - required=True, - type_hint=torch.Tensor, - description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.", - ), - InputParam( - "first_last_frame_latents", - required=True, - type_hint=torch.Tensor, - description="The first and last frame latents to use for the denoising process. Can be generated in prepare_first_last_frame_latents step.", + description="The image condition latents to use for the denoising process. Can be generated in prepare_first_frame_latents/prepare_first_last_frame_latents step.", ), InputParam( "dtype", @@ -147,7 +105,7 @@ def inputs(self) -> List[InputParam]: @torch.no_grad() def __call__(self, components: WanModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): block_state.latent_model_input = torch.cat( - [block_state.latents, block_state.first_last_frame_latents], dim=1 + [block_state.latents, block_state.image_condition_latents], dim=1 ).to(block_state.dtype) return components, block_state @@ -157,7 +115,7 @@ class WanLoopDenoiser(ModularPipelineBlocks): def __init__( self, - guider_input_fields: Dict[str, Any] = {"encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds")}, + guider_input_fields: dict[str, Any] = {"encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds")}, ): """Initialize a denoiser block that calls the denoiser model. This block is used in Wan2.1. @@ -178,7 +136,7 @@ def __init__( super().__init__() @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec( "guider", @@ -198,7 +156,7 @@ def description(self) -> str: ) @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: inputs = [ InputParam("attention_kwargs"), InputParam( @@ -267,7 +225,7 @@ class Wan22LoopDenoiser(ModularPipelineBlocks): def __init__( self, - guider_input_fields: Dict[str, Any] = {"encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds")}, + guider_input_fields: dict[str, Any] = {"encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds")}, ): """Initialize a denoiser block that calls the denoiser model. This block is used in Wan2.2. @@ -288,7 +246,7 @@ def __init__( super().__init__() @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec( "guider", @@ -315,7 +273,7 @@ def description(self) -> str: ) @property - def expected_configs(self) -> List[ConfigSpec]: + def expected_configs(self) -> list[ConfigSpec]: return [ ConfigSpec( name="boundary_ratio", @@ -325,7 +283,7 @@ def expected_configs(self) -> List[ConfigSpec]: ] @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: inputs = [ InputParam("attention_kwargs"), InputParam( @@ -401,7 +359,7 @@ class WanLoopAfterDenoiser(ModularPipelineBlocks): model_name = "wan" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", UniPCMultistepScheduler), ] @@ -442,13 +400,13 @@ def description(self) -> str: ) @property - def loop_expected_components(self) -> List[ComponentSpec]: + def loop_expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", UniPCMultistepScheduler), ] @property - def loop_inputs(self) -> List[InputParam]: + def loop_inputs(self) -> list[InputParam]: return [ InputParam( "timesteps", @@ -584,29 +542,3 @@ def description(self) -> str: " - `WanLoopAfterDenoiser`\n" "This block supports image-to-video tasks for Wan2.2." ) - - -class WanFLF2VDenoiseStep(WanDenoiseLoopWrapper): - block_classes = [ - WanFLF2VLoopBeforeDenoiser, - WanLoopDenoiser( - guider_input_fields={ - "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"), - "encoder_hidden_states_image": "image_embeds", - } - ), - WanLoopAfterDenoiser, - ] - block_names = ["before_denoiser", "denoiser", "after_denoiser"] - - @property - def description(self) -> str: - return ( - "Denoise step that iteratively denoise the latents. \n" - "Its loop logic is defined in `WanDenoiseLoopWrapper.__call__` method \n" - "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n" - " - `WanFLF2VLoopBeforeDenoiser`\n" - " - `WanLoopDenoiser`\n" - " - `WanLoopAfterDenoiser`\n" - "This block supports FLF2V tasks for wan2.1." - ) diff --git a/src/diffusers/modular_pipelines/wan/encoders.py b/src/diffusers/modular_pipelines/wan/encoders.py index 4fd69c6ca6ab..3e675a66e4f2 100644 --- a/src/diffusers/modular_pipelines/wan/encoders.py +++ b/src/diffusers/modular_pipelines/wan/encoders.py @@ -13,7 +13,6 @@ # limitations under the License. import html -from typing import List, Optional, Union import numpy as np import PIL @@ -62,7 +61,7 @@ def prompt_clean(text): def get_t5_prompt_embeds( text_encoder: UMT5EncoderModel, tokenizer: AutoTokenizer, - prompt: Union[str, List[str]], + prompt: str | list[str], max_sequence_length: int, device: torch.device, ): @@ -95,7 +94,7 @@ def encode_image( image: PipelineImageInput, image_processor: CLIPImageProcessor, image_encoder: CLIPVisionModel, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): image = image_processor(images=image, return_tensors="pt").to(device) image_embeds = image_encoder(**image, output_hidden_states=True) @@ -104,7 +103,7 @@ def encode_image( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -164,7 +163,7 @@ def description(self) -> str: return "Text Encoder step that generate text_embeddings to guide the video generation" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("text_encoder", UMT5EncoderModel), ComponentSpec("tokenizer", AutoTokenizer), @@ -177,7 +176,7 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("prompt"), InputParam("negative_prompt"), @@ -185,7 +184,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "prompt_embeds", @@ -212,22 +211,22 @@ def check_inputs(block_state): def encode_prompt( components, prompt: str, - device: Optional[torch.device] = None, + device: torch.device | None = None, prepare_unconditional_embeds: bool = True, - negative_prompt: Optional[str] = None, + negative_prompt: str | None = None, max_sequence_length: int = 512, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device prepare_unconditional_embeds (`bool`): whether to use prepare unconditional embeddings or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -307,7 +306,7 @@ def description(self) -> str: return "Image Resize step that resize the image to the target area (height * width) while maintaining the aspect ratio." @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("image", type_hint=PIL.Image.Image, required=True), InputParam("height", type_hint=int, default=480), @@ -315,7 +314,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam("resized_image", type_hint=PIL.Image.Image), ] @@ -343,7 +342,7 @@ def description(self) -> str: return "Image Resize step that resize the last_image to the same size of first frame image with center crop." @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam( "resized_image", type_hint=PIL.Image.Image, required=True, description="The resized first frame image" @@ -352,7 +351,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam("resized_last_image", type_hint=PIL.Image.Image), ] @@ -386,20 +385,20 @@ def description(self) -> str: return "Image Encoder step that generate image_embeds based on first frame image to guide the video generation" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("image_processor", CLIPImageProcessor), ComponentSpec("image_encoder", CLIPVisionModel), ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("resized_image", type_hint=PIL.Image.Image, required=True), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam("image_embeds", type_hint=torch.Tensor, description="The image embeddings"), ] @@ -430,21 +429,21 @@ def description(self) -> str: return "Image Encoder step that generate image_embeds based on first and last frame images to guide the video generation" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("image_processor", CLIPImageProcessor), ComponentSpec("image_encoder", CLIPVisionModel), ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("resized_image", type_hint=PIL.Image.Image, required=True), InputParam("resized_last_image", type_hint=PIL.Image.Image, required=True), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam("image_embeds", type_hint=torch.Tensor, description="The image embeddings"), ] @@ -468,7 +467,7 @@ def __call__(self, components: WanModularPipeline, state: PipelineState) -> Pipe return components, state -class WanVaeImageEncoderStep(ModularPipelineBlocks): +class WanVaeEncoderStep(ModularPipelineBlocks): model_name = "wan" @property @@ -476,7 +475,7 @@ def description(self) -> str: return "Vae Image Encoder step that generate condition_latents based on first frame image to guide the video generation" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("vae", AutoencoderKLWan), ComponentSpec( @@ -488,17 +487,17 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("resized_image", type_hint=PIL.Image.Image, required=True), InputParam("height"), InputParam("width"), - InputParam("num_frames"), + InputParam("num_frames", type_hint=int, default=81), InputParam("generator"), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "first_frame_latents", @@ -564,7 +563,51 @@ def __call__(self, components: WanModularPipeline, state: PipelineState) -> Pipe return components, state -class WanFirstLastFrameVaeImageEncoderStep(ModularPipelineBlocks): +class WanPrepareFirstFrameLatentsStep(ModularPipelineBlocks): + model_name = "wan" + + @property + def description(self) -> str: + return "step that prepares the masked first frame latents and add it to the latent condition" + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("first_frame_latents", type_hint=torch.Tensor | None), + InputParam("num_frames", required=True), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("image_condition_latents", type_hint=torch.Tensor | None), + ] + + def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + batch_size, _, _, latent_height, latent_width = block_state.first_frame_latents.shape + + mask_lat_size = torch.ones(batch_size, 1, block_state.num_frames, latent_height, latent_width) + mask_lat_size[:, :, list(range(1, block_state.num_frames))] = 0 + + first_frame_mask = mask_lat_size[:, :, 0:1] + first_frame_mask = torch.repeat_interleave( + first_frame_mask, dim=2, repeats=components.vae_scale_factor_temporal + ) + mask_lat_size = torch.concat([first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2) + mask_lat_size = mask_lat_size.view( + batch_size, -1, components.vae_scale_factor_temporal, latent_height, latent_width + ) + mask_lat_size = mask_lat_size.transpose(1, 2) + mask_lat_size = mask_lat_size.to(block_state.first_frame_latents.device) + block_state.image_condition_latents = torch.concat([mask_lat_size, block_state.first_frame_latents], dim=1) + + self.set_block_state(state, block_state) + return components, state + + +class WanFirstLastFrameVaeEncoderStep(ModularPipelineBlocks): model_name = "wan" @property @@ -572,7 +615,7 @@ def description(self) -> str: return "Vae Image Encoder step that generate condition_latents based on first and last frame images to guide the video generation" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("vae", AutoencoderKLWan), ComponentSpec( @@ -584,18 +627,18 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("resized_image", type_hint=PIL.Image.Image, required=True), InputParam("resized_last_image", type_hint=PIL.Image.Image, required=True), InputParam("height"), InputParam("width"), - InputParam("num_frames"), + InputParam("num_frames", type_hint=int, default=81), InputParam("generator"), ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "first_last_frame_latents", @@ -667,3 +710,49 @@ def __call__(self, components: WanModularPipeline, state: PipelineState) -> Pipe self.set_block_state(state, block_state) return components, state + + +class WanPrepareFirstLastFrameLatentsStep(ModularPipelineBlocks): + model_name = "wan" + + @property + def description(self) -> str: + return "step that prepares the masked latents with first and last frames and add it to the latent condition" + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("first_last_frame_latents", type_hint=torch.Tensor | None), + InputParam("num_frames", type_hint=int, required=True), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("image_condition_latents", type_hint=torch.Tensor | None), + ] + + def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + batch_size, _, _, latent_height, latent_width = block_state.first_last_frame_latents.shape + + mask_lat_size = torch.ones(batch_size, 1, block_state.num_frames, latent_height, latent_width) + mask_lat_size[:, :, list(range(1, block_state.num_frames - 1))] = 0 + + first_frame_mask = mask_lat_size[:, :, 0:1] + first_frame_mask = torch.repeat_interleave( + first_frame_mask, dim=2, repeats=components.vae_scale_factor_temporal + ) + mask_lat_size = torch.concat([first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2) + mask_lat_size = mask_lat_size.view( + batch_size, -1, components.vae_scale_factor_temporal, latent_height, latent_width + ) + mask_lat_size = mask_lat_size.transpose(1, 2) + mask_lat_size = mask_lat_size.to(block_state.first_last_frame_latents.device) + block_state.image_condition_latents = torch.concat( + [mask_lat_size, block_state.first_last_frame_latents], dim=1 + ) + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/wan/modular_blocks.py b/src/diffusers/modular_pipelines/wan/modular_blocks.py deleted file mode 100644 index b3b70b2f9be1..000000000000 --- a/src/diffusers/modular_pipelines/wan/modular_blocks.py +++ /dev/null @@ -1,474 +0,0 @@ -# Copyright 2025 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from ...utils import logging -from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks -from ..modular_pipeline_utils import InsertableDict -from .before_denoise import ( - WanAdditionalInputsStep, - WanPrepareFirstFrameLatentsStep, - WanPrepareFirstLastFrameLatentsStep, - WanPrepareLatentsStep, - WanSetTimestepsStep, - WanTextInputStep, -) -from .decoders import WanImageVaeDecoderStep -from .denoise import ( - Wan22DenoiseStep, - Wan22Image2VideoDenoiseStep, - WanDenoiseStep, - WanFLF2VDenoiseStep, - WanImage2VideoDenoiseStep, -) -from .encoders import ( - WanFirstLastFrameImageEncoderStep, - WanFirstLastFrameVaeImageEncoderStep, - WanImageCropResizeStep, - WanImageEncoderStep, - WanImageResizeStep, - WanTextEncoderStep, - WanVaeImageEncoderStep, -) - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -# wan2.1 -# wan2.1: text2vid -class WanCoreDenoiseStep(SequentialPipelineBlocks): - block_classes = [ - WanTextInputStep, - WanSetTimestepsStep, - WanPrepareLatentsStep, - WanDenoiseStep, - ] - block_names = ["input", "set_timesteps", "prepare_latents", "denoise"] - - @property - def description(self): - return ( - "denoise block that takes encoded conditions and runs the denoising process.\n" - + "This is a sequential pipeline blocks:\n" - + " - `WanTextInputStep` is used to adjust the batch size of the model inputs\n" - + " - `WanSetTimestepsStep` is used to set the timesteps\n" - + " - `WanPrepareLatentsStep` is used to prepare the latents\n" - + " - `WanDenoiseStep` is used to denoise the latents\n" - ) - - -# wan2.1: image2video -## image encoder -class WanImage2VideoImageEncoderStep(SequentialPipelineBlocks): - model_name = "wan" - block_classes = [WanImageResizeStep, WanImageEncoderStep] - block_names = ["image_resize", "image_encoder"] - - @property - def description(self): - return "Image2Video Image Encoder step that resize the image and encode the image to generate the image embeddings" - - -## vae encoder -class WanImage2VideoVaeImageEncoderStep(SequentialPipelineBlocks): - model_name = "wan" - block_classes = [WanImageResizeStep, WanVaeImageEncoderStep] - block_names = ["image_resize", "vae_image_encoder"] - - @property - def description(self): - return "Image2Video Vae Image Encoder step that resize the image and encode the first frame image to its latent representation" - - -## denoise -class WanImage2VideoCoreDenoiseStep(SequentialPipelineBlocks): - block_classes = [ - WanTextInputStep, - WanAdditionalInputsStep(image_latent_inputs=["first_frame_latents"]), - WanSetTimestepsStep, - WanPrepareLatentsStep, - WanPrepareFirstFrameLatentsStep, - WanImage2VideoDenoiseStep, - ] - block_names = [ - "input", - "additional_inputs", - "set_timesteps", - "prepare_latents", - "prepare_first_frame_latents", - "denoise", - ] - - @property - def description(self): - return ( - "denoise block that takes encoded text and image latent conditions and runs the denoising process.\n" - + "This is a sequential pipeline blocks:\n" - + " - `WanTextInputStep` is used to adjust the batch size of the model inputs\n" - + " - `WanAdditionalInputsStep` is used to adjust the batch size of the latent conditions\n" - + " - `WanSetTimestepsStep` is used to set the timesteps\n" - + " - `WanPrepareLatentsStep` is used to prepare the latents\n" - + " - `WanPrepareFirstFrameLatentsStep` is used to prepare the first frame latent conditions\n" - + " - `WanImage2VideoDenoiseStep` is used to denoise the latents\n" - ) - - -# wan2.1: FLF2v - - -## image encoder -class WanFLF2VImageEncoderStep(SequentialPipelineBlocks): - model_name = "wan" - block_classes = [WanImageResizeStep, WanImageCropResizeStep, WanFirstLastFrameImageEncoderStep] - block_names = ["image_resize", "last_image_resize", "image_encoder"] - - @property - def description(self): - return "FLF2V Image Encoder step that resize and encode and encode the first and last frame images to generate the image embeddings" - - -## vae encoder -class WanFLF2VVaeImageEncoderStep(SequentialPipelineBlocks): - model_name = "wan" - block_classes = [WanImageResizeStep, WanImageCropResizeStep, WanFirstLastFrameVaeImageEncoderStep] - block_names = ["image_resize", "last_image_resize", "vae_image_encoder"] - - @property - def description(self): - return "FLF2V Vae Image Encoder step that resize and encode and encode the first and last frame images to generate the latent conditions" - - -## denoise -class WanFLF2VCoreDenoiseStep(SequentialPipelineBlocks): - block_classes = [ - WanTextInputStep, - WanAdditionalInputsStep(image_latent_inputs=["first_last_frame_latents"]), - WanSetTimestepsStep, - WanPrepareLatentsStep, - WanPrepareFirstLastFrameLatentsStep, - WanFLF2VDenoiseStep, - ] - block_names = [ - "input", - "additional_inputs", - "set_timesteps", - "prepare_latents", - "prepare_first_last_frame_latents", - "denoise", - ] - - @property - def description(self): - return ( - "denoise block that takes encoded text and image latent conditions and runs the denoising process.\n" - + "This is a sequential pipeline blocks:\n" - + " - `WanTextInputStep` is used to adjust the batch size of the model inputs\n" - + " - `WanAdditionalInputsStep` is used to adjust the batch size of the latent conditions\n" - + " - `WanSetTimestepsStep` is used to set the timesteps\n" - + " - `WanPrepareLatentsStep` is used to prepare the latents\n" - + " - `WanPrepareFirstLastFrameLatentsStep` is used to prepare the latent conditions\n" - + " - `WanImage2VideoDenoiseStep` is used to denoise the latents\n" - ) - - -# wan2.1: auto blocks -## image encoder -class WanAutoImageEncoderStep(AutoPipelineBlocks): - block_classes = [WanFLF2VImageEncoderStep, WanImage2VideoImageEncoderStep] - block_names = ["flf2v_image_encoder", "image2video_image_encoder"] - block_trigger_inputs = ["last_image", "image"] - - @property - def description(self): - return ( - "Image Encoder step that encode the image to generate the image embeddings" - + "This is an auto pipeline block that works for image2video tasks." - + " - `WanFLF2VImageEncoderStep` (flf2v) is used when `last_image` is provided." - + " - `WanImage2VideoImageEncoderStep` (image2video) is used when `image` is provided." - + " - if `last_image` or `image` is not provided, step will be skipped." - ) - - -## vae encoder -class WanAutoVaeImageEncoderStep(AutoPipelineBlocks): - block_classes = [WanFLF2VVaeImageEncoderStep, WanImage2VideoVaeImageEncoderStep] - block_names = ["flf2v_vae_image_encoder", "image2video_vae_image_encoder"] - block_trigger_inputs = ["last_image", "image"] - - @property - def description(self): - return ( - "Vae Image Encoder step that encode the image to generate the image latents" - + "This is an auto pipeline block that works for image2video tasks." - + " - `WanFLF2VVaeImageEncoderStep` (flf2v) is used when `last_image` is provided." - + " - `WanImage2VideoVaeImageEncoderStep` (image2video) is used when `image` is provided." - + " - if `last_image` or `image` is not provided, step will be skipped." - ) - - -## denoise -class WanAutoDenoiseStep(AutoPipelineBlocks): - block_classes = [ - WanFLF2VCoreDenoiseStep, - WanImage2VideoCoreDenoiseStep, - WanCoreDenoiseStep, - ] - block_names = ["flf2v", "image2video", "text2video"] - block_trigger_inputs = ["first_last_frame_latents", "first_frame_latents", None] - - @property - def description(self) -> str: - return ( - "Denoise step that iteratively denoise the latents. " - "This is a auto pipeline block that works for text2video and image2video tasks." - " - `WanCoreDenoiseStep` (text2video) for text2vid tasks." - " - `WanCoreImage2VideoCoreDenoiseStep` (image2video) for image2video tasks." - + " - if `first_frame_latents` is provided, `WanCoreImage2VideoDenoiseStep` will be used.\n" - + " - if `first_frame_latents` is not provided, `WanCoreDenoiseStep` will be used.\n" - ) - - -# auto pipeline blocks -class WanAutoBlocks(SequentialPipelineBlocks): - block_classes = [ - WanTextEncoderStep, - WanAutoImageEncoderStep, - WanAutoVaeImageEncoderStep, - WanAutoDenoiseStep, - WanImageVaeDecoderStep, - ] - block_names = [ - "text_encoder", - "image_encoder", - "vae_image_encoder", - "denoise", - "decode", - ] - - @property - def description(self): - return ( - "Auto Modular pipeline for text-to-video using Wan.\n" - + "- for text-to-video generation, all you need to provide is `prompt`" - ) - - -# wan22 -# wan2.2: text2vid - - -## denoise -class Wan22CoreDenoiseStep(SequentialPipelineBlocks): - block_classes = [ - WanTextInputStep, - WanSetTimestepsStep, - WanPrepareLatentsStep, - Wan22DenoiseStep, - ] - block_names = ["input", "set_timesteps", "prepare_latents", "denoise"] - - @property - def description(self): - return ( - "denoise block that takes encoded conditions and runs the denoising process.\n" - + "This is a sequential pipeline blocks:\n" - + " - `WanTextInputStep` is used to adjust the batch size of the model inputs\n" - + " - `WanSetTimestepsStep` is used to set the timesteps\n" - + " - `WanPrepareLatentsStep` is used to prepare the latents\n" - + " - `Wan22DenoiseStep` is used to denoise the latents in wan2.2\n" - ) - - -# wan2.2: image2video -## denoise -class Wan22Image2VideoCoreDenoiseStep(SequentialPipelineBlocks): - block_classes = [ - WanTextInputStep, - WanAdditionalInputsStep(image_latent_inputs=["first_frame_latents"]), - WanSetTimestepsStep, - WanPrepareLatentsStep, - WanPrepareFirstFrameLatentsStep, - Wan22Image2VideoDenoiseStep, - ] - block_names = [ - "input", - "additional_inputs", - "set_timesteps", - "prepare_latents", - "prepare_first_frame_latents", - "denoise", - ] - - @property - def description(self): - return ( - "denoise block that takes encoded text and image latent conditions and runs the denoising process.\n" - + "This is a sequential pipeline blocks:\n" - + " - `WanTextInputStep` is used to adjust the batch size of the model inputs\n" - + " - `WanAdditionalInputsStep` is used to adjust the batch size of the latent conditions\n" - + " - `WanSetTimestepsStep` is used to set the timesteps\n" - + " - `WanPrepareLatentsStep` is used to prepare the latents\n" - + " - `WanPrepareFirstFrameLatentsStep` is used to prepare the first frame latent conditions\n" - + " - `Wan22Image2VideoDenoiseStep` is used to denoise the latents in wan2.2\n" - ) - - -class Wan22AutoDenoiseStep(AutoPipelineBlocks): - block_classes = [ - Wan22Image2VideoCoreDenoiseStep, - Wan22CoreDenoiseStep, - ] - block_names = ["image2video", "text2video"] - block_trigger_inputs = ["first_frame_latents", None] - - @property - def description(self) -> str: - return ( - "Denoise step that iteratively denoise the latents. " - "This is a auto pipeline block that works for text2video and image2video tasks." - " - `Wan22Image2VideoCoreDenoiseStep` (image2video) for image2video tasks." - " - `Wan22CoreDenoiseStep` (text2video) for text2vid tasks." - + " - if `first_frame_latents` is provided, `Wan22Image2VideoCoreDenoiseStep` will be used.\n" - + " - if `first_frame_latents` is not provided, `Wan22CoreDenoiseStep` will be used.\n" - ) - - -class Wan22AutoBlocks(SequentialPipelineBlocks): - block_classes = [ - WanTextEncoderStep, - WanAutoVaeImageEncoderStep, - Wan22AutoDenoiseStep, - WanImageVaeDecoderStep, - ] - block_names = [ - "text_encoder", - "vae_image_encoder", - "denoise", - "decode", - ] - - @property - def description(self): - return ( - "Auto Modular pipeline for text-to-video using Wan2.2.\n" - + "- for text-to-video generation, all you need to provide is `prompt`" - ) - - -# presets for wan2.1 and wan2.2 -# YiYi Notes: should we move these to doc? -# wan2.1 -TEXT2VIDEO_BLOCKS = InsertableDict( - [ - ("text_encoder", WanTextEncoderStep), - ("input", WanTextInputStep), - ("set_timesteps", WanSetTimestepsStep), - ("prepare_latents", WanPrepareLatentsStep), - ("denoise", WanDenoiseStep), - ("decode", WanImageVaeDecoderStep), - ] -) - -IMAGE2VIDEO_BLOCKS = InsertableDict( - [ - ("image_resize", WanImageResizeStep), - ("image_encoder", WanImage2VideoImageEncoderStep), - ("vae_image_encoder", WanImage2VideoVaeImageEncoderStep), - ("input", WanTextInputStep), - ("additional_inputs", WanAdditionalInputsStep(image_latent_inputs=["first_frame_latents"])), - ("set_timesteps", WanSetTimestepsStep), - ("prepare_latents", WanPrepareLatentsStep), - ("prepare_first_frame_latents", WanPrepareFirstFrameLatentsStep), - ("denoise", WanImage2VideoDenoiseStep), - ("decode", WanImageVaeDecoderStep), - ] -) - - -FLF2V_BLOCKS = InsertableDict( - [ - ("image_resize", WanImageResizeStep), - ("last_image_resize", WanImageCropResizeStep), - ("image_encoder", WanFLF2VImageEncoderStep), - ("vae_image_encoder", WanFLF2VVaeImageEncoderStep), - ("input", WanTextInputStep), - ("additional_inputs", WanAdditionalInputsStep(image_latent_inputs=["first_last_frame_latents"])), - ("set_timesteps", WanSetTimestepsStep), - ("prepare_latents", WanPrepareLatentsStep), - ("prepare_first_last_frame_latents", WanPrepareFirstLastFrameLatentsStep), - ("denoise", WanFLF2VDenoiseStep), - ("decode", WanImageVaeDecoderStep), - ] -) - -AUTO_BLOCKS = InsertableDict( - [ - ("text_encoder", WanTextEncoderStep), - ("image_encoder", WanAutoImageEncoderStep), - ("vae_image_encoder", WanAutoVaeImageEncoderStep), - ("denoise", WanAutoDenoiseStep), - ("decode", WanImageVaeDecoderStep), - ] -) - -# wan2.2 presets - -TEXT2VIDEO_BLOCKS_WAN22 = InsertableDict( - [ - ("text_encoder", WanTextEncoderStep), - ("input", WanTextInputStep), - ("set_timesteps", WanSetTimestepsStep), - ("prepare_latents", WanPrepareLatentsStep), - ("denoise", Wan22DenoiseStep), - ("decode", WanImageVaeDecoderStep), - ] -) - -IMAGE2VIDEO_BLOCKS_WAN22 = InsertableDict( - [ - ("image_resize", WanImageResizeStep), - ("vae_image_encoder", WanImage2VideoVaeImageEncoderStep), - ("input", WanTextInputStep), - ("set_timesteps", WanSetTimestepsStep), - ("prepare_latents", WanPrepareLatentsStep), - ("denoise", Wan22DenoiseStep), - ("decode", WanImageVaeDecoderStep), - ] -) - -AUTO_BLOCKS_WAN22 = InsertableDict( - [ - ("text_encoder", WanTextEncoderStep), - ("vae_image_encoder", WanAutoVaeImageEncoderStep), - ("denoise", Wan22AutoDenoiseStep), - ("decode", WanImageVaeDecoderStep), - ] -) - -# presets all blocks (wan and wan22) - - -ALL_BLOCKS = { - "wan2.1": { - "text2video": TEXT2VIDEO_BLOCKS, - "image2video": IMAGE2VIDEO_BLOCKS, - "flf2v": FLF2V_BLOCKS, - "auto": AUTO_BLOCKS, - }, - "wan2.2": { - "text2video": TEXT2VIDEO_BLOCKS_WAN22, - "image2video": IMAGE2VIDEO_BLOCKS_WAN22, - "auto": AUTO_BLOCKS_WAN22, - }, -} diff --git a/src/diffusers/modular_pipelines/wan/modular_blocks_wan.py b/src/diffusers/modular_pipelines/wan/modular_blocks_wan.py new file mode 100644 index 000000000000..b641c6cd7fcc --- /dev/null +++ b/src/diffusers/modular_pipelines/wan/modular_blocks_wan.py @@ -0,0 +1,162 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...utils import logging +from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline_utils import OutputParam +from .before_denoise import ( + WanPrepareLatentsStep, + WanSetTimestepsStep, + WanTextInputStep, +) +from .decoders import WanVaeDecoderStep +from .denoise import ( + WanDenoiseStep, +) +from .encoders import ( + WanTextEncoderStep, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# ==================== +# 1. DENOISE +# ==================== + + +# inputs(text) -> set_timesteps -> prepare_latents -> denoise +# auto_docstring +class WanCoreDenoiseStep(SequentialPipelineBlocks): + """ + denoise block that takes encoded conditions and runs the denoising process. + + Components: + transformer (`WanTransformer3DModel`) scheduler (`UniPCMultistepScheduler`) guider (`ClassifierFreeGuidance`) + + Inputs: + num_videos_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + Pre-generated negative text embeddings. Can be generated from text_encoder step. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + num_frames (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + attention_kwargs (`None`, *optional*): + TODO: Add description. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "wan" + block_classes = [ + WanTextInputStep, + WanSetTimestepsStep, + WanPrepareLatentsStep, + WanDenoiseStep, + ] + block_names = ["input", "set_timesteps", "prepare_latents", "denoise"] + + @property + def description(self): + return "denoise block that takes encoded conditions and runs the denoising process." + + @property + def outputs(self): + return [OutputParam.template("latents")] + + +# ==================== +# 2. BLOCKS (Wan2.1 text2video) +# ==================== + + +# auto_docstring +class WanBlocks(SequentialPipelineBlocks): + """ + Modular pipeline blocks for Wan2.1. + + Components: + text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) guider (`ClassifierFreeGuidance`) transformer + (`WanTransformer3DModel`) scheduler (`UniPCMultistepScheduler`) vae (`AutoencoderKLWan`) video_processor + (`VideoProcessor`) + + Inputs: + prompt (`None`, *optional*): + TODO: Add description. + negative_prompt (`None`, *optional*): + TODO: Add description. + max_sequence_length (`None`, *optional*, defaults to 512): + TODO: Add description. + num_videos_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + num_frames (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + attention_kwargs (`None`, *optional*): + TODO: Add description. + output_type (`str`, *optional*, defaults to np): + The output type of the decoded videos + + Outputs: + videos (`list`): + The generated videos. + """ + + model_name = "wan" + block_classes = [ + WanTextEncoderStep, + WanCoreDenoiseStep, + WanVaeDecoderStep, + ] + block_names = ["text_encoder", "denoise", "decode"] + + @property + def description(self): + return "Modular pipeline blocks for Wan2.1." + + @property + def outputs(self): + return [OutputParam.template("videos")] diff --git a/src/diffusers/modular_pipelines/wan/modular_blocks_wan22.py b/src/diffusers/modular_pipelines/wan/modular_blocks_wan22.py new file mode 100644 index 000000000000..9f602c24713b --- /dev/null +++ b/src/diffusers/modular_pipelines/wan/modular_blocks_wan22.py @@ -0,0 +1,176 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...utils import logging +from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline_utils import OutputParam +from .before_denoise import ( + WanPrepareLatentsStep, + WanSetTimestepsStep, + WanTextInputStep, +) +from .decoders import WanVaeDecoderStep +from .denoise import ( + Wan22DenoiseStep, +) +from .encoders import ( + WanTextEncoderStep, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# ==================== +# 1. DENOISE +# ==================== + +# inputs(text) -> set_timesteps -> prepare_latents -> denoise + + +# auto_docstring +class Wan22CoreDenoiseStep(SequentialPipelineBlocks): + """ + denoise block that takes encoded conditions and runs the denoising process. + + Components: + transformer (`WanTransformer3DModel`) scheduler (`UniPCMultistepScheduler`) guider (`ClassifierFreeGuidance`) + guider_2 (`ClassifierFreeGuidance`) transformer_2 (`WanTransformer3DModel`) + + Configs: + boundary_ratio (default: 0.875): The boundary ratio to divide the denoising loop into high noise and low + noise stages. + + Inputs: + num_videos_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + Pre-generated negative text embeddings. Can be generated from text_encoder step. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + num_frames (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + attention_kwargs (`None`, *optional*): + TODO: Add description. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "wan" + block_classes = [ + WanTextInputStep, + WanSetTimestepsStep, + WanPrepareLatentsStep, + Wan22DenoiseStep, + ] + block_names = ["input", "set_timesteps", "prepare_latents", "denoise"] + + @property + def description(self): + return "denoise block that takes encoded conditions and runs the denoising process." + + @property + def outputs(self): + return [OutputParam.template("latents")] + + +# ==================== +# 2. BLOCKS (Wan2.2 text2video) +# ==================== + + +# auto_docstring +class Wan22Blocks(SequentialPipelineBlocks): + """ + Modular pipeline for text-to-video using Wan2.2. + + Components: + text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) guider (`ClassifierFreeGuidance`) transformer + (`WanTransformer3DModel`) scheduler (`UniPCMultistepScheduler`) guider_2 (`ClassifierFreeGuidance`) + transformer_2 (`WanTransformer3DModel`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) + + Configs: + boundary_ratio (default: 0.875): The boundary ratio to divide the denoising loop into high noise and low + noise stages. + + Inputs: + prompt (`None`, *optional*): + TODO: Add description. + negative_prompt (`None`, *optional*): + TODO: Add description. + max_sequence_length (`None`, *optional*, defaults to 512): + TODO: Add description. + num_videos_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + num_frames (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + attention_kwargs (`None`, *optional*): + TODO: Add description. + output_type (`str`, *optional*, defaults to np): + The output type of the decoded videos + + Outputs: + videos (`list`): + The generated videos. + """ + + model_name = "wan" + block_classes = [ + WanTextEncoderStep, + Wan22CoreDenoiseStep, + WanVaeDecoderStep, + ] + block_names = [ + "text_encoder", + "denoise", + "decode", + ] + + @property + def description(self): + return "Modular pipeline for text-to-video using Wan2.2." + + @property + def outputs(self): + return [OutputParam.template("videos")] diff --git a/src/diffusers/modular_pipelines/wan/modular_blocks_wan22_i2v.py b/src/diffusers/modular_pipelines/wan/modular_blocks_wan22_i2v.py new file mode 100644 index 000000000000..8e55b7a50f08 --- /dev/null +++ b/src/diffusers/modular_pipelines/wan/modular_blocks_wan22_i2v.py @@ -0,0 +1,236 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...utils import logging +from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline_utils import OutputParam +from .before_denoise import ( + WanAdditionalInputsStep, + WanPrepareLatentsStep, + WanSetTimestepsStep, + WanTextInputStep, +) +from .decoders import WanVaeDecoderStep +from .denoise import ( + Wan22Image2VideoDenoiseStep, +) +from .encoders import ( + WanImageResizeStep, + WanPrepareFirstFrameLatentsStep, + WanTextEncoderStep, + WanVaeEncoderStep, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# ==================== +# 1. VAE ENCODER +# ==================== + + +# auto_docstring +class WanImage2VideoVaeEncoderStep(SequentialPipelineBlocks): + """ + Image2Video Vae Image Encoder step that resize the image and encode the first frame image to its latent + representation + + Components: + vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) + + Inputs: + image (`Image`): + TODO: Add description. + height (`int`, *optional*, defaults to 480): + TODO: Add description. + width (`int`, *optional*, defaults to 832): + TODO: Add description. + num_frames (`int`, *optional*, defaults to 81): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + resized_image (`Image`): + TODO: Add description. + first_frame_latents (`Tensor`): + video latent representation with the first frame image condition + image_condition_latents (`Tensor | NoneType`): + TODO: Add description. + """ + + model_name = "wan-i2v" + block_classes = [WanImageResizeStep, WanVaeEncoderStep, WanPrepareFirstFrameLatentsStep] + block_names = ["image_resize", "vae_encoder", "prepare_first_frame_latents"] + + @property + def description(self): + return "Image2Video Vae Image Encoder step that resize the image and encode the first frame image to its latent representation" + + +# ==================== +# 2. DENOISE +# ==================== + + +# inputs (text + image_condition_latents) -> set_timesteps -> prepare_latents -> denoise (latents) +# auto_docstring +class Wan22Image2VideoCoreDenoiseStep(SequentialPipelineBlocks): + """ + denoise block that takes encoded text and image latent conditions and runs the denoising process. + + Components: + transformer (`WanTransformer3DModel`) scheduler (`UniPCMultistepScheduler`) guider (`ClassifierFreeGuidance`) + guider_2 (`ClassifierFreeGuidance`) transformer_2 (`WanTransformer3DModel`) + + Configs: + boundary_ratio (default: 0.875): The boundary ratio to divide the denoising loop into high noise and low + noise stages. + + Inputs: + num_videos_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + Pre-generated negative text embeddings. Can be generated from text_encoder step. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + num_frames (`None`, *optional*): + TODO: Add description. + image_condition_latents (`None`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + attention_kwargs (`None`, *optional*): + TODO: Add description. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "wan-i2v" + block_classes = [ + WanTextInputStep, + WanAdditionalInputsStep(image_latent_inputs=["image_condition_latents"]), + WanSetTimestepsStep, + WanPrepareLatentsStep, + Wan22Image2VideoDenoiseStep, + ] + block_names = [ + "input", + "additional_inputs", + "set_timesteps", + "prepare_latents", + "denoise", + ] + + @property + def description(self): + return "denoise block that takes encoded text and image latent conditions and runs the denoising process." + + @property + def outputs(self): + return [OutputParam.template("latents")] + + +# ==================== +# 3. BLOCKS (Wan2.2 Image2Video) +# ==================== + + +# auto_docstring +class Wan22Image2VideoBlocks(SequentialPipelineBlocks): + """ + Modular pipeline for image-to-video using Wan2.2. + + Components: + text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) guider (`ClassifierFreeGuidance`) vae + (`AutoencoderKLWan`) video_processor (`VideoProcessor`) transformer (`WanTransformer3DModel`) scheduler + (`UniPCMultistepScheduler`) guider_2 (`ClassifierFreeGuidance`) transformer_2 (`WanTransformer3DModel`) + + Configs: + boundary_ratio (default: 0.875): The boundary ratio to divide the denoising loop into high noise and low + noise stages. + + Inputs: + prompt (`None`, *optional*): + TODO: Add description. + negative_prompt (`None`, *optional*): + TODO: Add description. + max_sequence_length (`None`, *optional*, defaults to 512): + TODO: Add description. + image (`Image`): + TODO: Add description. + height (`int`, *optional*, defaults to 480): + TODO: Add description. + width (`int`, *optional*, defaults to 832): + TODO: Add description. + num_frames (`int`, *optional*, defaults to 81): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_videos_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + attention_kwargs (`None`, *optional*): + TODO: Add description. + output_type (`str`, *optional*, defaults to np): + The output type of the decoded videos + + Outputs: + videos (`list`): + The generated videos. + """ + + model_name = "wan-i2v" + block_classes = [ + WanTextEncoderStep, + WanImage2VideoVaeEncoderStep, + Wan22Image2VideoCoreDenoiseStep, + WanVaeDecoderStep, + ] + block_names = [ + "text_encoder", + "vae_encoder", + "denoise", + "decode", + ] + + @property + def description(self): + return "Modular pipeline for image-to-video using Wan2.2." + + @property + def outputs(self): + return [OutputParam.template("videos")] diff --git a/src/diffusers/modular_pipelines/wan/modular_blocks_wan_i2v.py b/src/diffusers/modular_pipelines/wan/modular_blocks_wan_i2v.py new file mode 100644 index 000000000000..c08db62c469a --- /dev/null +++ b/src/diffusers/modular_pipelines/wan/modular_blocks_wan_i2v.py @@ -0,0 +1,481 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...utils import logging +from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import OutputParam +from .before_denoise import ( + WanAdditionalInputsStep, + WanPrepareLatentsStep, + WanSetTimestepsStep, + WanTextInputStep, +) +from .decoders import WanVaeDecoderStep +from .denoise import ( + WanImage2VideoDenoiseStep, +) +from .encoders import ( + WanFirstLastFrameImageEncoderStep, + WanFirstLastFrameVaeEncoderStep, + WanImageCropResizeStep, + WanImageEncoderStep, + WanImageResizeStep, + WanPrepareFirstFrameLatentsStep, + WanPrepareFirstLastFrameLatentsStep, + WanTextEncoderStep, + WanVaeEncoderStep, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +# ==================== +# 1. IMAGE ENCODER +# ==================== + + +# wan2.1 I2V (first frame only) +# auto_docstring +class WanImage2VideoImageEncoderStep(SequentialPipelineBlocks): + """ + Image2Video Image Encoder step that resize the image and encode the image to generate the image embeddings + + Components: + image_processor (`CLIPImageProcessor`) image_encoder (`CLIPVisionModel`) + + Inputs: + image (`Image`): + TODO: Add description. + height (`int`, *optional*, defaults to 480): + TODO: Add description. + width (`int`, *optional*, defaults to 832): + TODO: Add description. + + Outputs: + resized_image (`Image`): + TODO: Add description. + image_embeds (`Tensor`): + The image embeddings + """ + + model_name = "wan-i2v" + block_classes = [WanImageResizeStep, WanImageEncoderStep] + block_names = ["image_resize", "image_encoder"] + + @property + def description(self): + return "Image2Video Image Encoder step that resize the image and encode the image to generate the image embeddings" + + +# wan2.1 FLF2V (first and last frame) +# auto_docstring +class WanFLF2VImageEncoderStep(SequentialPipelineBlocks): + """ + FLF2V Image Encoder step that resize and encode and encode the first and last frame images to generate the image + embeddings + + Components: + image_processor (`CLIPImageProcessor`) image_encoder (`CLIPVisionModel`) + + Inputs: + image (`Image`): + TODO: Add description. + height (`int`, *optional*, defaults to 480): + TODO: Add description. + width (`int`, *optional*, defaults to 832): + TODO: Add description. + last_image (`Image`): + The last frameimage + + Outputs: + resized_image (`Image`): + TODO: Add description. + resized_last_image (`Image`): + TODO: Add description. + image_embeds (`Tensor`): + The image embeddings + """ + + model_name = "wan-i2v" + block_classes = [WanImageResizeStep, WanImageCropResizeStep, WanFirstLastFrameImageEncoderStep] + block_names = ["image_resize", "last_image_resize", "image_encoder"] + + @property + def description(self): + return "FLF2V Image Encoder step that resize and encode and encode the first and last frame images to generate the image embeddings" + + +# wan2.1 Auto Image Encoder +# auto_docstring +class WanAutoImageEncoderStep(AutoPipelineBlocks): + """ + Image Encoder step that encode the image to generate the image embeddingsThis is an auto pipeline block that works + for image2video tasks. - `WanFLF2VImageEncoderStep` (flf2v) is used when `last_image` is provided. - + `WanImage2VideoImageEncoderStep` (image2video) is used when `image` is provided. - if `last_image` or `image` is + not provided, step will be skipped. + + Components: + image_processor (`CLIPImageProcessor`) image_encoder (`CLIPVisionModel`) + + Inputs: + image (`Image`, *optional*): + TODO: Add description. + height (`int`, *optional*, defaults to 480): + TODO: Add description. + width (`int`, *optional*, defaults to 832): + TODO: Add description. + last_image (`Image`, *optional*): + The last frameimage + + Outputs: + resized_image (`Image`): + TODO: Add description. + resized_last_image (`Image`): + TODO: Add description. + image_embeds (`Tensor`): + The image embeddings + """ + + block_classes = [WanFLF2VImageEncoderStep, WanImage2VideoImageEncoderStep] + block_names = ["flf2v_image_encoder", "image2video_image_encoder"] + block_trigger_inputs = ["last_image", "image"] + model_name = "wan-i2v" + + @property + def description(self): + return ( + "Image Encoder step that encode the image to generate the image embeddings" + + "This is an auto pipeline block that works for image2video tasks." + + " - `WanFLF2VImageEncoderStep` (flf2v) is used when `last_image` is provided." + + " - `WanImage2VideoImageEncoderStep` (image2video) is used when `image` is provided." + + " - if `last_image` or `image` is not provided, step will be skipped." + ) + + +# ==================== +# 2. VAE ENCODER +# ==================== + + +# wan2.1 I2V (first frame only) +# auto_docstring +class WanImage2VideoVaeEncoderStep(SequentialPipelineBlocks): + """ + Image2Video Vae Image Encoder step that resize the image and encode the first frame image to its latent + representation + + Components: + vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) + + Inputs: + image (`Image`): + TODO: Add description. + height (`int`, *optional*, defaults to 480): + TODO: Add description. + width (`int`, *optional*, defaults to 832): + TODO: Add description. + num_frames (`int`, *optional*, defaults to 81): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + resized_image (`Image`): + TODO: Add description. + first_frame_latents (`Tensor`): + video latent representation with the first frame image condition + image_condition_latents (`Tensor | NoneType`): + TODO: Add description. + """ + + model_name = "wan-i2v" + block_classes = [WanImageResizeStep, WanVaeEncoderStep, WanPrepareFirstFrameLatentsStep] + block_names = ["image_resize", "vae_encoder", "prepare_first_frame_latents"] + + @property + def description(self): + return "Image2Video Vae Image Encoder step that resize the image and encode the first frame image to its latent representation" + + +# wan2.1 FLF2V (first and last frame) +# auto_docstring +class WanFLF2VVaeEncoderStep(SequentialPipelineBlocks): + """ + FLF2V Vae Image Encoder step that resize and encode and encode the first and last frame images to generate the + latent conditions + + Components: + vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) + + Inputs: + image (`Image`): + TODO: Add description. + height (`int`, *optional*, defaults to 480): + TODO: Add description. + width (`int`, *optional*, defaults to 832): + TODO: Add description. + last_image (`Image`): + The last frameimage + num_frames (`int`, *optional*, defaults to 81): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + resized_image (`Image`): + TODO: Add description. + resized_last_image (`Image`): + TODO: Add description. + first_last_frame_latents (`Tensor`): + video latent representation with the first and last frame images condition + image_condition_latents (`Tensor | NoneType`): + TODO: Add description. + """ + + model_name = "wan-i2v" + block_classes = [ + WanImageResizeStep, + WanImageCropResizeStep, + WanFirstLastFrameVaeEncoderStep, + WanPrepareFirstLastFrameLatentsStep, + ] + block_names = ["image_resize", "last_image_resize", "vae_encoder", "prepare_first_last_frame_latents"] + + @property + def description(self): + return "FLF2V Vae Image Encoder step that resize and encode and encode the first and last frame images to generate the latent conditions" + + +# wan2.1 Auto Vae Encoder +# auto_docstring +class WanAutoVaeEncoderStep(AutoPipelineBlocks): + """ + Vae Image Encoder step that encode the image to generate the image latentsThis is an auto pipeline block that works + for image2video tasks. - `WanFLF2VVaeEncoderStep` (flf2v) is used when `last_image` is provided. - + `WanImage2VideoVaeEncoderStep` (image2video) is used when `image` is provided. - if `last_image` or `image` is not + provided, step will be skipped. + + Components: + vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) + + Inputs: + image (`Image`, *optional*): + TODO: Add description. + height (`int`, *optional*, defaults to 480): + TODO: Add description. + width (`int`, *optional*, defaults to 832): + TODO: Add description. + last_image (`Image`, *optional*): + The last frameimage + num_frames (`int`, *optional*, defaults to 81): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + resized_image (`Image`): + TODO: Add description. + resized_last_image (`Image`): + TODO: Add description. + first_last_frame_latents (`Tensor`): + video latent representation with the first and last frame images condition + image_condition_latents (`Tensor | NoneType`): + TODO: Add description. + first_frame_latents (`Tensor`): + video latent representation with the first frame image condition + """ + + model_name = "wan-i2v" + block_classes = [WanFLF2VVaeEncoderStep, WanImage2VideoVaeEncoderStep] + block_names = ["flf2v_vae_encoder", "image2video_vae_encoder"] + block_trigger_inputs = ["last_image", "image"] + + @property + def description(self): + return ( + "Vae Image Encoder step that encode the image to generate the image latents" + + "This is an auto pipeline block that works for image2video tasks." + + " - `WanFLF2VVaeEncoderStep` (flf2v) is used when `last_image` is provided." + + " - `WanImage2VideoVaeEncoderStep` (image2video) is used when `image` is provided." + + " - if `last_image` or `image` is not provided, step will be skipped." + ) + + +# ==================== +# 3. DENOISE (inputs -> set_timesteps -> prepare_latents -> denoise) +# ==================== + + +# wan2.1 I2V core denoise (support both I2V and FLF2V) +# inputs (text + image_condition_latents) -> set_timesteps -> prepare_latents -> denoise (latents) +# auto_docstring +class WanImage2VideoCoreDenoiseStep(SequentialPipelineBlocks): + """ + denoise block that takes encoded text and image latent conditions and runs the denoising process. + + Components: + transformer (`WanTransformer3DModel`) scheduler (`UniPCMultistepScheduler`) guider (`ClassifierFreeGuidance`) + + Inputs: + num_videos_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + Pre-generated negative text embeddings. Can be generated from text_encoder step. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + num_frames (`None`, *optional*): + TODO: Add description. + image_condition_latents (`None`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + attention_kwargs (`None`, *optional*): + TODO: Add description. + image_embeds (`Tensor`): + TODO: Add description. + + Outputs: + batch_size (`int`): + Number of prompts, the final batch size of model inputs should be batch_size * num_videos_per_prompt + dtype (`dtype`): + Data type of model tensor inputs (determined by `transformer.dtype`) + latents (`Tensor`): + The initial latents to use for the denoising process + """ + + model_name = "wan-i2v" + block_classes = [ + WanTextInputStep, + WanAdditionalInputsStep(image_latent_inputs=["image_condition_latents"]), + WanSetTimestepsStep, + WanPrepareLatentsStep, + WanImage2VideoDenoiseStep, + ] + block_names = [ + "input", + "additional_inputs", + "set_timesteps", + "prepare_latents", + "denoise", + ] + + @property + def description(self): + return "denoise block that takes encoded text and image latent conditions and runs the denoising process." + + +# ==================== +# 4. BLOCKS (Wan2.1 Image2Video) +# ==================== + + +# wan2.1 Image2Video Auto Blocks +# auto_docstring +class WanImage2VideoAutoBlocks(SequentialPipelineBlocks): + """ + Auto Modular pipeline for image-to-video using Wan. + + Supported workflows: + - `image2video`: requires `image`, `prompt` + - `flf2v`: requires `last_image`, `image`, `prompt` + + Components: + text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) guider (`ClassifierFreeGuidance`) + image_processor (`CLIPImageProcessor`) image_encoder (`CLIPVisionModel`) vae (`AutoencoderKLWan`) + video_processor (`VideoProcessor`) transformer (`WanTransformer3DModel`) scheduler + (`UniPCMultistepScheduler`) + + Inputs: + prompt (`None`, *optional*): + TODO: Add description. + negative_prompt (`None`, *optional*): + TODO: Add description. + max_sequence_length (`None`, *optional*, defaults to 512): + TODO: Add description. + image (`Image`, *optional*): + TODO: Add description. + height (`int`, *optional*, defaults to 480): + TODO: Add description. + width (`int`, *optional*, defaults to 832): + TODO: Add description. + last_image (`Image`, *optional*): + The last frameimage + num_frames (`int`, *optional*, defaults to 81): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_videos_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + image_condition_latents (`None`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + attention_kwargs (`None`, *optional*): + TODO: Add description. + image_embeds (`Tensor`): + TODO: Add description. + output_type (`str`, *optional*, defaults to np): + The output type of the decoded videos + + Outputs: + videos (`list`): + The generated videos. + """ + + model_name = "wan-i2v" + block_classes = [ + WanTextEncoderStep, + WanAutoImageEncoderStep, + WanAutoVaeEncoderStep, + WanImage2VideoCoreDenoiseStep, + WanVaeDecoderStep, + ] + block_names = [ + "text_encoder", + "image_encoder", + "vae_encoder", + "denoise", + "decode", + ] + + _workflow_map = { + "image2video": {"image": True, "prompt": True}, + "flf2v": {"last_image": True, "image": True, "prompt": True}, + } + + @property + def description(self): + return "Auto Modular pipeline for image-to-video using Wan." + + @property + def outputs(self): + return [OutputParam.template("videos")] diff --git a/src/diffusers/modular_pipelines/wan/modular_pipeline.py b/src/diffusers/modular_pipelines/wan/modular_pipeline.py index 930b25e4b905..0e52026a51bf 100644 --- a/src/diffusers/modular_pipelines/wan/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/wan/modular_pipeline.py @@ -13,8 +13,6 @@ # limitations under the License. -from typing import Any, Dict, Optional - from ...loaders import WanLoraLoaderMixin from ...pipelines.pipeline_utils import StableDiffusionMixin from ...utils import logging @@ -30,19 +28,12 @@ class WanModularPipeline( WanLoraLoaderMixin, ): """ - A ModularPipeline for Wan. + A ModularPipeline for Wan2.1 text2video. > [!WARNING] > This is an experimental feature and is likely to change in the future. """ - default_blocks_name = "WanAutoBlocks" - - # override the default_blocks_name in base class, which is just return self.default_blocks_name - def get_default_blocks_name(self, config_dict: Optional[Dict[str, Any]]) -> Optional[str]: - if config_dict is not None and "boundary_ratio" in config_dict and config_dict["boundary_ratio"] is not None: - return "Wan22AutoBlocks" - else: - return "WanAutoBlocks" + default_blocks_name = "WanBlocks" @property def default_height(self): @@ -118,3 +109,33 @@ def num_train_timesteps(self): if hasattr(self, "scheduler") and self.scheduler is not None: num_train_timesteps = self.scheduler.config.num_train_timesteps return num_train_timesteps + + +class WanImage2VideoModularPipeline(WanModularPipeline): + """ + A ModularPipeline for Wan2.1 image2video (both I2V and FLF2V). + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "WanImage2VideoAutoBlocks" + + +class Wan22ModularPipeline(WanModularPipeline): + """ + A ModularPipeline for Wan2.2 text2video. + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "Wan22Blocks" + + +class Wan22Image2VideoModularPipeline(Wan22ModularPipeline): + """ + A ModularPipeline for Wan2.2 image2video. + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "Wan22Image2VideoBlocks" diff --git a/src/diffusers/modular_pipelines/z_image/__init__.py b/src/diffusers/modular_pipelines/z_image/__init__.py index c8a8c14396c0..5c04008d3305 100644 --- a/src/diffusers/modular_pipelines/z_image/__init__.py +++ b/src/diffusers/modular_pipelines/z_image/__init__.py @@ -21,12 +21,7 @@ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) else: - _import_structure["decoders"] = ["ZImageVaeDecoderStep"] - _import_structure["encoders"] = ["ZImageTextEncoderStep", "ZImageVaeImageEncoderStep"] - _import_structure["modular_blocks"] = [ - "ALL_BLOCKS", - "ZImageAutoBlocks", - ] + _import_structure["modular_blocks_z_image"] = ["ZImageAutoBlocks"] _import_structure["modular_pipeline"] = ["ZImageModularPipeline"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: @@ -36,12 +31,7 @@ except OptionalDependencyNotAvailable: from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 else: - from .decoders import ZImageVaeDecoderStep - from .encoders import ZImageTextEncoderStep - from .modular_blocks import ( - ALL_BLOCKS, - ZImageAutoBlocks, - ) + from .modular_blocks_z_image import ZImageAutoBlocks from .modular_pipeline import ZImageModularPipeline else: import sys diff --git a/src/diffusers/modular_pipelines/z_image/before_denoise.py b/src/diffusers/modular_pipelines/z_image/before_denoise.py index 35ea768f12c3..8558f2c67f65 100644 --- a/src/diffusers/modular_pipelines/z_image/before_denoise.py +++ b/src/diffusers/modular_pipelines/z_image/before_denoise.py @@ -13,7 +13,6 @@ # limitations under the License. import inspect -from typing import List, Optional, Tuple, Union import torch @@ -91,7 +90,7 @@ def repeat_tensor_to_batch_size( return input_tensor -def calculate_dimension_from_latents(latents: torch.Tensor, vae_scale_factor_spatial: int) -> Tuple[int, int]: +def calculate_dimension_from_latents(latents: torch.Tensor, vae_scale_factor_spatial: int) -> tuple[int, int]: """Calculate image dimensions from latent tensor dimensions. This function converts latent spatial dimensions to image spatial dimensions by multiplying the latent height/width @@ -103,7 +102,7 @@ def calculate_dimension_from_latents(latents: torch.Tensor, vae_scale_factor_spa vae_scale_factor (int): The scale factor used by the VAE to compress image spatial dimension. By default, it is 16 Returns: - Tuple[int, int]: The calculated image dimensions as (height, width) + tuple[int, int]: The calculated image dimensions as (height, width) """ latent_height, latent_width = latents.shape[2:] height = latent_height * vae_scale_factor_spatial // 2 @@ -129,10 +128,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -147,15 +146,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -201,30 +200,30 @@ def description(self) -> str: ) @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("transformer", ZImageTransformer2DModel), ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("num_images_per_prompt", default=1), InputParam( "prompt_embeds", required=True, - type_hint=List[torch.Tensor], + type_hint=list[torch.Tensor], description="Pre-generated text embeddings. Can be generated from text_encoder step.", ), InputParam( "negative_prompt_embeds", - type_hint=List[torch.Tensor], + type_hint=list[torch.Tensor], description="Pre-generated negative text embeddings. Can be generated from text_encoder step.", ), ] @property - def intermediate_outputs(self) -> List[str]: + def intermediate_outputs(self) -> list[str]: return [ OutputParam( "batch_size", @@ -283,8 +282,8 @@ class ZImageAdditionalInputsStep(ModularPipelineBlocks): def __init__( self, - image_latent_inputs: List[str] = ["image_latents"], - additional_batch_inputs: List[str] = [], + image_latent_inputs: list[str] = ["image_latents"], + additional_batch_inputs: list[str] = [], ): """Initialize a configurable step that standardizes the inputs for the denoising step. It:\n" @@ -295,10 +294,10 @@ def __init__( This is a dynamic block that allows you to configure which inputs to process. Args: - image_latent_inputs (List[str], optional): Names of image latent tensors to process. + image_latent_inputs (list[str], optional): Names of image latent tensors to process. In additional to adjust batch size of these inputs, they will be used to determine height/width. Can be a single string or list of strings. Defaults to ["image_latents"]. - additional_batch_inputs (List[str], optional): + additional_batch_inputs (list[str], optional): Names of additional conditional input tensors to expand batch size. These tensors will only have their batch dimensions adjusted to match the final batch size. Can be a single string or list of strings. Defaults to []. @@ -346,7 +345,7 @@ def description(self) -> str: return summary_section + inputs_info + placement_section @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: inputs = [ InputParam(name="num_images_per_prompt", default=1), InputParam(name="batch_size", required=True), @@ -406,11 +405,11 @@ def description(self) -> str: return "Prepare latents step that prepares the latents for the text-to-video generation process" @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("height", type_hint=int), InputParam("width", type_hint=int), - InputParam("latents", type_hint=Optional[torch.Tensor]), + InputParam("latents", type_hint=torch.Tensor | None), InputParam("num_images_per_prompt", type_hint=int, default=1), InputParam("generator"), InputParam( @@ -423,7 +422,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process" @@ -496,7 +495,7 @@ class ZImageSetTimestepsStep(ModularPipelineBlocks): model_name = "z-image" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), ] @@ -506,7 +505,7 @@ def description(self) -> str: return "Step that sets the scheduler's timesteps for inference. Need to run after prepare latents step." @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("latents", required=True), InputParam("num_inference_steps", default=9), @@ -514,7 +513,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "timesteps", type_hint=torch.Tensor, description="The timesteps to use for the denoising process" @@ -554,7 +553,7 @@ class ZImageSetTimestepsWithStrengthStep(ModularPipelineBlocks): model_name = "z-image" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), ] @@ -564,7 +563,7 @@ def description(self) -> str: return "Step that sets the scheduler's timesteps for inference with strength. Need to run after set timesteps step." @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("timesteps", required=True), InputParam("num_inference_steps", required=True), @@ -602,7 +601,7 @@ def description(self) -> str: return "step that prepares the latents with image condition, need to run after set timesteps and prepare latents step." @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("latents", required=True), InputParam("image_latents", required=True), diff --git a/src/diffusers/modular_pipelines/z_image/decoders.py b/src/diffusers/modular_pipelines/z_image/decoders.py index cdb6a2e5eac1..353253102376 100644 --- a/src/diffusers/modular_pipelines/z_image/decoders.py +++ b/src/diffusers/modular_pipelines/z_image/decoders.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, List, Tuple, Union +from typing import Any import numpy as np import PIL @@ -33,7 +33,7 @@ class ZImageVaeDecoderStep(ModularPipelineBlocks): model_name = "z-image" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("vae", AutoencoderKL), ComponentSpec( @@ -49,7 +49,7 @@ def description(self) -> str: return "Step that decodes the denoised latents into images" @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: return [ InputParam( "latents", @@ -64,11 +64,11 @@ def inputs(self) -> List[Tuple[str, Any]]: ] @property - def intermediate_outputs(self) -> List[str]: + def intermediate_outputs(self) -> list[str]: return [ OutputParam( "images", - type_hint=Union[List[PIL.Image.Image], List[torch.Tensor], List[np.ndarray]], + type_hint=list[PIL.Image.Image, list[torch.Tensor], list[np.ndarray]], description="The generated images, can be a PIL.Image.Image, torch.Tensor or a numpy array", ) ] diff --git a/src/diffusers/modular_pipelines/z_image/denoise.py b/src/diffusers/modular_pipelines/z_image/denoise.py index 3d5a00a9df50..863df312389a 100644 --- a/src/diffusers/modular_pipelines/z_image/denoise.py +++ b/src/diffusers/modular_pipelines/z_image/denoise.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Dict, List, Tuple +from typing import Any import torch @@ -46,7 +46,7 @@ def description(self) -> str: ) @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam( "latents", @@ -80,7 +80,7 @@ class ZImageLoopDenoiser(ModularPipelineBlocks): def __init__( self, - guider_input_fields: Dict[str, Any] = {"cap_feats": ("prompt_embeds", "negative_prompt_embeds")}, + guider_input_fields: dict[str, Any] = {"cap_feats": ("prompt_embeds", "negative_prompt_embeds")}, ): """Initialize a denoiser block that calls the denoiser model. This block is used in Z-Image. @@ -101,7 +101,7 @@ def __init__( super().__init__() @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec( "guider", @@ -121,7 +121,7 @@ def description(self) -> str: ) @property - def inputs(self) -> List[Tuple[str, Any]]: + def inputs(self) -> list[tuple[str, Any]]: inputs = [ InputParam( "num_inference_steps", @@ -131,7 +131,7 @@ def inputs(self) -> List[Tuple[str, Any]]: ), InputParam( kwargs_type="denoiser_input_fields", - description="conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.", + description="The conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.", ), ] guider_input_names = [] @@ -205,7 +205,7 @@ class ZImageLoopAfterDenoiser(ModularPipelineBlocks): model_name = "z-image" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), ] @@ -246,13 +246,13 @@ def description(self) -> str: ) @property - def loop_expected_components(self) -> List[ComponentSpec]: + def loop_expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), ] @property - def loop_inputs(self) -> List[InputParam]: + def loop_inputs(self) -> list[InputParam]: return [ InputParam( "timesteps", diff --git a/src/diffusers/modular_pipelines/z_image/encoders.py b/src/diffusers/modular_pipelines/z_image/encoders.py index f5769fe2deec..06deb8236893 100644 --- a/src/diffusers/modular_pipelines/z_image/encoders.py +++ b/src/diffusers/modular_pipelines/z_image/encoders.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Optional, Union import PIL import torch @@ -37,10 +36,10 @@ def get_qwen_prompt_embeds( text_encoder: Qwen3Model, tokenizer: Qwen2Tokenizer, - prompt: Union[str, List[str]], + prompt: str | list[str], device: torch.device, max_sequence_length: int = 512, -) -> List[torch.Tensor]: +) -> list[torch.Tensor]: prompt = [prompt] if isinstance(prompt, str) else prompt for i, prompt_item in enumerate(prompt): @@ -82,7 +81,7 @@ def get_qwen_prompt_embeds( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -134,7 +133,7 @@ def description(self) -> str: return "Text Encoder step that generate text_embeddings to guide the video generation" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("text_encoder", Qwen3Model), ComponentSpec("tokenizer", Qwen2Tokenizer), @@ -147,7 +146,7 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("prompt"), InputParam("negative_prompt"), @@ -155,17 +154,17 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "prompt_embeds", - type_hint=List[torch.Tensor], + type_hint=list[torch.Tensor], kwargs_type="denoiser_input_fields", description="text embeddings used to guide the image generation", ), OutputParam( "negative_prompt_embeds", - type_hint=List[torch.Tensor], + type_hint=list[torch.Tensor], kwargs_type="denoiser_input_fields", description="negative text embeddings used to guide the image generation", ), @@ -182,22 +181,22 @@ def check_inputs(block_state): def encode_prompt( components, prompt: str, - device: Optional[torch.device] = None, + device: torch.device | None = None, prepare_unconditional_embeds: bool = True, - negative_prompt: Optional[str] = None, + negative_prompt: str | None = None, max_sequence_length: int = 512, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device prepare_unconditional_embeds (`bool`): whether to use prepare unconditional embeddings or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -278,7 +277,7 @@ def description(self) -> str: return "Vae Image Encoder step that generate condition_latents based on image to guide the image generation" @property - def expected_components(self) -> List[ComponentSpec]: + def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("vae", AutoencoderKL), ComponentSpec( @@ -290,7 +289,7 @@ def expected_components(self) -> List[ComponentSpec]: ] @property - def inputs(self) -> List[InputParam]: + def inputs(self) -> list[InputParam]: return [ InputParam("image", type_hint=PIL.Image.Image, required=True), InputParam("height"), @@ -299,7 +298,7 @@ def inputs(self) -> List[InputParam]: ] @property - def intermediate_outputs(self) -> List[OutputParam]: + def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "image_latents", diff --git a/src/diffusers/modular_pipelines/z_image/modular_blocks.py b/src/diffusers/modular_pipelines/z_image/modular_blocks.py deleted file mode 100644 index a54baeccaf0c..000000000000 --- a/src/diffusers/modular_pipelines/z_image/modular_blocks.py +++ /dev/null @@ -1,191 +0,0 @@ -# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from ...utils import logging -from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks -from ..modular_pipeline_utils import InsertableDict -from .before_denoise import ( - ZImageAdditionalInputsStep, - ZImagePrepareLatentsStep, - ZImagePrepareLatentswithImageStep, - ZImageSetTimestepsStep, - ZImageSetTimestepsWithStrengthStep, - ZImageTextInputStep, -) -from .decoders import ZImageVaeDecoderStep -from .denoise import ( - ZImageDenoiseStep, -) -from .encoders import ( - ZImageTextEncoderStep, - ZImageVaeImageEncoderStep, -) - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -# z-image -# text2image -class ZImageCoreDenoiseStep(SequentialPipelineBlocks): - block_classes = [ - ZImageTextInputStep, - ZImagePrepareLatentsStep, - ZImageSetTimestepsStep, - ZImageDenoiseStep, - ] - block_names = ["input", "prepare_latents", "set_timesteps", "denoise"] - - @property - def description(self): - return ( - "denoise block that takes encoded conditions and runs the denoising process.\n" - + "This is a sequential pipeline blocks:\n" - + " - `ZImageTextInputStep` is used to adjust the batch size of the model inputs\n" - + " - `ZImagePrepareLatentsStep` is used to prepare the latents\n" - + " - `ZImageSetTimestepsStep` is used to set the timesteps\n" - + " - `ZImageDenoiseStep` is used to denoise the latents\n" - ) - - -# z-image: image2image -## denoise -class ZImageImage2ImageCoreDenoiseStep(SequentialPipelineBlocks): - block_classes = [ - ZImageTextInputStep, - ZImageAdditionalInputsStep(image_latent_inputs=["image_latents"]), - ZImagePrepareLatentsStep, - ZImageSetTimestepsStep, - ZImageSetTimestepsWithStrengthStep, - ZImagePrepareLatentswithImageStep, - ZImageDenoiseStep, - ] - block_names = [ - "input", - "additional_inputs", - "prepare_latents", - "set_timesteps", - "set_timesteps_with_strength", - "prepare_latents_with_image", - "denoise", - ] - - @property - def description(self): - return ( - "denoise block that takes encoded text and image latent conditions and runs the denoising process.\n" - + "This is a sequential pipeline blocks:\n" - + " - `ZImageTextInputStep` is used to adjust the batch size of the model inputs\n" - + " - `ZImageAdditionalInputsStep` is used to adjust the batch size of the latent conditions\n" - + " - `ZImagePrepareLatentsStep` is used to prepare the latents\n" - + " - `ZImageSetTimestepsStep` is used to set the timesteps\n" - + " - `ZImageSetTimestepsWithStrengthStep` is used to set the timesteps with strength\n" - + " - `ZImagePrepareLatentswithImageStep` is used to prepare the latents with image\n" - + " - `ZImageDenoiseStep` is used to denoise the latents\n" - ) - - -## auto blocks -class ZImageAutoDenoiseStep(AutoPipelineBlocks): - block_classes = [ - ZImageImage2ImageCoreDenoiseStep, - ZImageCoreDenoiseStep, - ] - block_names = ["image2image", "text2image"] - block_trigger_inputs = ["image_latents", None] - - @property - def description(self) -> str: - return ( - "Denoise step that iteratively denoise the latents. " - "This is a auto pipeline block that works for text2image and image2image tasks." - " - `ZImageCoreDenoiseStep` (text2image) for text2image tasks." - " - `ZImageImage2ImageCoreDenoiseStep` (image2image) for image2image tasks." - + " - if `image_latents` is provided, `ZImageImage2ImageCoreDenoiseStep` will be used.\n" - + " - if `image_latents` is not provided, `ZImageCoreDenoiseStep` will be used.\n" - ) - - -class ZImageAutoVaeImageEncoderStep(AutoPipelineBlocks): - block_classes = [ZImageVaeImageEncoderStep] - block_names = ["vae_encoder"] - block_trigger_inputs = ["image"] - - @property - def description(self) -> str: - return "Vae Image Encoder step that encode the image to generate the image latents" - +"This is an auto pipeline block that works for image2image tasks." - +" - `ZImageVaeImageEncoderStep` is used when `image` is provided." - +" - if `image` is not provided, step will be skipped." - - -class ZImageAutoBlocks(SequentialPipelineBlocks): - block_classes = [ - ZImageTextEncoderStep, - ZImageAutoVaeImageEncoderStep, - ZImageAutoDenoiseStep, - ZImageVaeDecoderStep, - ] - block_names = ["text_encoder", "vae_encoder", "denoise", "decode"] - - @property - def description(self) -> str: - return "Auto Modular pipeline for text-to-image and image-to-image using ZImage.\n" - +" - for text-to-image generation, all you need to provide is `prompt`\n" - +" - for image-to-image generation, you need to provide `image`\n" - +" - if `image` is not provided, step will be skipped." - - -# presets -TEXT2IMAGE_BLOCKS = InsertableDict( - [ - ("text_encoder", ZImageTextEncoderStep), - ("input", ZImageTextInputStep), - ("prepare_latents", ZImagePrepareLatentsStep), - ("set_timesteps", ZImageSetTimestepsStep), - ("denoise", ZImageDenoiseStep), - ("decode", ZImageVaeDecoderStep), - ] -) - -IMAGE2IMAGE_BLOCKS = InsertableDict( - [ - ("text_encoder", ZImageTextEncoderStep), - ("vae_encoder", ZImageVaeImageEncoderStep), - ("input", ZImageTextInputStep), - ("additional_inputs", ZImageAdditionalInputsStep(image_latent_inputs=["image_latents"])), - ("prepare_latents", ZImagePrepareLatentsStep), - ("set_timesteps", ZImageSetTimestepsStep), - ("set_timesteps_with_strength", ZImageSetTimestepsWithStrengthStep), - ("prepare_latents_with_image", ZImagePrepareLatentswithImageStep), - ("denoise", ZImageDenoiseStep), - ("decode", ZImageVaeDecoderStep), - ] -) - - -AUTO_BLOCKS = InsertableDict( - [ - ("text_encoder", ZImageTextEncoderStep), - ("vae_encoder", ZImageAutoVaeImageEncoderStep), - ("denoise", ZImageAutoDenoiseStep), - ("decode", ZImageVaeDecoderStep), - ] -) - -ALL_BLOCKS = { - "text2image": TEXT2IMAGE_BLOCKS, - "image2image": IMAGE2IMAGE_BLOCKS, - "auto": AUTO_BLOCKS, -} diff --git a/src/diffusers/modular_pipelines/z_image/modular_blocks_z_image.py b/src/diffusers/modular_pipelines/z_image/modular_blocks_z_image.py new file mode 100644 index 000000000000..23e20d55fb1e --- /dev/null +++ b/src/diffusers/modular_pipelines/z_image/modular_blocks_z_image.py @@ -0,0 +1,334 @@ +# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...utils import logging +from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import OutputParam +from .before_denoise import ( + ZImageAdditionalInputsStep, + ZImagePrepareLatentsStep, + ZImagePrepareLatentswithImageStep, + ZImageSetTimestepsStep, + ZImageSetTimestepsWithStrengthStep, + ZImageTextInputStep, +) +from .decoders import ZImageVaeDecoderStep +from .denoise import ( + ZImageDenoiseStep, +) +from .encoders import ( + ZImageTextEncoderStep, + ZImageVaeImageEncoderStep, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# ==================== +# 1. DENOISE +# ==================== + + +# text2image: inputs(text) -> set_timesteps -> prepare_latents -> denoise +# auto_docstring +class ZImageCoreDenoiseStep(SequentialPipelineBlocks): + """ + denoise block that takes encoded conditions and runs the denoising process. + + Components: + transformer (`ZImageTransformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`ClassifierFreeGuidance`) + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`list`): + Pre-generated text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`list`, *optional*): + Pre-generated negative text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 9): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + **denoiser_input_fields (`None`, *optional*): + The conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + block_classes = [ + ZImageTextInputStep, + ZImagePrepareLatentsStep, + ZImageSetTimestepsStep, + ZImageDenoiseStep, + ] + block_names = ["input", "prepare_latents", "set_timesteps", "denoise"] + + @property + def description(self): + return "denoise block that takes encoded conditions and runs the denoising process." + + @property + def outputs(self): + return [OutputParam.template("latents")] + + +# image2image: inputs(text + image_latents) -> prepare_latents -> set_timesteps -> set_timesteps_with_strength -> prepare_latents_with_image -> denoise +# auto_docstring +class ZImageImage2ImageCoreDenoiseStep(SequentialPipelineBlocks): + """ + denoise block that takes encoded text and image latent conditions and runs the denoising process. + + Components: + transformer (`ZImageTransformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`ClassifierFreeGuidance`) + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`list`): + Pre-generated text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`list`, *optional*): + Pre-generated negative text embeddings. Can be generated from text_encoder step. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + image_latents (`None`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 9): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + strength (`None`, *optional*, defaults to 0.6): + TODO: Add description. + **denoiser_input_fields (`None`, *optional*): + The conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + block_classes = [ + ZImageTextInputStep, + ZImageAdditionalInputsStep(image_latent_inputs=["image_latents"]), + ZImagePrepareLatentsStep, + ZImageSetTimestepsStep, + ZImageSetTimestepsWithStrengthStep, + ZImagePrepareLatentswithImageStep, + ZImageDenoiseStep, + ] + block_names = [ + "input", + "additional_inputs", + "prepare_latents", + "set_timesteps", + "set_timesteps_with_strength", + "prepare_latents_with_image", + "denoise", + ] + + @property + def description(self): + return "denoise block that takes encoded text and image latent conditions and runs the denoising process." + + @property + def outputs(self): + return [OutputParam.template("latents")] + + +# auto_docstring +class ZImageAutoDenoiseStep(AutoPipelineBlocks): + """ + Denoise step that iteratively denoise the latents. This is a auto pipeline block that works for text2image and + image2image tasks. - `ZImageCoreDenoiseStep` (text2image) for text2image tasks. - + `ZImageImage2ImageCoreDenoiseStep` (image2image) for image2image tasks. - if `image_latents` is provided, + `ZImageImage2ImageCoreDenoiseStep` will be used. + - if `image_latents` is not provided, `ZImageCoreDenoiseStep` will be used. + + Components: + transformer (`ZImageTransformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`ClassifierFreeGuidance`) + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`list`): + Pre-generated text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`list`, *optional*): + Pre-generated negative text embeddings. Can be generated from text_encoder step. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + image_latents (`None`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`None`): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + strength (`None`, *optional*, defaults to 0.6): + TODO: Add description. + **denoiser_input_fields (`None`, *optional*): + The conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + block_classes = [ + ZImageImage2ImageCoreDenoiseStep, + ZImageCoreDenoiseStep, + ] + block_names = ["image2image", "text2image"] + block_trigger_inputs = ["image_latents", None] + + @property + def description(self) -> str: + return ( + "Denoise step that iteratively denoise the latents. " + "This is a auto pipeline block that works for text2image and image2image tasks." + " - `ZImageCoreDenoiseStep` (text2image) for text2image tasks." + " - `ZImageImage2ImageCoreDenoiseStep` (image2image) for image2image tasks." + + " - if `image_latents` is provided, `ZImageImage2ImageCoreDenoiseStep` will be used.\n" + + " - if `image_latents` is not provided, `ZImageCoreDenoiseStep` will be used.\n" + ) + + +# auto_docstring +class ZImageAutoVaeImageEncoderStep(AutoPipelineBlocks): + """ + Vae Image Encoder step that encode the image to generate the image latents + + Components: + vae (`AutoencoderKL`) image_processor (`VaeImageProcessor`) + + Inputs: + image (`Image`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + image_latents (`Tensor`): + video latent representation with the first frame image condition + """ + + block_classes = [ZImageVaeImageEncoderStep] + block_names = ["vae_encoder"] + block_trigger_inputs = ["image"] + + @property + def description(self) -> str: + return "Vae Image Encoder step that encode the image to generate the image latents" + +"This is an auto pipeline block that works for image2image tasks." + +" - `ZImageVaeImageEncoderStep` is used when `image` is provided." + +" - if `image` is not provided, step will be skipped." + + +# auto_docstring +class ZImageAutoBlocks(SequentialPipelineBlocks): + """ + Auto Modular pipeline for text-to-image and image-to-image using ZImage. + + Supported workflows: + - `text2image`: requires `prompt` + - `image2image`: requires `image`, `prompt` + + Components: + text_encoder (`Qwen3Model`) tokenizer (`Qwen2Tokenizer`) guider (`ClassifierFreeGuidance`) vae + (`AutoencoderKL`) image_processor (`VaeImageProcessor`) transformer (`ZImageTransformer2DModel`) scheduler + (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + prompt (`None`, *optional*): + TODO: Add description. + negative_prompt (`None`, *optional*): + TODO: Add description. + max_sequence_length (`None`, *optional*, defaults to 512): + TODO: Add description. + image (`Image`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + image_latents (`None`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`): + TODO: Add description. + num_inference_steps (`None`): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + strength (`None`, *optional*, defaults to 0.6): + TODO: Add description. + **denoiser_input_fields (`None`, *optional*): + The conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + output_type (`str`, *optional*, defaults to pil): + The type of the output images, can be 'pil', 'np', 'pt' + + Outputs: + images (`list`): + Generated images. + """ + + block_classes = [ + ZImageTextEncoderStep, + ZImageAutoVaeImageEncoderStep, + ZImageAutoDenoiseStep, + ZImageVaeDecoderStep, + ] + block_names = ["text_encoder", "vae_encoder", "denoise", "decode"] + _workflow_map = { + "text2image": {"prompt": True}, + "image2image": {"image": True, "prompt": True}, + } + + @property + def description(self) -> str: + return "Auto Modular pipeline for text-to-image and image-to-image using ZImage." + + @property + def outputs(self): + return [OutputParam.template("images")] diff --git a/src/diffusers/optimization.py b/src/diffusers/optimization.py index e0b3576e4426..044bb0db1908 100644 --- a/src/diffusers/optimization.py +++ b/src/diffusers/optimization.py @@ -16,7 +16,6 @@ import math from enum import Enum -from typing import Optional, Union from torch.optim import Optimizer from torch.optim.lr_scheduler import LambdaLR @@ -287,11 +286,11 @@ def lr_lambda(current_step: int): def get_scheduler( - name: Union[str, SchedulerType], + name: str | SchedulerType, optimizer: Optimizer, - step_rules: Optional[str] = None, - num_warmup_steps: Optional[int] = None, - num_training_steps: Optional[int] = None, + step_rules: str | None = None, + num_warmup_steps: int | None = None, + num_training_steps: int | None = None, num_cycles: int = 1, power: float = 1.0, last_epoch: int = -1, diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index f7615c1a4439..2ee30a58b576 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -6,7 +6,6 @@ _LazyModule, get_objects_from_module, is_flax_available, - is_k_diffusion_available, is_librosa_available, is_note_seq_available, is_onnx_available, @@ -15,6 +14,7 @@ is_torch_available, is_torch_npu_available, is_transformers_available, + is_transformers_version, ) @@ -128,8 +128,8 @@ "AnimateDiffVideoToVideoControlNetPipeline", ] _import_structure["bria"] = ["BriaPipeline"] - _import_structure["bria_fibo"] = ["BriaFiboPipeline"] - _import_structure["flux2"] = ["Flux2Pipeline"] + _import_structure["bria_fibo"] = ["BriaFiboPipeline", "BriaFiboEditPipeline"] + _import_structure["flux2"] = ["Flux2Pipeline", "Flux2KleinPipeline"] _import_structure["flux"] = [ "FluxControlPipeline", "FluxControlInpaintPipeline", @@ -154,7 +154,7 @@ "AudioLDM2UNet2DConditionModel", ] _import_structure["blip_diffusion"] = ["BlipDiffusionPipeline"] - _import_structure["chroma"] = ["ChromaPipeline", "ChromaImg2ImgPipeline"] + _import_structure["chroma"] = ["ChromaPipeline", "ChromaImg2ImgPipeline", "ChromaInpaintPipeline", "ChromaRadiancePipeline"] _import_structure["cogvideo"] = [ "CogVideoXPipeline", "CogVideoXImageToVideoPipeline", @@ -166,6 +166,7 @@ _import_structure["consisid"] = ["ConsisIDPipeline"] _import_structure["cosmos"] = [ "Cosmos2_5_PredictBasePipeline", + "Cosmos2_5_TransferPipeline", "Cosmos2TextToImagePipeline", "CosmosTextToWorldPipeline", "CosmosVideoToWorldPipeline", @@ -236,6 +237,7 @@ "EasyAnimateInpaintPipeline", "EasyAnimateControlPipeline", ] + _import_structure["helios"] = ["HeliosPipeline", "HeliosPyramidPipeline"] _import_structure["hidream_image"] = ["HiDreamImagePipeline"] _import_structure["hunyuandit"] = ["HunyuanDiTPipeline"] _import_structure["hunyuan_video"] = [ @@ -288,6 +290,13 @@ "LTXImageToVideoPipeline", "LTXConditionPipeline", "LTXLatentUpsamplePipeline", + "LTXI2VLongMultiPromptPipeline", + ] + _import_structure["ltx2"] = [ + "LTX2Pipeline", + "LTX2ConditionPipeline", + "LTX2ImageToVideoPipeline", + "LTX2LatentUpsamplePipeline", ] _import_structure["lumina"] = ["LuminaPipeline", "LuminaText2ImgPipeline"] _import_structure["lumina2"] = ["Lumina2Pipeline", "Lumina2Text2ImgPipeline"] @@ -407,11 +416,12 @@ "Kandinsky5I2IPipeline", ] _import_structure["z_image"] = [ - "ZImageImg2ImgPipeline", - "ZImagePipeline", - "ZImageControlNetPipeline", "ZImageControlNetInpaintPipeline", + "ZImageControlNetPipeline", + "ZImageImg2ImgPipeline", + "ZImageInpaintPipeline", "ZImageOmniPipeline", + "ZImagePipeline", ] _import_structure["skyreels_v2"] = [ "SkyReelsV2DiffusionForcingPipeline", @@ -432,6 +442,8 @@ "QwenImageLayeredPipeline", ] _import_structure["chronoedit"] = ["ChronoEditPipeline"] + _import_structure["glm_image"] = ["GlmImagePipeline"] + try: if not is_onnx_available(): raise OptionalDependencyNotAvailable() @@ -459,21 +471,6 @@ ] ) -try: - if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ..utils import ( - dummy_torch_and_transformers_and_k_diffusion_objects, - ) - - _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_and_k_diffusion_objects)) -else: - _import_structure["stable_diffusion_k_diffusion"] = [ - "StableDiffusionKDiffusionPipeline", - "StableDiffusionXLKDiffusionPipeline", - ] - try: if not (is_torch_available() and is_transformers_available() and is_sentencepiece_available()): raise OptionalDependencyNotAvailable() @@ -592,8 +589,8 @@ from .aura_flow import AuraFlowPipeline from .blip_diffusion import BlipDiffusionPipeline from .bria import BriaPipeline - from .bria_fibo import BriaFiboPipeline - from .chroma import ChromaImg2ImgPipeline, ChromaPipeline + from .bria_fibo import BriaFiboEditPipeline, BriaFiboPipeline + from .chroma import ChromaImg2ImgPipeline, ChromaInpaintPipeline, ChromaPipeline from .chronoedit import ChronoEditPipeline from .cogvideo import ( CogVideoXFunControlPipeline, @@ -625,6 +622,7 @@ ) from .cosmos import ( Cosmos2_5_PredictBasePipeline, + Cosmos2_5_TransferPipeline, Cosmos2TextToImagePipeline, Cosmos2VideoToWorldPipeline, CosmosTextToWorldPipeline, @@ -673,7 +671,9 @@ FluxPriorReduxPipeline, ReduxImageEncoder, ) - from .flux2 import Flux2Pipeline + from .flux2 import Flux2KleinPipeline, Flux2Pipeline + from .glm_image import GlmImagePipeline + from .helios import HeliosPipeline, HeliosPyramidPipeline from .hidream_image import HiDreamImagePipeline from .hunyuan_image import HunyuanImagePipeline, HunyuanImageRefinerPipeline from .hunyuan_video import ( @@ -729,7 +729,14 @@ LEditsPPPipelineStableDiffusionXL, ) from .longcat_image import LongCatImageEditPipeline, LongCatImagePipeline - from .ltx import LTXConditionPipeline, LTXImageToVideoPipeline, LTXLatentUpsamplePipeline, LTXPipeline + from .ltx import ( + LTXConditionPipeline, + LTXI2VLongMultiPromptPipeline, + LTXImageToVideoPipeline, + LTXLatentUpsamplePipeline, + LTXPipeline, + ) + from .ltx2 import LTX2ConditionPipeline, LTX2ImageToVideoPipeline, LTX2LatentUpsamplePipeline, LTX2Pipeline from .lucy import LucyEditPipeline from .lumina import LuminaPipeline, LuminaText2ImgPipeline from .lumina2 import Lumina2Pipeline, Lumina2Text2ImgPipeline @@ -857,6 +864,7 @@ ZImageControlNetInpaintPipeline, ZImageControlNetPipeline, ZImageImg2ImgPipeline, + ZImageInpaintPipeline, ZImageOmniPipeline, ZImagePipeline, ) @@ -884,17 +892,6 @@ StableDiffusionOnnxPipeline, ) - try: - if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()): - raise OptionalDependencyNotAvailable() - except OptionalDependencyNotAvailable: - from ..utils.dummy_torch_and_transformers_and_k_diffusion_objects import * - else: - from .stable_diffusion_k_diffusion import ( - StableDiffusionKDiffusionPipeline, - StableDiffusionXLKDiffusionPipeline, - ) - try: if not (is_torch_available() and is_transformers_available() and is_sentencepiece_available()): raise OptionalDependencyNotAvailable() diff --git a/src/diffusers/pipelines/allegro/pipeline_allegro.py b/src/diffusers/pipelines/allegro/pipeline_allegro.py index 3be0129088fb..e54e9ed20739 100644 --- a/src/diffusers/pipelines/allegro/pipeline_allegro.py +++ b/src/diffusers/pipelines/allegro/pipeline_allegro.py @@ -18,7 +18,7 @@ import math import re import urllib.parse as ul -from typing import Callable, Dict, List, Optional, Tuple, Union +from typing import Callable import torch from transformers import T5EncoderModel, T5Tokenizer @@ -84,10 +84,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -102,15 +102,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -214,15 +214,15 @@ def __init__( # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.encode_prompt with 120->512, num_images_per_prompt->num_videos_per_prompt def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], do_classifier_free_guidance: bool = True, negative_prompt: str = "", num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, + device: torch.device | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, clean_caption: bool = False, max_sequence_length: int = 512, **kwargs, @@ -231,9 +231,9 @@ def encode_prompt( Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For PixArt-Alpha, this should be "". @@ -719,46 +719,44 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, negative_prompt: str = "", num_inference_steps: int = 100, - timesteps: List[int] = None, + timesteps: list[int] = None, guidance_scale: float = 7.5, - num_frames: Optional[int] = None, - height: Optional[int] = None, - width: Optional[int] = None, + num_frames: int | None = None, + height: int | None = None, + width: int | None = None, num_videos_per_prompt: int = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], clean_caption: bool = True, max_sequence_length: int = 512, - ) -> Union[AllegroPipelineOutput, Tuple]: + ) -> AllegroPipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the video generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_inference_steps (`int`, *optional*, defaults to 100): The number of denoising steps. More denoising steps usually lead to a higher quality video at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps` timesteps are used. Must be in descending order. guidance_scale (`float`, *optional*, defaults to 7.5): @@ -778,7 +776,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -887,7 +885,13 @@ def __call__( prompt_embeds = prompt_embeds.unsqueeze(1) # b l d -> b 1 l d # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, timesteps + ) self.scheduler.set_timesteps(num_inference_steps, device=device) # 5. Prepare latents. diff --git a/src/diffusers/pipelines/allegro/pipeline_output.py b/src/diffusers/pipelines/allegro/pipeline_output.py index 6a721783ca86..bf85a4954ce9 100644 --- a/src/diffusers/pipelines/allegro/pipeline_output.py +++ b/src/diffusers/pipelines/allegro/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL @@ -14,10 +13,10 @@ class AllegroPipelineOutput(BaseOutput): Output class for Allegro pipelines. Args: - frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): - List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): + list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape `(batch_size, num_frames, channels, height, width)`. """ - frames: Union[torch.Tensor, np.ndarray, List[List[PIL.Image.Image]]] + frames: torch.Tensor | np.ndarray | list[list[PIL.Image.Image]] diff --git a/src/diffusers/pipelines/amused/pipeline_amused.py b/src/diffusers/pipelines/amused/pipeline_amused.py index 131e34d1a4a1..b23adf0d2152 100644 --- a/src/diffusers/pipelines/amused/pipeline_amused.py +++ b/src/diffusers/pipelines/amused/pipeline_amused.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from transformers import CLIPTextModelWithProjection, CLIPTokenizer @@ -84,33 +84,33 @@ def __init__( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Optional[Union[List[str], str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: list[str] | str | None = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 12, guidance_scale: float = 10.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[torch.Generator] = None, - latents: Optional[torch.IntTensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_encoder_hidden_states: Optional[torch.Tensor] = None, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | None = None, + latents: torch.IntTensor | None = None, + prompt_embeds: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_encoder_hidden_states: torch.Tensor | None = None, output_type="pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, micro_conditioning_aesthetic_score: int = 6, - micro_conditioning_crop_coord: Tuple[int, int] = (0, 0), - temperature: Union[int, Tuple[int, int], List[int]] = (2, 0), + micro_conditioning_crop_coord: tuple[int, int] = (0, 0), + temperature: int | tuple[int, int] | list[int] = (2, 0), ): """ The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.transformer.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -122,7 +122,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 10.0): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -162,10 +162,10 @@ def __call__( The targeted aesthetic score according to the laion aesthetic classifier. See https://laion.ai/blog/laion-aesthetics/ and the micro-conditioning section of https://huggingface.co/papers/2307.01952. - micro_conditioning_crop_coord (`Tuple[int]`, *optional*, defaults to (0, 0)): + micro_conditioning_crop_coord (`tuple[int]`, *optional*, defaults to (0, 0)): The targeted height, width crop coordinates. See the micro-conditioning section of https://huggingface.co/papers/2307.01952. - temperature (`Union[int, Tuple[int, int], List[int]]`, *optional*, defaults to (2, 0)): + temperature (`int | tuple[int, int, list[int]]`, *optional*, defaults to (2, 0)): Configures the temperature scheduler on `self.scheduler` see `AmusedScheduler#set_timesteps`. Examples: diff --git a/src/diffusers/pipelines/amused/pipeline_amused_img2img.py b/src/diffusers/pipelines/amused/pipeline_amused_img2img.py index a122c12236dd..79ebd96dedeb 100644 --- a/src/diffusers/pipelines/amused/pipeline_amused_img2img.py +++ b/src/diffusers/pipelines/amused/pipeline_amused_img2img.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from transformers import CLIPTextModelWithProjection, CLIPTokenizer @@ -99,34 +99,34 @@ def __init__( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Optional[Union[List[str], str]] = None, + prompt: list[str] | str | None = None, image: PipelineImageInput = None, strength: float = 0.5, num_inference_steps: int = 12, guidance_scale: float = 10.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[torch.Generator] = None, - prompt_embeds: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_encoder_hidden_states: Optional[torch.Tensor] = None, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | None = None, + prompt_embeds: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_encoder_hidden_states: torch.Tensor | None = None, output_type="pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, micro_conditioning_aesthetic_score: int = 6, - micro_conditioning_crop_coord: Tuple[int, int] = (0, 0), - temperature: Union[int, Tuple[int, int], List[int]] = (2, 0), + micro_conditioning_crop_coord: tuple[int, int] = (0, 0), + temperature: int | tuple[int, int] | list[int] = (2, 0), ): """ The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a @@ -144,7 +144,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 10.0): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -181,10 +181,10 @@ def __call__( The targeted aesthetic score according to the laion aesthetic classifier. See https://laion.ai/blog/laion-aesthetics/ and the micro-conditioning section of https://huggingface.co/papers/2307.01952. - micro_conditioning_crop_coord (`Tuple[int]`, *optional*, defaults to (0, 0)): + micro_conditioning_crop_coord (`tuple[int]`, *optional*, defaults to (0, 0)): The targeted height, width crop coordinates. See the micro-conditioning section of https://huggingface.co/papers/2307.01952. - temperature (`Union[int, Tuple[int, int], List[int]]`, *optional*, defaults to (2, 0)): + temperature (`int | tuple[int, int, list[int]]`, *optional*, defaults to (2, 0)): Configures the temperature scheduler on `self.scheduler` see `AmusedScheduler#set_timesteps`. Examples: diff --git a/src/diffusers/pipelines/amused/pipeline_amused_inpaint.py b/src/diffusers/pipelines/amused/pipeline_amused_inpaint.py index f4bd4944ff9a..55302401832c 100644 --- a/src/diffusers/pipelines/amused/pipeline_amused_inpaint.py +++ b/src/diffusers/pipelines/amused/pipeline_amused_inpaint.py @@ -13,7 +13,7 @@ # limitations under the License. -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from transformers import CLIPTextModelWithProjection, CLIPTokenizer @@ -115,41 +115,41 @@ def __init__( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Optional[Union[List[str], str]] = None, + prompt: list[str] | str | None = None, image: PipelineImageInput = None, mask_image: PipelineImageInput = None, strength: float = 1.0, num_inference_steps: int = 12, guidance_scale: float = 10.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[torch.Generator] = None, - prompt_embeds: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_encoder_hidden_states: Optional[torch.Tensor] = None, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | None = None, + prompt_embeds: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_encoder_hidden_states: torch.Tensor | None = None, output_type="pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, micro_conditioning_aesthetic_score: int = 6, - micro_conditioning_crop_coord: Tuple[int, int] = (0, 0), - temperature: Union[int, Tuple[int, int], List[int]] = (2, 0), + micro_conditioning_crop_coord: tuple[int, int] = (0, 0), + temperature: int | tuple[int, int] | list[int] = (2, 0), ): """ The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image latents as `image`, but if passing latents directly it is not encoded again. - mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one @@ -168,7 +168,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 10.0): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -205,10 +205,10 @@ def __call__( The targeted aesthetic score according to the laion aesthetic classifier. See https://laion.ai/blog/laion-aesthetics/ and the micro-conditioning section of https://huggingface.co/papers/2307.01952. - micro_conditioning_crop_coord (`Tuple[int]`, *optional*, defaults to (0, 0)): + micro_conditioning_crop_coord (`tuple[int]`, *optional*, defaults to (0, 0)): The targeted height, width crop coordinates. See the micro-conditioning section of https://huggingface.co/papers/2307.01952. - temperature (`Union[int, Tuple[int, int], List[int]]`, *optional*, defaults to (2, 0)): + temperature (`int | tuple[int, int, list[int]]`, *optional*, defaults to (2, 0)): Configures the temperature scheduler on `self.scheduler` see `AmusedScheduler#set_timesteps`. Examples: diff --git a/src/diffusers/pipelines/animatediff/pipeline_animatediff.py b/src/diffusers/pipelines/animatediff/pipeline_animatediff.py index 091b6db713ba..4d7477bc8754 100644 --- a/src/diffusers/pipelines/animatediff/pipeline_animatediff.py +++ b/src/diffusers/pipelines/animatediff/pipeline_animatediff.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection @@ -122,16 +122,14 @@ def __init__( vae: AutoencoderKL, text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, - unet: Union[UNet2DConditionModel, UNetMotionModel], + unet: UNet2DConditionModel | UNetMotionModel, motion_adapter: MotionAdapter, - scheduler: Union[ - DDIMScheduler, - PNDMScheduler, - LMSDiscreteScheduler, - EulerDiscreteScheduler, - EulerAncestralDiscreteScheduler, - DPMSolverMultistepScheduler, - ], + scheduler: DDIMScheduler + | PNDMScheduler + | LMSDiscreteScheduler + | EulerDiscreteScheduler + | EulerAncestralDiscreteScheduler + | DPMSolverMultistepScheduler, feature_extractor: CLIPImageProcessor = None, image_encoder: CLIPVisionModelWithProjection = None, ): @@ -160,16 +158,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -177,7 +175,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -276,7 +274,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -574,27 +572,27 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Optional[Union[str, List[str]]] = None, - num_frames: Optional[int] = 16, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] | None = None, + num_frames: int | None = 16, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_videos_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_videos_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + cross_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], decode_chunk_size: int = 16, **kwargs, ): @@ -602,7 +600,7 @@ def __call__( The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated video. @@ -617,13 +615,13 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -639,7 +637,7 @@ def __call__( not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -660,7 +658,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/animatediff/pipeline_animatediff_controlnet.py b/src/diffusers/pipelines/animatediff/pipeline_animatediff_controlnet.py index 70180ccf0650..eb511129cc6f 100644 --- a/src/diffusers/pipelines/animatediff/pipeline_animatediff_controlnet.py +++ b/src/diffusers/pipelines/animatediff/pipeline_animatediff_controlnet.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch import torch.nn.functional as F @@ -164,12 +164,12 @@ def __init__( vae: AutoencoderKL, text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, - unet: Union[UNet2DConditionModel, UNetMotionModel], + unet: UNet2DConditionModel | UNetMotionModel, motion_adapter: MotionAdapter, - controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel], + controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel, scheduler: KarrasDiffusionSchedulers, - feature_extractor: Optional[CLIPImageProcessor] = None, - image_encoder: Optional[CLIPVisionModelWithProjection] = None, + feature_extractor: CLIPImageProcessor | None = None, + image_encoder: CLIPVisionModelWithProjection | None = None, ): super().__init__() if isinstance(unet, UNet2DConditionModel): @@ -203,16 +203,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -220,7 +220,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -319,7 +319,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -721,39 +721,39 @@ def interrupt(self): @torch.no_grad() def __call__( self, - prompt: Union[str, List[str]] = None, - num_frames: Optional[int] = 16, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + num_frames: int | None = 16, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_videos_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_videos_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[PipelineImageInput] = None, - conditioning_frames: Optional[List[PipelineImageInput]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: PipelineImageInput | None = None, + conditioning_frames: list[PipelineImageInput] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, + cross_attention_kwargs: dict[str, Any] | None = None, + controlnet_conditioning_scale: float | list[float] = 1.0, guess_mode: bool = False, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], decode_chunk_size: int = 16, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated video. @@ -768,13 +768,13 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -790,12 +790,12 @@ def __call__( not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not provided, embeddings are computed from the `ip_adapter_image` input argument. - conditioning_frames (`List[PipelineImageInput]`, *optional*): + conditioning_frames (`list[PipelineImageInput]`, *optional*): The ControlNet input condition to provide guidance to the `unet` for generation. If multiple ControlNets are specified, images must be passed as a list such that each element of the list can be correctly batched for input to a single ControlNet. @@ -807,16 +807,16 @@ def __call__( cross_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set the corresponding scale as a list. guess_mode (`bool`, *optional*, defaults to `False`): The ControlNet encoder tries to recognize the content of the input image even if you remove all prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended. - control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the ControlNet starts applying. - control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the ControlNet stops applying. clip_skip (`int`, *optional*): Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that @@ -826,7 +826,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py b/src/diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py index 56d319027595..68ce7c92896a 100644 --- a/src/diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py +++ b/src/diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from transformers import ( @@ -150,10 +150,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -168,15 +168,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -284,16 +284,14 @@ def __init__( text_encoder_2: CLIPTextModelWithProjection, tokenizer: CLIPTokenizer, tokenizer_2: CLIPTokenizer, - unet: Union[UNet2DConditionModel, UNetMotionModel], + unet: UNet2DConditionModel | UNetMotionModel, motion_adapter: MotionAdapter, - scheduler: Union[ - DDIMScheduler, - PNDMScheduler, - LMSDiscreteScheduler, - EulerDiscreteScheduler, - EulerAncestralDiscreteScheduler, - DPMSolverMultistepScheduler, - ], + scheduler: DDIMScheduler + | PNDMScheduler + | LMSDiscreteScheduler + | EulerDiscreteScheduler + | EulerAncestralDiscreteScheduler + | DPMSolverMultistepScheduler, image_encoder: CLIPVisionModelWithProjection = None, feature_extractor: CLIPImageProcessor = None, force_zeros_for_empty_prompt: bool = True, @@ -329,26 +327,26 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, + prompt_2: str | None = None, + device: torch.device | None = None, num_videos_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders device: (`torch.device`): @@ -357,11 +355,11 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -479,7 +477,7 @@ def encode_prompt( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) - uncond_tokens: List[str] + uncond_tokens: list[str] if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -870,50 +868,50 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, num_frames: int = 16, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_end: Optional[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, + denoising_end: float | None = None, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_videos_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + num_videos_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, guidance_rescale: float = 0.0, - original_size: Optional[Tuple[int, int]] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Optional[Tuple[int, int]] = None, - negative_original_size: Optional[Tuple[int, int]] = None, - negative_crops_coords_top_left: Tuple[int, int] = (0, 0), - negative_target_size: Optional[Tuple[int, int]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + original_size: tuple[int, int] | None = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] | None = None, + negative_original_size: tuple[int, int] | None = None, + negative_crops_coords_top_left: tuple[int, int] = (0, 0), + negative_target_size: tuple[int, int] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders num_frames: @@ -932,11 +930,11 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality video at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -953,11 +951,11 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower video quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the video generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the video generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders num_videos_per_prompt (`int`, *optional*, defaults to 1): @@ -965,7 +963,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -988,7 +986,7 @@ def __call__( input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. If not provided, embeddings are computed from the `ip_adapter_image` input argument. output_type (`str`, *optional*, defaults to `"pil"`): @@ -1007,31 +1005,31 @@ def __call__( [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when using zero terminal SNR. - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a specific image resolution. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a target image resolution. It should be as same as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more @@ -1041,7 +1039,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/animatediff/pipeline_animatediff_sparsectrl.py b/src/diffusers/pipelines/animatediff/pipeline_animatediff_sparsectrl.py index 46d650efe8b6..8b3eb8fc3c03 100644 --- a/src/diffusers/pipelines/animatediff/pipeline_animatediff_sparsectrl.py +++ b/src/diffusers/pipelines/animatediff/pipeline_animatediff_sparsectrl.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL @@ -117,7 +117,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -176,7 +176,7 @@ def __init__( vae: AutoencoderKL, text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, - unet: Union[UNet2DConditionModel, UNetMotionModel], + unet: UNet2DConditionModel | UNetMotionModel, motion_adapter: MotionAdapter, controlnet: SparseControlNetModel, scheduler: KarrasDiffusionSchedulers, @@ -212,16 +212,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -229,7 +229,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -328,7 +328,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -675,7 +675,7 @@ def prepare_sparse_control_conditioning( controlnet_frame_indices: int, device: torch.device, dtype: torch.dtype, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: assert conditioning_frames.shape[2] >= len(controlnet_frame_indices) batch_size, channels, _, height, width = conditioning_frames.shape @@ -713,37 +713,37 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, num_frames: int = 16, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, + negative_prompt: str | list[str] | None = None, num_videos_per_prompt: int = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - conditioning_frames: Optional[List[PipelineImageInput]] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + conditioning_frames: list[PipelineImageInput] | None = None, output_type: str = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, - controlnet_frame_indices: List[int] = [0], + cross_attention_kwargs: dict[str, Any] | None = None, + controlnet_conditioning_scale: float | list[float] = 1.0, + controlnet_frame_indices: list[int] = [0], guess_mode: bool = False, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated video. @@ -758,13 +758,13 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -780,12 +780,12 @@ def __call__( not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not provided, embeddings are computed from the `ip_adapter_image` input argument. - conditioning_frames (`List[PipelineImageInput]`, *optional*): + conditioning_frames (`list[PipelineImageInput]`, *optional*): The SparseControlNet input to provide guidance to the `unet` for generation. output_type (`str`, *optional*, defaults to `"pil"`): The output format of the generated video. Choose between `torch.Tensor`, `PIL.Image` or `np.array`. @@ -795,11 +795,11 @@ def __call__( cross_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set the corresponding scale as a list. - controlnet_frame_indices (`List[int]`): + controlnet_frame_indices (`list[int]`): The indices where the conditioning frames must be applied for generation. Multiple frames can be provided to guide the model to generate similar structure outputs, where the `unet` can "fill-in-the-gaps" for interpolation videos, or a single frame could be provided for general expected @@ -812,7 +812,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py b/src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py index 6f3a609aba4a..4e7cd21fc25d 100644 --- a/src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py +++ b/src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection @@ -106,7 +106,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -121,10 +121,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -139,15 +139,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -225,16 +225,14 @@ def __init__( vae: AutoencoderKL, text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, - unet: Union[UNet2DConditionModel, UNetMotionModel], + unet: UNet2DConditionModel | UNetMotionModel, motion_adapter: MotionAdapter, - scheduler: Union[ - DDIMScheduler, - PNDMScheduler, - LMSDiscreteScheduler, - EulerDiscreteScheduler, - EulerAncestralDiscreteScheduler, - DPMSolverMultistepScheduler, - ], + scheduler: DDIMScheduler + | PNDMScheduler + | LMSDiscreteScheduler + | EulerDiscreteScheduler + | EulerAncestralDiscreteScheduler + | DPMSolverMultistepScheduler, feature_extractor: CLIPImageProcessor = None, image_encoder: CLIPVisionModelWithProjection = None, ): @@ -262,16 +260,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -279,7 +277,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -378,7 +376,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -636,16 +634,16 @@ def get_timesteps(self, num_inference_steps, timesteps, strength, device): def prepare_latents( self, - video: Optional[torch.Tensor] = None, + video: torch.Tensor | None = None, height: int = 64, width: int = 64, num_channels_latents: int = 4, batch_size: int = 1, - timestep: Optional[int] = None, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + timestep: int | None = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, decode_chunk_size: int = 16, add_noise: bool = False, ) -> torch.Tensor: @@ -746,40 +744,40 @@ def interrupt(self): @torch.no_grad() def __call__( self, - video: List[List[PipelineImageInput]] = None, - prompt: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + video: list[list[PipelineImageInput]] = None, + prompt: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, enforce_inference_steps: bool = False, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, guidance_scale: float = 7.5, strength: float = 0.8, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_videos_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_videos_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + cross_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], decode_chunk_size: int = 16, ): r""" The call function to the pipeline for generation. Args: - video (`List[PipelineImageInput]`): + video (`list[PipelineImageInput]`): The input video to condition the generation on. Must be a list of images/frames of the video. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated video. @@ -788,11 +786,11 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality videos at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -801,13 +799,13 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -823,7 +821,7 @@ def __call__( not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -843,7 +841,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -897,16 +895,20 @@ def __call__( dtype = self.dtype # 3. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device if not enforce_inference_steps: timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, timesteps, strength, device) latent_timestep = timesteps[:1].repeat(batch_size * num_videos_per_prompt) else: denoising_inference_steps = int(num_inference_steps / strength) timesteps, denoising_inference_steps = retrieve_timesteps( - self.scheduler, denoising_inference_steps, device, timesteps, sigmas + self.scheduler, denoising_inference_steps, timestep_device, timesteps, sigmas ) timesteps = timesteps[-num_inference_steps:] latent_timestep = timesteps[:1].repeat(batch_size * num_videos_per_prompt) diff --git a/src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video_controlnet.py b/src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video_controlnet.py index b00f344598ad..56ed5e23c1db 100644 --- a/src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video_controlnet.py +++ b/src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video_controlnet.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch import torch.nn.functional as F @@ -124,7 +124,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -139,10 +139,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -157,15 +157,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -229,7 +229,7 @@ class AnimateDiffVideoToVideoControlNetPipeline( A [`UNet2DConditionModel`] used to create a UNetMotionModel to denoise the encoded video latents. motion_adapter ([`MotionAdapter`]): A [`MotionAdapter`] to be used in combination with `unet` to denoise the encoded video latents. - controlnet ([`ControlNetModel`] or `List[ControlNetModel]` or `Tuple[ControlNetModel]` or `MultiControlNetModel`): + controlnet ([`ControlNetModel`] or `list[ControlNetModel]` or `tuple[ControlNetModel]` or `MultiControlNetModel`): Provides additional conditioning to the `unet` during the denoising process. If you set multiple ControlNets as a list, the outputs from each ControlNet are added together to create one combined additional conditioning. @@ -247,17 +247,15 @@ def __init__( vae: AutoencoderKL, text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, - unet: Union[UNet2DConditionModel, UNetMotionModel], + unet: UNet2DConditionModel | UNetMotionModel, motion_adapter: MotionAdapter, - controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel], - scheduler: Union[ - DDIMScheduler, - PNDMScheduler, - LMSDiscreteScheduler, - EulerDiscreteScheduler, - EulerAncestralDiscreteScheduler, - DPMSolverMultistepScheduler, - ], + controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel, + scheduler: DDIMScheduler + | PNDMScheduler + | LMSDiscreteScheduler + | EulerDiscreteScheduler + | EulerAncestralDiscreteScheduler + | DPMSolverMultistepScheduler, feature_extractor: CLIPImageProcessor = None, image_encoder: CLIPVisionModelWithProjection = None, ): @@ -293,16 +291,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -310,7 +308,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -409,7 +407,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -768,16 +766,16 @@ def get_timesteps(self, num_inference_steps, timesteps, strength, device): # Copied from diffusers.pipelines.animatediff.pipeline_animatediff_video2video.AnimateDiffVideoToVideoPipeline.prepare_latents def prepare_latents( self, - video: Optional[torch.Tensor] = None, + video: torch.Tensor | None = None, height: int = 64, width: int = 64, num_channels_latents: int = 4, batch_size: int = 1, - timestep: Optional[int] = None, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + timestep: int | None = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, decode_chunk_size: int = 16, add_noise: bool = False, ) -> torch.Tensor: @@ -911,45 +909,45 @@ def interrupt(self): @torch.no_grad() def __call__( self, - video: List[List[PipelineImageInput]] = None, - prompt: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + video: list[list[PipelineImageInput]] = None, + prompt: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, enforce_inference_steps: bool = False, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, guidance_scale: float = 7.5, strength: float = 0.8, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_videos_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_videos_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - conditioning_frames: Optional[List[PipelineImageInput]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + conditioning_frames: list[PipelineImageInput] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, + cross_attention_kwargs: dict[str, Any] | None = None, + controlnet_conditioning_scale: float | list[float] = 1.0, guess_mode: bool = False, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], decode_chunk_size: int = 16, ): r""" The call function to the pipeline for generation. Args: - video (`List[PipelineImageInput]`): + video (`list[PipelineImageInput]`): The input video to condition the generation on. Must be a list of images/frames of the video. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated video. @@ -958,11 +956,11 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality videos at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -971,13 +969,13 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -993,12 +991,12 @@ def __call__( not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not provided, embeddings are computed from the `ip_adapter_image` input argument. - conditioning_frames (`List[PipelineImageInput]`, *optional*): + conditioning_frames (`list[PipelineImageInput]`, *optional*): The ControlNet input condition to provide guidance to the `unet` for generation. If multiple ControlNets are specified, images must be passed as a list such that each element of the list can be correctly batched for input to a single ControlNet. @@ -1009,16 +1007,16 @@ def __call__( cross_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set the corresponding scale as a list. guess_mode (`bool`, *optional*, defaults to `False`): The ControlNet encoder tries to recognize the content of the input image even if you remove all prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended. - control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the ControlNet starts applying. - control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the ControlNet stops applying. clip_skip (`int`, *optional*): Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that @@ -1028,7 +1026,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1100,16 +1098,20 @@ def __call__( dtype = self.dtype # 3. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device if not enforce_inference_steps: timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, timesteps, strength, device) latent_timestep = timesteps[:1].repeat(batch_size * num_videos_per_prompt) else: denoising_inference_steps = int(num_inference_steps / strength) timesteps, denoising_inference_steps = retrieve_timesteps( - self.scheduler, denoising_inference_steps, device, timesteps, sigmas + self.scheduler, denoising_inference_steps, timestep_device, timesteps, sigmas ) timesteps = timesteps[-num_inference_steps:] latent_timestep = timesteps[:1].repeat(batch_size * num_videos_per_prompt) diff --git a/src/diffusers/pipelines/animatediff/pipeline_output.py b/src/diffusers/pipelines/animatediff/pipeline_output.py index 2417223cf95e..436a20f455bf 100644 --- a/src/diffusers/pipelines/animatediff/pipeline_output.py +++ b/src/diffusers/pipelines/animatediff/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -14,11 +13,11 @@ class AnimateDiffPipelineOutput(BaseOutput): Output class for AnimateDiff pipelines. Args: - frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): - List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): + list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape `(batch_size, num_frames, channels, height, width)` """ - frames: Union[torch.Tensor, np.ndarray, List[List[PIL.Image.Image]]] + frames: torch.Tensor | np.ndarray | list[list[PIL.Image.Image]] diff --git a/src/diffusers/pipelines/audioldm/pipeline_audioldm.py b/src/diffusers/pipelines/audioldm/pipeline_audioldm.py index 6a70f00c76c7..357c3582b21c 100644 --- a/src/diffusers/pipelines/audioldm/pipeline_audioldm.py +++ b/src/diffusers/pipelines/audioldm/pipeline_audioldm.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -88,7 +88,7 @@ def __init__( self, vae: AutoencoderKL, text_encoder: ClapTextModelWithProjection, - tokenizer: Union[RobertaTokenizer, RobertaTokenizerFast], + tokenizer: RobertaTokenizer | RobertaTokenizerFast, unet: UNet2DConditionModel, scheduler: KarrasDiffusionSchedulers, vocoder: SpeechT5HifiGan, @@ -112,14 +112,14 @@ def _encode_prompt( num_waveforms_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device (`torch.device`): torch device @@ -127,7 +127,7 @@ def _encode_prompt( number of waveforms that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the audio generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -189,7 +189,7 @@ def _encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -361,28 +361,28 @@ def prepare_latents(self, batch_size, num_channels_latents, height, dtype, devic @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - audio_length_in_s: Optional[float] = None, + prompt: str | list[str] = None, + audio_length_in_s: float | None = None, num_inference_steps: int = 10, guidance_scale: float = 2.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_waveforms_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_waveforms_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: Optional[int] = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - output_type: Optional[str] = "np", + callback: Callable[[int, int, torch.Tensor], None] | None = None, + callback_steps: int | None = 1, + cross_attention_kwargs: dict[str, Any] | None = None, + output_type: str | None = "np", ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide audio generation. If not defined, you need to pass `prompt_embeds`. audio_length_in_s (`int`, *optional*, defaults to 5.12): The length of the generated audio sample in seconds. @@ -392,7 +392,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 2.5): A higher guidance scale value encourages the model to generate audio that is closely linked to the text `prompt` at the expense of lower sound quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in audio generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_waveforms_per_prompt (`int`, *optional*, defaults to 1): @@ -400,7 +400,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py b/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py index 878f6f08db42..09aa0ad17003 100644 --- a/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py +++ b/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py @@ -13,7 +13,7 @@ # limitations under the License. from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import torch import torch.nn as nn @@ -28,10 +28,7 @@ AttnAddedKVProcessor, AttnProcessor, ) -from ...models.embeddings import ( - TimestepEmbedding, - Timesteps, -) +from ...models.embeddings import TimestepEmbedding, Timesteps from ...models.modeling_utils import ModelMixin from ...models.resnet import Downsample2D, ResnetBlock2D, Upsample2D from ...models.transformers.transformer_2d import Transformer2DModel @@ -75,7 +72,7 @@ class AudioLDM2ProjectionModelOutput(BaseOutput): """ hidden_states: torch.Tensor - attention_mask: Optional[torch.LongTensor] = None + attention_mask: torch.LongTensor | None = None class AudioLDM2ProjectionModel(ModelMixin, ConfigMixin): @@ -124,10 +121,10 @@ def __init__( def forward( self, - hidden_states: Optional[torch.Tensor] = None, - hidden_states_1: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.LongTensor] = None, - attention_mask_1: Optional[torch.LongTensor] = None, + hidden_states: torch.Tensor | None = None, + hidden_states_1: torch.Tensor | None = None, + attention_mask: torch.LongTensor | None = None, + attention_mask_1: torch.LongTensor | None = None, ): hidden_states = self.projection(hidden_states) hidden_states, attention_mask = add_special_tokens( @@ -174,23 +171,23 @@ class AudioLDM2UNet2DConditionModel(ModelMixin, AttentionMixin, ConfigMixin, UNe for all models (such as downloading or saving). Parameters: - sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`): + sample_size (`int` or `tuple[int, int]`, *optional*, defaults to `None`): Height and width of input/output sample. in_channels (`int`, *optional*, defaults to 4): Number of channels in the input sample. out_channels (`int`, *optional*, defaults to 4): Number of channels in the output. flip_sin_to_cos (`bool`, *optional*, defaults to `False`): Whether to flip the sin to cos in the time embedding. freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding. - down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`): + down_block_types (`tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`): The tuple of downsample blocks to use. mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`): Block type for middle of UNet, it can only be `UNetMidBlock2DCrossAttn` for AudioLDM2. - up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`): + up_block_types (`tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`): The tuple of upsample blocks to use. - only_cross_attention (`bool` or `Tuple[bool]`, *optional*, default to `False`): + only_cross_attention (`bool` or `tuple[bool]`, *optional*, default to `False`): Whether to include self-attention in the basic transformer blocks, see [`~models.attention.BasicTransformerBlock`]. - block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): + block_out_channels (`tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): The tuple of output channels for each block. layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block. downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution. @@ -199,9 +196,9 @@ class AudioLDM2UNet2DConditionModel(ModelMixin, AttentionMixin, ConfigMixin, UNe norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization. If `None`, normalization and activation layers is skipped in post-processing. norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization. - cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280): + cross_attention_dim (`int` or `tuple[int]`, *optional*, defaults to 1280): The dimension of the cross attention features. - transformer_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 1): + transformer_layers_per_block (`int` or `tuple[int]`, *optional*, defaults to 1): The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`], [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`]. @@ -240,49 +237,44 @@ class conditioning with `class_embed_type` equal to `None`. @register_to_config def __init__( self, - sample_size: Optional[int] = None, + sample_size: int | None = None, in_channels: int = 4, out_channels: int = 4, flip_sin_to_cos: bool = True, freq_shift: int = 0, - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str] = ( "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D", ), - mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn", - up_block_types: Tuple[str, ...] = ( - "UpBlock2D", - "CrossAttnUpBlock2D", - "CrossAttnUpBlock2D", - "CrossAttnUpBlock2D", - ), - only_cross_attention: Union[bool, Tuple[bool]] = False, - block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280), - layers_per_block: Union[int, Tuple[int]] = 2, + mid_block_type: str = "UNetMidBlock2DCrossAttn", + up_block_types: tuple[str] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"), + only_cross_attention: bool | tuple[bool] = False, + block_out_channels: tuple[int] = (320, 640, 1280, 1280), + layers_per_block: int | tuple[int] = 2, downsample_padding: int = 1, mid_block_scale_factor: float = 1, act_fn: str = "silu", - norm_num_groups: Optional[int] = 32, + norm_num_groups: int | None = 32, norm_eps: float = 1e-5, - cross_attention_dim: Union[int, Tuple[int]] = 1280, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, - attention_head_dim: Union[int, Tuple[int]] = 8, - num_attention_heads: Optional[Union[int, Tuple[int]]] = None, + cross_attention_dim: int | tuple[int] = 1280, + transformer_layers_per_block: int | tuple[int] = 1, + attention_head_dim: int | tuple[int] = 8, + num_attention_heads: int | tuple[int] | None = None, use_linear_projection: bool = False, - class_embed_type: Optional[str] = None, - num_class_embeds: Optional[int] = None, + class_embed_type: str | None = None, + num_class_embeds: int | None = None, upcast_attention: bool = False, resnet_time_scale_shift: str = "default", time_embedding_type: str = "positional", - time_embedding_dim: Optional[int] = None, - time_embedding_act_fn: Optional[str] = None, - timestep_post_act: Optional[str] = None, - time_cond_proj_dim: Optional[int] = None, + time_embedding_dim: int | None = None, + time_embedding_act_fn: str | None = None, + timestep_post_act: str | None = None, + time_cond_proj_dim: int | None = None, conv_in_kernel: int = 3, conv_out_kernel: int = 3, - projection_class_embeddings_input_dim: Optional[int] = None, + projection_class_embeddings_input_dim: int | None = None, class_embeddings_concat: bool = False, ): super().__init__() @@ -606,7 +598,7 @@ def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): # Recursively walk through all the children. # Any children which exposes the set_attention_slice method # gets the message - def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]): + def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: list[int]): if hasattr(module, "set_attention_slice"): module.set_attention_slice(slice_size.pop()) @@ -620,17 +612,17 @@ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[i def forward( self, sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], + timestep: torch.Tensor | float | int, encoder_hidden_states: torch.Tensor, - class_labels: Optional[torch.Tensor] = None, - timestep_cond: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, + class_labels: torch.Tensor | None = None, + timestep_cond: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + encoder_attention_mask: torch.Tensor | None = None, return_dict: bool = True, - encoder_hidden_states_1: Optional[torch.Tensor] = None, - encoder_attention_mask_1: Optional[torch.Tensor] = None, - ) -> Union[UNet2DConditionOutput, Tuple]: + encoder_hidden_states_1: torch.Tensor | None = None, + encoder_attention_mask_1: torch.Tensor | None = None, + ) -> UNet2DConditionOutput | tuple: r""" The [`AudioLDM2UNet2DConditionModel`] forward method. @@ -1032,13 +1024,13 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - encoder_hidden_states_1: Optional[torch.Tensor] = None, - encoder_attention_mask_1: Optional[torch.Tensor] = None, + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + encoder_attention_mask: torch.Tensor | None = None, + encoder_hidden_states_1: torch.Tensor | None = None, + encoder_attention_mask_1: torch.Tensor | None = None, ): output_states = () num_layers = len(self.resnets) @@ -1194,13 +1186,13 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - encoder_hidden_states_1: Optional[torch.Tensor] = None, - encoder_attention_mask_1: Optional[torch.Tensor] = None, + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + encoder_attention_mask: torch.Tensor | None = None, + encoder_hidden_states_1: torch.Tensor | None = None, + encoder_attention_mask_1: torch.Tensor | None = None, ) -> torch.Tensor: hidden_states = self.resnets[0](hidden_states, temb) num_attention_per_layer = len(self.attentions) // (len(self.resnets) - 1) @@ -1344,15 +1336,15 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - upsample_size: Optional[int] = None, - attention_mask: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - encoder_hidden_states_1: Optional[torch.Tensor] = None, - encoder_attention_mask_1: Optional[torch.Tensor] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + upsample_size: int | None = None, + attention_mask: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + encoder_hidden_states_1: torch.Tensor | None = None, + encoder_attention_mask_1: torch.Tensor | None = None, ): num_layers = len(self.resnets) num_attention_per_layer = len(self.attentions) // num_layers diff --git a/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py b/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py index 452fc3c01b27..b023974a33dd 100644 --- a/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py +++ b/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -196,11 +196,11 @@ def __init__( self, vae: AutoencoderKL, text_encoder: ClapModel, - text_encoder_2: Union[T5EncoderModel, VitsModel], + text_encoder_2: T5EncoderModel | VitsModel, projection_model: AudioLDM2ProjectionModel, language_model: GPT2LMHeadModel, - tokenizer: Union[RobertaTokenizer, RobertaTokenizerFast], - tokenizer_2: Union[T5Tokenizer, T5TokenizerFast, VitsTokenizer], + tokenizer: RobertaTokenizer | RobertaTokenizerFast, + tokenizer_2: T5Tokenizer | T5TokenizerFast | VitsTokenizer, feature_extractor: ClapFeatureExtractor, unet: AudioLDM2UNet2DConditionModel, scheduler: KarrasDiffusionSchedulers, @@ -251,7 +251,7 @@ def disable_vae_slicing(self): ) self.vae.disable_slicing() - def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = "cuda"): + def enable_model_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = "cuda"): r""" Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward` @@ -316,7 +316,7 @@ def generate_language_model( The sequence used as a prompt for the generation. max_new_tokens (`int`): Number of new tokens to generate. - model_kwargs (`Dict[str, Any]`, *optional*): + model_kwargs (`dict[str, Any]`, *optional*): Ad hoc parametrization of additional model-specific kwargs that will be forwarded to the `forward` function of the model. @@ -361,21 +361,21 @@ def encode_prompt( do_classifier_free_guidance, transcription=None, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - generated_prompt_embeds: Optional[torch.Tensor] = None, - negative_generated_prompt_embeds: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.LongTensor] = None, - negative_attention_mask: Optional[torch.LongTensor] = None, - max_new_tokens: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + generated_prompt_embeds: torch.Tensor | None = None, + negative_generated_prompt_embeds: torch.Tensor | None = None, + attention_mask: torch.LongTensor | None = None, + negative_attention_mask: torch.LongTensor | None = None, + max_new_tokens: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - transcription (`str` or `List[str]`): + transcription (`str` or `list[str]`): transcription of text to speech device (`torch.device`): torch device @@ -383,7 +383,7 @@ def encode_prompt( number of waveforms that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the audio generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -502,6 +502,10 @@ def encode_prompt( text_input_ids, attention_mask=attention_mask, ) + # Extract the pooler output if it's a BaseModelOutputWithPooling (Transformers v5+) + # otherwise use it directly (Transformers v4) + if hasattr(prompt_embeds, "pooler_output"): + prompt_embeds = prompt_embeds.pooler_output # append the seq-len dim: (bs, hidden_size) -> (bs, seq_len, hidden_size) prompt_embeds = prompt_embeds[:, None, :] # make sure that we attend to this single hidden-state @@ -569,7 +573,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -610,6 +614,10 @@ def encode_prompt( uncond_input_ids, attention_mask=negative_attention_mask, ) + # Extract the pooler output if it's a BaseModelOutputWithPooling (Transformers v5+) + # otherwise use it directly (Transformers v4) + if hasattr(negative_prompt_embeds, "pooler_output"): + negative_prompt_embeds = negative_prompt_embeds.pooler_output # append the seq-len dim: (bs, hidden_size) -> (bs, seq_len, hidden_size) negative_prompt_embeds = negative_prompt_embeds[:, None, :] # make sure that we attend to this single hidden-state @@ -862,36 +870,36 @@ def prepare_latents(self, batch_size, num_channels_latents, height, dtype, devic @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - transcription: Union[str, List[str]] = None, - audio_length_in_s: Optional[float] = None, + prompt: str | list[str] = None, + transcription: str | list[str] = None, + audio_length_in_s: float | None = None, num_inference_steps: int = 200, guidance_scale: float = 3.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_waveforms_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_waveforms_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - generated_prompt_embeds: Optional[torch.Tensor] = None, - negative_generated_prompt_embeds: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.LongTensor] = None, - negative_attention_mask: Optional[torch.LongTensor] = None, - max_new_tokens: Optional[int] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + generated_prompt_embeds: torch.Tensor | None = None, + negative_generated_prompt_embeds: torch.Tensor | None = None, + attention_mask: torch.LongTensor | None = None, + negative_attention_mask: torch.LongTensor | None = None, + max_new_tokens: int | None = None, return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: Optional[int] = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - output_type: Optional[str] = "np", + callback: Callable[[int, int, torch.Tensor], None] | None = None, + callback_steps: int | None = 1, + cross_attention_kwargs: dict[str, Any] | None = None, + output_type: str | None = "np", ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide audio generation. If not defined, you need to pass `prompt_embeds`. - transcription (`str` or `List[str]`, *optional*):\ + transcription (`str` or `list[str]`, *optional*):\ The transcript for text to speech. audio_length_in_s (`int`, *optional*, defaults to 10.24): The length of the generated audio sample in seconds. @@ -901,7 +909,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 3.5): A higher guidance scale value encourages the model to generate audio that is closely linked to the text `prompt` at the expense of lower sound quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in audio generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_waveforms_per_prompt (`int`, *optional*, defaults to 1): @@ -912,7 +920,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py b/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py index bb9884e41381..c609bc8b29f8 100644 --- a/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py +++ b/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from transformers import T5Tokenizer, UMT5EncoderModel @@ -63,10 +63,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -81,15 +81,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -231,25 +231,25 @@ def check_inputs( def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Union[str, List[str]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] = None, do_classifier_free_guidance: bool = True, num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, + device: torch.device | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, max_sequence_length: int = 256, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): @@ -429,37 +429,35 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, num_inference_steps: int = 50, - sigmas: List[float] = None, + sigmas: list[float] = None, guidance_scale: float = 3.5, - num_images_per_prompt: Optional[int] = 1, - height: Optional[int] = 1024, - width: Optional[int] = 1024, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, + num_images_per_prompt: int | None = 1, + height: int | None = 1024, + width: int | None = 1024, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, max_sequence_length: int = 256, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], - ) -> Union[ImagePipelineOutput, Tuple]: + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + ) -> ImagePipelineOutput | tuple: r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -470,7 +468,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. guidance_scale (`float`, *optional*, defaults to 5.0): @@ -481,7 +479,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -514,7 +512,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -586,7 +584,13 @@ def __call__( # 4. Prepare timesteps # sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, sigmas=sigmas) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, sigmas=sigmas + ) # 5. Prepare latents. latent_channels = self.transformer.config.in_channels diff --git a/src/diffusers/pipelines/auto_pipeline.py b/src/diffusers/pipelines/auto_pipeline.py index c14910250b54..2ea865f56306 100644 --- a/src/diffusers/pipelines/auto_pipeline.py +++ b/src/diffusers/pipelines/auto_pipeline.py @@ -21,7 +21,7 @@ from ..models.controlnets import ControlNetUnionModel from ..utils import is_sentencepiece_available from .aura_flow import AuraFlowPipeline -from .chroma import ChromaPipeline +from .chroma import ChromaPipeline, ChromaRadiancePipeline from .cogview3 import CogView3PlusPipeline from .cogview4 import CogView4ControlPipeline, CogView4Pipeline from .controlnet import ( @@ -52,6 +52,9 @@ FluxKontextPipeline, FluxPipeline, ) +from .flux2 import Flux2KleinPipeline, Flux2Pipeline +from .glm_image import GlmImagePipeline +from .helios import HeliosPipeline, HeliosPyramidPipeline from .hunyuandit import HunyuanDiTPipeline from .kandinsky import ( KandinskyCombinedPipeline, @@ -99,6 +102,7 @@ QwenImageEditPlusPipeline, QwenImageImg2ImgPipeline, QwenImageInpaintPipeline, + QwenImageLayeredPipeline, QwenImagePipeline, ) from .sana import SanaPipeline @@ -124,6 +128,7 @@ ZImageControlNetInpaintPipeline, ZImageControlNetPipeline, ZImageImg2ImgPipeline, + ZImageInpaintPipeline, ZImageOmniPipeline, ZImagePipeline, ) @@ -162,11 +167,17 @@ ("flux-control", FluxControlPipeline), ("flux-controlnet", FluxControlNetPipeline), ("flux-kontext", FluxKontextPipeline), + ("flux2-klein", Flux2KleinPipeline), + ("flux2", Flux2Pipeline), ("lumina", LuminaPipeline), ("lumina2", Lumina2Pipeline), ("chroma", ChromaPipeline), + ("chroma-radiance", ChromaRadiancePipeline), ("cogview3", CogView3PlusPipeline), ("cogview4", CogView4Pipeline), + ("glm_image", GlmImagePipeline), + ("helios", HeliosPipeline), + ("helios-pyramid", HeliosPyramidPipeline), ("cogview4-control", CogView4ControlPipeline), ("qwenimage", QwenImagePipeline), ("qwenimage-controlnet", QwenImageControlNetPipeline), @@ -199,9 +210,12 @@ ("flux-controlnet", FluxControlNetImg2ImgPipeline), ("flux-control", FluxControlImg2ImgPipeline), ("flux-kontext", FluxKontextPipeline), + ("flux2-klein", Flux2KleinPipeline), + ("flux2", Flux2Pipeline), ("qwenimage", QwenImageImg2ImgPipeline), ("qwenimage-edit", QwenImageEditPipeline), ("qwenimage-edit-plus", QwenImageEditPlusPipeline), + ("qwenimage-layered", QwenImageLayeredPipeline), ("z-image", ZImageImg2ImgPipeline), ] ) @@ -226,6 +240,7 @@ ("stable-diffusion-pag", StableDiffusionPAGInpaintPipeline), ("qwenimage", QwenImageInpaintPipeline), ("qwenimage-edit", QwenImageEditInpaintPipeline), + ("z-image", ZImageInpaintPipeline), ] ) @@ -237,7 +252,7 @@ AUTO_IMAGE2VIDEO_PIPELINES_MAPPING = OrderedDict( [ - ("wan", WanImageToVideoPipeline), + ("wan-i2v", WanImageToVideoPipeline), ] ) @@ -385,11 +400,11 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs): force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. output_loading_info(`bool`, *optional*, defaults to `False`): @@ -411,7 +426,7 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs): Mirror source to resolve accessibility issues if you’re downloading a model in China. We do not guarantee the timeliness or safety of the source, and you should refer to the mirror site for more information. - device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*): + device_map (`str` or `dict[str, int | str | torch.device]`, *optional*): A map that specifies where each submodule should go. It doesn’t need to be defined for each parameter/buffer name; once a given module name is inside, every submodule of it will be sent to the same device. @@ -676,11 +691,11 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs): force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. output_loading_info(`bool`, *optional*, defaults to `False`): @@ -702,7 +717,7 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs): Mirror source to resolve accessibility issues if you’re downloading a model in China. We do not guarantee the timeliness or safety of the source, and you should refer to the mirror site for more information. - device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*): + device_map (`str` or `dict[str, int | str | torch.device]`, *optional*): A map that specifies where each submodule should go. It doesn’t need to be defined for each parameter/buffer name; once a given module name is inside, every submodule of it will be sent to the same device. @@ -982,11 +997,11 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs): force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. output_loading_info(`bool`, *optional*, defaults to `False`): @@ -1008,7 +1023,7 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs): Mirror source to resolve accessibility issues if you’re downloading a model in China. We do not guarantee the timeliness or safety of the source, and you should refer to the mirror site for more information. - device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*): + device_map (`str` or `dict[str, int | str | torch.device]`, *optional*): A map that specifies where each submodule should go. It doesn’t need to be defined for each parameter/buffer name; once a given module name is inside, every submodule of it will be sent to the same device. diff --git a/src/diffusers/pipelines/blip_diffusion/__init__.py b/src/diffusers/pipelines/blip_diffusion/__init__.py index af6c879d5ce8..c245313e2f8a 100644 --- a/src/diffusers/pipelines/blip_diffusion/__init__.py +++ b/src/diffusers/pipelines/blip_diffusion/__init__.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Optional, Union import numpy as np import PIL diff --git a/src/diffusers/pipelines/blip_diffusion/blip_image_processing.py b/src/diffusers/pipelines/blip_diffusion/blip_image_processing.py index e45f431d0b9d..5d2a0186f041 100644 --- a/src/diffusers/pipelines/blip_diffusion/blip_image_processing.py +++ b/src/diffusers/pipelines/blip_diffusion/blip_image_processing.py @@ -14,8 +14,6 @@ # limitations under the License. """Image processor class for BLIP.""" -from typing import Dict, List, Optional, Union - import numpy as np import torch from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict @@ -69,11 +67,11 @@ class BlipImageProcessor(BaseImageProcessor): do_normalize (`bool`, *optional*, defaults to `True`): Whether to normalize the image. Can be overridden by the `do_normalize` parameter in the `preprocess` method. Can be overridden by the `do_normalize` parameter in the `preprocess` method. - image_mean (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_MEAN`): + image_mean (`float` or `list[float]`, *optional*, defaults to `IMAGENET_STANDARD_MEAN`): Mean to use if normalizing the image. This is a float or list of floats the length of the number of channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method. Can be overridden by the `image_mean` parameter in the `preprocess` method. - image_std (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_STD`): + image_std (`float` or `list[float]`, *optional*, defaults to `IMAGENET_STANDARD_STD`): Standard deviation to use if normalizing the image. This is a float or list of floats the length of the number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method. Can be overridden by the `image_std` parameter in the `preprocess` method. @@ -86,13 +84,13 @@ class BlipImageProcessor(BaseImageProcessor): def __init__( self, do_resize: bool = True, - size: Dict[str, int] = None, + size: dict[str, int] = None, resample: PILImageResampling = PILImageResampling.BICUBIC, do_rescale: bool = True, - rescale_factor: Union[int, float] = 1 / 255, + rescale_factor: int | float = 1 / 255, do_normalize: bool = True, - image_mean: Optional[Union[float, List[float]]] = None, - image_std: Optional[Union[float, List[float]]] = None, + image_mean: float | list[float] | None = None, + image_std: float | list[float] | None = None, do_convert_rgb: bool = True, do_center_crop: bool = True, **kwargs, @@ -116,10 +114,10 @@ def __init__( def resize( self, image: np.ndarray, - size: Dict[str, int], + size: dict[str, int], resample: PILImageResampling = PILImageResampling.BICUBIC, - data_format: Optional[Union[str, ChannelDimension]] = None, - input_data_format: Optional[Union[str, ChannelDimension]] = None, + data_format: str | ChannelDimension | None = None, + input_data_format: str | ChannelDimension | None = None, **kwargs, ) -> np.ndarray: """ @@ -128,7 +126,7 @@ def resize( Args: image (`np.ndarray`): Image to resize. - size (`Dict[str, int]`): + size (`dict[str, int]`): Dictionary in the format `{"height": int, "width": int}` specifying the size of the output image. resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BICUBIC`): `PILImageResampling` filter to use when resizing the image e.g. `PILImageResampling.BICUBIC`. @@ -164,19 +162,19 @@ def resize( def preprocess( self, images: ImageInput, - do_resize: Optional[bool] = None, - size: Optional[Dict[str, int]] = None, + do_resize: bool | None = None, + size: dict[str, int] | None = None, resample: PILImageResampling = None, - do_rescale: Optional[bool] = None, - do_center_crop: Optional[bool] = None, - rescale_factor: Optional[float] = None, - do_normalize: Optional[bool] = None, - image_mean: Optional[Union[float, List[float]]] = None, - image_std: Optional[Union[float, List[float]]] = None, - return_tensors: Optional[Union[str, TensorType]] = None, + do_rescale: bool | None = None, + do_center_crop: bool | None = None, + rescale_factor: float | None = None, + do_normalize: bool | None = None, + image_mean: float | list[float] | None = None, + image_std: float | list[float] | None = None, + return_tensors: str | TensorType | None = None, do_convert_rgb: bool = None, data_format: ChannelDimension = ChannelDimension.FIRST, - input_data_format: Optional[Union[str, ChannelDimension]] = None, + input_data_format: str | ChannelDimension | None = None, **kwargs, ) -> PIL.Image.Image: """ @@ -188,7 +186,7 @@ def preprocess( passing in images with pixel values between 0 and 1, set `do_rescale=False`. do_resize (`bool`, *optional*, defaults to `self.do_resize`): Whether to resize the image. - size (`Dict[str, int]`, *optional*, defaults to `self.size`): + size (`dict[str, int]`, *optional*, defaults to `self.size`): Controls the size of the image after `resize`. The shortest edge of the image is resized to `size["shortest_edge"]` whilst preserving the aspect ratio. If the longest edge of this resized image is > `int(size["shortest_edge"] * (1333 / 800))`, then the image is resized again to make the longest @@ -201,9 +199,9 @@ def preprocess( Rescale factor to rescale the image by if `do_rescale` is set to `True`. do_normalize (`bool`, *optional*, defaults to `self.do_normalize`): Whether to normalize the image. - image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`): + image_mean (`float` or `list[float]`, *optional*, defaults to `self.image_mean`): Image mean to normalize the image by if `do_normalize` is set to `True`. - image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`): + image_std (`float` or `list[float]`, *optional*, defaults to `self.image_std`): Image standard deviation to normalize the image by if `do_normalize` is set to `True`. do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`): Whether to convert the image to RGB. diff --git a/src/diffusers/pipelines/blip_diffusion/modeling_blip2.py b/src/diffusers/pipelines/blip_diffusion/modeling_blip2.py index b061ac2636a5..c434ccdaccca 100644 --- a/src/diffusers/pipelines/blip_diffusion/modeling_blip2.py +++ b/src/diffusers/pipelines/blip_diffusion/modeling_blip2.py @@ -11,8 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Tuple, Union - import torch from torch import nn from transformers import BertTokenizer @@ -372,11 +370,11 @@ def __init__(self, config: Blip2VisionConfig): @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=Blip2VisionConfig) def forward( self, - pixel_values: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple, BaseModelOutputWithPooling]: + pixel_values: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> tuple | BaseModelOutputWithPooling: r""" Returns: @@ -464,7 +462,7 @@ class PreTrainedModel def get_extended_attention_mask( self, attention_mask: torch.Tensor, - input_shape: Tuple[int], + input_shape: tuple[int], device: torch.device, has_query: bool = False, ) -> torch.Tensor: @@ -474,7 +472,7 @@ def get_extended_attention_mask( Arguments: attention_mask (`torch.Tensor`): Mask with ones indicating tokens to attend to, zeros for tokens to ignore. - input_shape (`Tuple[int]`): + input_shape (`tuple[int]`): The shape of the input to the model. device (`torch.device`): The device of the input to the model. diff --git a/src/diffusers/pipelines/blip_diffusion/modeling_ctx_clip.py b/src/diffusers/pipelines/blip_diffusion/modeling_ctx_clip.py index 1b0342ce7a56..c5364f8985aa 100644 --- a/src/diffusers/pipelines/blip_diffusion/modeling_ctx_clip.py +++ b/src/diffusers/pipelines/blip_diffusion/modeling_ctx_clip.py @@ -12,8 +12,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Tuple, Union - import torch from torch import nn from transformers import CLIPPreTrainedModel @@ -22,7 +20,7 @@ from transformers.models.clip.modeling_clip import CLIPEncoder -def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): +def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: int | None = None): """ Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`. """ @@ -54,13 +52,13 @@ def forward( self, ctx_embeddings: torch.Tensor = None, ctx_begin_pos: list = None, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple, BaseModelOutputWithPooling]: + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> tuple | BaseModelOutputWithPooling: return self.text_model( ctx_embeddings=ctx_embeddings, ctx_begin_pos=ctx_begin_pos, @@ -86,13 +84,13 @@ def forward( self, ctx_embeddings: torch.Tensor, ctx_begin_pos: list, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple, BaseModelOutputWithPooling]: + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> tuple | BaseModelOutputWithPooling: r""" Returns: @@ -184,9 +182,9 @@ def forward( self, ctx_embeddings: torch.Tensor, ctx_begin_pos: list, - input_ids: Optional[torch.LongTensor] = None, - position_ids: Optional[torch.LongTensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, + input_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.Tensor | None = None, ) -> torch.Tensor: if ctx_embeddings is None: ctx_len = 0 diff --git a/src/diffusers/pipelines/blip_diffusion/pipeline_blip_diffusion.py b/src/diffusers/pipelines/blip_diffusion/pipeline_blip_diffusion.py index 705d930b59fe..aa3dbdae966b 100644 --- a/src/diffusers/pipelines/blip_diffusion/pipeline_blip_diffusion.py +++ b/src/diffusers/pipelines/blip_diffusion/pipeline_blip_diffusion.py @@ -11,8 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Optional, Union - import PIL.Image import torch from transformers import CLIPTokenizer @@ -116,8 +114,8 @@ def __init__( qformer: Blip2QFormerModel, image_processor: BlipImageProcessor, ctx_begin_pos: int = 2, - mean: List[float] = None, - std: List[float] = None, + mean: list[float] = None, + std: list[float] = None, ): super().__init__() @@ -193,33 +191,33 @@ def encode_prompt(self, query_embeds, prompt, device=None): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: List[str], + prompt: list[str], reference_image: PIL.Image.Image, - source_subject_category: List[str], - target_subject_category: List[str], - latents: Optional[torch.Tensor] = None, + source_subject_category: list[str], + target_subject_category: list[str], + latents: torch.Tensor | None = None, guidance_scale: float = 7.5, height: int = 512, width: int = 512, num_inference_steps: int = 50, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - neg_prompt: Optional[str] = "", + generator: torch.Generator | list[torch.Generator] | None = None, + neg_prompt: str | None = "", prompt_strength: float = 1.0, prompt_reps: int = 20, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, ): """ Function invoked when calling the pipeline for generation. Args: - prompt (`List[str]`): + prompt (`list[str]`): The prompt or prompts to guide the image generation. reference_image (`PIL.Image.Image`): The reference image to condition the generation on. - source_subject_category (`List[str]`): + source_subject_category (`list[str]`): The source subject category. - target_subject_category (`List[str]`): + target_subject_category (`list[str]`): The target subject category. latents (`torch.Tensor`, *optional*): Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image @@ -238,7 +236,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. neg_prompt (`str`, *optional*, defaults to ""): diff --git a/src/diffusers/pipelines/bria/pipeline_bria.py b/src/diffusers/pipelines/bria/pipeline_bria.py index a22a756005ac..95ae9ce96e7e 100644 --- a/src/diffusers/pipelines/bria/pipeline_bria.py +++ b/src/diffusers/pipelines/bria/pipeline_bria.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -116,7 +116,7 @@ class BriaPipeline(DiffusionPipeline): def __init__( self, transformer: BriaTransformer2DModel, - scheduler: Union[FlowMatchEulerDiscreteScheduler, KarrasDiffusionSchedulers], + scheduler: FlowMatchEulerDiscreteScheduler | KarrasDiffusionSchedulers, vae: AutoencoderKL, text_encoder: T5EncoderModel, tokenizer: T5TokenizerFast, @@ -145,20 +145,20 @@ def __init__( def encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt: str | list[str] | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 128, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -166,7 +166,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -320,10 +320,10 @@ def check_inputs( def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 128, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): tokenizer = self.tokenizer text_encoder = self.text_encoder @@ -449,97 +449,109 @@ def _prepare_latent_image_ids(batch_size, height, width, device, dtype): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 30, - timesteps: List[int] = None, + timesteps: list[int] = None, guidance_scale: float = 5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 128, - clip_value: Union[None, float] = None, + clip_value: None | float = None, normalize: bool = False, ): r""" - Function invoked when calling the pipeline for generation. - - Args: - prompt (`str` or `List[str]`, *optional*): - The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. - instead. - height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): - The height in pixels of the generated image. This is set to 1024 by default for the best results. - width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): - The width in pixels of the generated image. This is set to 1024 by default for the best results. - num_inference_steps (`int`, *optional*, defaults to 50): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. - timesteps (`List[int]`, *optional*): - Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument - in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is - passed will be used. Must be in descending order. - guidance_scale (`float`, *optional*, defaults to 5.0): - Guidance scale as defined in [Classifier-Free Diffusion - Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. - of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting - `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to - the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): - The prompt or prompts not to guide the image generation. If not defined, one has to pass - `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is - less than `1`). - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): - One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) - to make generation deterministic. - latents (`torch.FloatTensor`, *optional*): - Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image - generation. Can be used to tweak the same generation with different prompts. If not provided, a latents - tensor will be generated by sampling using the supplied random `generator`. - prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not - provided, text embeddings will be generated from `prompt` input argument. - negative_prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input - argument. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between - [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.bria.BriaPipelineOutput`] instead of a plain tuple. - attention_kwargs (`dict`, *optional*): - A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under - `self.processor` in - [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - callback_on_step_end (`Callable`, *optional*): - A function that calls at the end of each denoising steps during the inference. The function is called - with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, - callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by - `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): - The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list - will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the - `._callback_tensor_inputs` attribute of your pipeline class. - max_sequence_length (`int` defaults to 256): Maximum sequence length to use with the `prompt`. - - Examples: - - Returns: - [`~pipelines.bria.BriaPipelineOutput`] or `tuple`: [`~pipelines.bria.BriaPipelineOutput`] if `return_dict` - is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated - images. + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `list[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass + `prompt_embeds`. instead. + height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The height in pixels of the generated image. This is set to 1024 by default for the best + results. + width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The width in pixels of the generated image. This is set to 1024 by default for the best + results. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at + the expense of slower inference. + timesteps (`list[int]`, *optional*): + Custom timesteps to use for the denoising process with schedulers which support a `timesteps` + argument in their `set_timesteps` method. If not defined, the default behavior when + `num_inference_steps` is passed will be used. Must be in descending order. + guidance_scale (`float`, *optional*, defaults to 5.0): + <<<<<<< HEAD + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://arxiv.org/abs/2207.12598). `guidance_scale` is defined as `w` of equation 2. + of [Imagen Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting + `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely + linked to the text `prompt`, usually at the expense of lower image quality. + negative_prompt (`str` or `list[str]`, *optional*): + ======= + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of + equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is + enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images + that are closely linked to the text `prompt`, usually at the expense of lower image quality. + negative_prompt (`str` or `list[str]`, *optional*): + >>>>>>> main + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if + `guidance_scale` is less than `1`). + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + One or a list of [torch + generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make + generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for + image generation. Can be used to tweak the same generation with different prompts. If not + provided, a latents tensor will be generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` + input argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.bria.BriaPipelineOutput`] instead of a plain tuple. + attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined + under `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is + called with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, + timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as + specified by `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`list`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the + list will be passed as `callback_kwargs` argument. You will only be able to include variables + listed in the `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int` defaults to 256): Maximum sequence length to use with the `prompt`. + + Examples: + + Returns: + [`~pipelines.bria.BriaPipelineOutput`] or `tuple`: [`~pipelines.bria.BriaPipelineOutput`] if + `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list + with the generated images. """ height = height or self.default_sample_size * self.vae_scale_factor diff --git a/src/diffusers/pipelines/bria/pipeline_output.py b/src/diffusers/pipelines/bria/pipeline_output.py index 54eed0623371..4bf2ed949c26 100644 --- a/src/diffusers/pipelines/bria/pipeline_output.py +++ b/src/diffusers/pipelines/bria/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -13,9 +12,9 @@ class BriaPipelineOutput(BaseOutput): Output class for Bria pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image] | np.ndarray diff --git a/src/diffusers/pipelines/bria_fibo/__init__.py b/src/diffusers/pipelines/bria_fibo/__init__.py index 206a463b394b..8dd77270902c 100644 --- a/src/diffusers/pipelines/bria_fibo/__init__.py +++ b/src/diffusers/pipelines/bria_fibo/__init__.py @@ -23,6 +23,8 @@ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) else: _import_structure["pipeline_bria_fibo"] = ["BriaFiboPipeline"] + _import_structure["pipeline_bria_fibo_edit"] = ["BriaFiboEditPipeline"] + if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: @@ -33,6 +35,7 @@ from ...utils.dummy_torch_and_transformers_objects import * else: from .pipeline_bria_fibo import BriaFiboPipeline + from .pipeline_bria_fibo_edit import BriaFiboEditPipeline else: import sys diff --git a/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py index 8fd29756b290..1f178066b17d 100644 --- a/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py +++ b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py @@ -8,7 +8,7 @@ # # See the license for further details. -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -94,7 +94,7 @@ class BriaFiboPipeline(DiffusionPipeline, FluxLoraLoaderMixin): def __init__( self, transformer: BriaFiboTransformer2DModel, - scheduler: Union[FlowMatchEulerDiscreteScheduler, KarrasDiffusionSchedulers], + scheduler: FlowMatchEulerDiscreteScheduler | KarrasDiffusionSchedulers, vae: AutoencoderKLWan, text_encoder: SmolLM3ForCausalLM, tokenizer: AutoTokenizer, @@ -113,11 +113,11 @@ def __init__( def get_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, max_sequence_length: int = 2048, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -200,19 +200,19 @@ def pad_embedding(prompt_embeds, max_tokens, attention_mask=None): def encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + device: torch.device | None = None, num_images_per_prompt: int = 1, guidance_scale: float = 5, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt: str | list[str] | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 3000, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -220,7 +220,7 @@ def encode_prompt( number of images that should be generated per prompt guidance_scale (`float`): Guidance scale for classifier free guidance. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -459,23 +459,23 @@ def _prepare_attention_mask(attention_mask): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 30, - timesteps: List[int] = None, + timesteps: list[int] = None, guidance_scale: float = 5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 3000, do_patching=False, ): @@ -483,7 +483,7 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): @@ -493,7 +493,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. @@ -503,13 +503,13 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): diff --git a/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py new file mode 100644 index 000000000000..c2327bbce1c7 --- /dev/null +++ b/src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo_edit.py @@ -0,0 +1,1133 @@ +# Copyright (c) Bria.ai. All rights reserved. +# +# This file is licensed under the Creative Commons Attribution-NonCommercial 4.0 International Public License (CC-BY-NC-4.0). +# You may obtain a copy of the license at https://creativecommons.org/licenses/by-nc/4.0/ +# +# You are free to share and adapt this material for non-commercial purposes provided you give appropriate credit, +# indicate if changes were made, and do not use the material for commercial purposes. +# +# See the license for further details. + +import json +from typing import Any, Callable, Dict, List, Optional, Union + +import numpy as np +import torch +from PIL import Image +from transformers import AutoTokenizer +from transformers.models.smollm3.modeling_smollm3 import SmolLM3ForCausalLM + +from ...image_processor import PipelineImageInput, VaeImageProcessor +from ...loaders import FluxLoraLoaderMixin +from ...models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan +from ...models.transformers.transformer_bria_fibo import BriaFiboTransformer2DModel +from ...pipelines.bria_fibo.pipeline_output import BriaFiboPipelineOutput +from ...pipelines.flux.pipeline_flux import calculate_shift, retrieve_timesteps +from ...pipelines.pipeline_utils import DiffusionPipeline +from ...schedulers import FlowMatchEulerDiscreteScheduler, KarrasDiffusionSchedulers +from ...utils import ( + USE_PEFT_BACKEND, + is_torch_xla_available, + logging, + replace_example_docstring, + scale_lora_layers, + unscale_lora_layers, +) +from ...utils.torch_utils import randn_tensor + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +PipelineMaskInput = Union[ + torch.FloatTensor, Image.Image, List[Image.Image], List[torch.FloatTensor], np.ndarray, List[np.ndarray] +] + +# TODO: Update example docstring +EXAMPLE_DOC_STRING = """ + Example: + ```python + import torch + from diffusers import BriaFiboEditPipeline + from diffusers.modular_pipelines import ModularPipeline + + torch.set_grad_enabled(False) + vlm_pipe = ModularPipelineBlocks.from_pretrained("briaai/FIBO-VLM-prompt-to-JSON", trust_remote_code=True) + vlm_pipe = vlm_pipe.init_pipeline() + + pipe = BriaFiboEditPipeline.from_pretrained( + "briaai/fibo-edit", + torch_dtype=torch.bfloat16, + ) + pipe.to("cuda") + + output = vlm_pipe( + prompt="A hyper-detailed, ultra-fluffy owl sitting in the trees at night, looking directly at the camera with wide, adorable, expressive eyes. Its feathers are soft and voluminous, catching the cool moonlight with subtle silver highlights. The owl's gaze is curious and full of charm, giving it a whimsical, storybook-like personality." + ) + json_prompt_generate = json.loads(output.values["json_prompt"]) + + image = Image.open("image_generate.png") + + edit_prompt = "Make the owl to be a cat" + + json_prompt_generate["edit_instruction"] = edit_prompt + + results_generate = pipe( + prompt=json_prompt_generate, num_inference_steps=50, guidance_scale=3.5, image=image, output_type="np" + ) + ``` +""" + +PREFERRED_RESOLUTION = { + 256 * 256: [(208, 304), (224, 288), (256, 256), (288, 224), (304, 208), (320, 192), (336, 192)], + 512 * 512: [ + (416, 624), + (432, 592), + (464, 560), + (512, 512), + (544, 480), + (576, 448), + (592, 432), + (608, 416), + (624, 416), + (640, 400), + (672, 384), + (704, 368), + ], + 1024 * 1024: [ + (832, 1248), + (880, 1184), + (912, 1136), + (1024, 1024), + (1136, 912), + (1184, 880), + (1216, 848), + (1248, 832), + (1248, 832), + (1264, 816), + (1296, 800), + (1360, 768), + ], +} + + +def is_valid_edit_json(json_input: str | dict): + """ + Check if the input is a valid JSON string or dict with an "edit_instruction" key. + + Args: + json_input (`str` or `dict`): + The JSON string or dict to check. + + Returns: + `bool`: True if the input is a valid JSON string or dict with an "edit_instruction" key, False otherwise. + """ + try: + if isinstance(json_input, str) and "edit_instruction" in json_input: + json.loads(json_input) + return True + elif isinstance(json_input, dict) and "edit_instruction" in json_input: + return True + else: + return False + except json.JSONDecodeError: + return False + + +def is_valid_mask(mask: PipelineMaskInput): + """ + Check if the mask is a valid mask. + """ + if isinstance(mask, torch.Tensor): + return True + elif isinstance(mask, Image.Image): + return True + elif isinstance(mask, list): + return all(isinstance(m, (torch.Tensor, Image.Image, np.ndarray)) for m in mask) + elif isinstance(mask, np.ndarray): + return mask.ndim in [2, 3] and mask.min() >= 0 and mask.max() <= 1 + else: + return False + + +def get_mask_size(mask: PipelineMaskInput): + """ + Get the size of the mask. + """ + if isinstance(mask, torch.Tensor): + return mask.shape[-2:] + elif isinstance(mask, Image.Image): + return mask.size[::-1] # (height, width) + elif isinstance(mask, list): + return [get_mask_size(m) for m in mask] + elif isinstance(mask, np.ndarray): + return mask.shape[-2:] + else: + return None + + +def get_image_size(image: PipelineImageInput): + """ + Get the size of the image. + """ + if isinstance(image, torch.Tensor): + return image.shape[-2:] + elif isinstance(image, Image.Image): + return image.size[::-1] # (height, width) + elif isinstance(image, list): + return [get_image_size(i) for i in image] + else: + return None + + +def paste_mask_on_image(mask: PipelineMaskInput, image: PipelineImageInput): + """convert mask and image to PIL Images and paste the mask on the image""" + if isinstance(mask, torch.Tensor): + if mask.ndim == 3 and mask.shape[0] == 1: + mask = mask.squeeze(0) + mask = Image.fromarray((mask.cpu().numpy() * 255).astype(np.uint8)) + elif isinstance(mask, Image.Image): + pass + elif isinstance(mask, list): + mask = mask[0] + if isinstance(mask, torch.Tensor): + if mask.ndim == 3 and mask.shape[0] == 1: + mask = mask.squeeze(0) + mask = Image.fromarray((mask.cpu().numpy() * 255).astype(np.uint8)) + elif isinstance(mask, np.ndarray): + mask = Image.fromarray((mask * 255).astype(np.uint8)) + elif isinstance(mask, np.ndarray): + mask = Image.fromarray((mask * 255).astype(np.uint8)) + + if isinstance(image, torch.Tensor): + if image.ndim == 3: + image = image.permute(1, 2, 0) + image = Image.fromarray((image.cpu().numpy() * 255).astype(np.uint8)) + elif isinstance(image, Image.Image): + pass + elif isinstance(image, list): + image = image[0] + if isinstance(image, torch.Tensor): + if image.ndim == 3: + image = image.permute(1, 2, 0) + image = Image.fromarray((image.cpu().numpy() * 255).astype(np.uint8)) + elif isinstance(image, np.ndarray): + image = Image.fromarray((image * 255).astype(np.uint8)) + elif isinstance(image, np.ndarray): + image = Image.fromarray((image * 255).astype(np.uint8)) + + mask = mask.convert("L") + image = image.convert("RGB") + gray_color = (128, 128, 128) + gray_img = Image.new("RGB", image.size, gray_color) + image = Image.composite(gray_img, image, mask) + return image + + +class BriaFiboEditPipeline(DiffusionPipeline, FluxLoraLoaderMixin): + r""" + Args: + transformer (`BriaFiboTransformer2DModel`): + The transformer model for 2D diffusion modeling. + scheduler (`FlowMatchEulerDiscreteScheduler` or `KarrasDiffusionSchedulers`): + Scheduler to be used with `transformer` to denoise the encoded latents. + vae (`AutoencoderKLWan`): + Variational Auto-Encoder for encoding and decoding images to and from latent representations. + text_encoder (`SmolLM3ForCausalLM`): + Text encoder for processing input prompts. + tokenizer (`AutoTokenizer`): + Tokenizer used for processing the input text prompts for the text_encoder. + """ + + model_cpu_offload_seq = "text_encoder->image_encoder->transformer->vae" + _callback_tensor_inputs = ["latents", "prompt_embeds"] + + def __init__( + self, + transformer: BriaFiboTransformer2DModel, + scheduler: Union[FlowMatchEulerDiscreteScheduler, KarrasDiffusionSchedulers], + vae: AutoencoderKLWan, + text_encoder: SmolLM3ForCausalLM, + tokenizer: AutoTokenizer, + ): + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + scheduler=scheduler, + ) + + self.vae_scale_factor = 16 + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) # * 2) + self.default_sample_size = 32 # 64 + + def get_prompt_embeds( + self, + prompt: Union[str, List[str]], + num_images_per_prompt: int = 1, + max_sequence_length: int = 2048, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + if not prompt: + raise ValueError("`prompt` must be a non-empty string or list of strings.") + + batch_size = len(prompt) + bot_token_id = 128000 + + text_encoder_device = device if device is not None else torch.device("cpu") + if not isinstance(text_encoder_device, torch.device): + text_encoder_device = torch.device(text_encoder_device) + + if all(p == "" for p in prompt): + input_ids = torch.full((batch_size, 1), bot_token_id, dtype=torch.long, device=text_encoder_device) + attention_mask = torch.ones_like(input_ids) + else: + tokenized = self.tokenizer( + prompt, + padding="longest", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + input_ids = tokenized.input_ids.to(text_encoder_device) + attention_mask = tokenized.attention_mask.to(text_encoder_device) + + if any(p == "" for p in prompt): + empty_rows = torch.tensor([p == "" for p in prompt], dtype=torch.bool, device=text_encoder_device) + input_ids[empty_rows] = bot_token_id + attention_mask[empty_rows] = 1 + + encoder_outputs = self.text_encoder( + input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + ) + hidden_states = encoder_outputs.hidden_states + + prompt_embeds = torch.cat([hidden_states[-1], hidden_states[-2]], dim=-1) + prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + + prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) + hidden_states = tuple( + layer.repeat_interleave(num_images_per_prompt, dim=0).to(device=device) for layer in hidden_states + ) + attention_mask = attention_mask.repeat_interleave(num_images_per_prompt, dim=0).to(device=device) + + return prompt_embeds, hidden_states, attention_mask + + @staticmethod + def pad_embedding(prompt_embeds, max_tokens, attention_mask=None): + # Pad embeddings to `max_tokens` while preserving the mask of real tokens. + batch_size, seq_len, dim = prompt_embeds.shape + + if attention_mask is None: + attention_mask = torch.ones((batch_size, seq_len), dtype=prompt_embeds.dtype, device=prompt_embeds.device) + else: + attention_mask = attention_mask.to(device=prompt_embeds.device, dtype=prompt_embeds.dtype) + + if max_tokens < seq_len: + raise ValueError("`max_tokens` must be greater or equal to the current sequence length.") + + if max_tokens > seq_len: + pad_length = max_tokens - seq_len + padding = torch.zeros( + (batch_size, pad_length, dim), dtype=prompt_embeds.dtype, device=prompt_embeds.device + ) + prompt_embeds = torch.cat([prompt_embeds, padding], dim=1) + + mask_padding = torch.zeros( + (batch_size, pad_length), dtype=prompt_embeds.dtype, device=prompt_embeds.device + ) + attention_mask = torch.cat([attention_mask, mask_padding], dim=1) + + return prompt_embeds, attention_mask + + def encode_prompt( + self, + prompt: Union[str, List[str]], + device: torch.device | None = None, + num_images_per_prompt: int = 1, + guidance_scale: float = 5, + negative_prompt: Optional[Union[str, List[str]]] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + max_sequence_length: int = 3000, + lora_scale: bool | None = None, + ): + r""" + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + device: (`torch.device`): + torch device + num_images_per_prompt (`int`): + number of images that should be generated per prompt + guidance_scale (`float`): + Guidance scale for classifier free guidance. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + """ + device = device or self._execution_device + + # set lora scale so that monkey patched LoRA + # function of text encoder can correctly access it + if lora_scale is not None and isinstance(self, FluxLoraLoaderMixin): + self._lora_scale = lora_scale + + # dynamically adjust the LoRA scale + if self.text_encoder is not None and USE_PEFT_BACKEND: + scale_lora_layers(self.text_encoder, lora_scale) + + prompt = [prompt] if isinstance(prompt, str) else prompt + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + prompt_attention_mask = None + negative_prompt_attention_mask = None + if prompt_embeds is None: + prompt_embeds, prompt_layers, prompt_attention_mask = self.get_prompt_embeds( + prompt=prompt, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + ) + prompt_embeds = prompt_embeds.to(dtype=self.transformer.dtype) + prompt_layers = [tensor.to(dtype=self.transformer.dtype) for tensor in prompt_layers] + + if guidance_scale > 1: + if isinstance(negative_prompt, list) and negative_prompt[0] is None: + negative_prompt = "" + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + if prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + + negative_prompt_embeds, negative_prompt_layers, negative_prompt_attention_mask = self.get_prompt_embeds( + prompt=negative_prompt, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + ) + negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.transformer.dtype) + negative_prompt_layers = [tensor.to(dtype=self.transformer.dtype) for tensor in negative_prompt_layers] + + if self.text_encoder is not None: + if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND: + # Retrieve the original scale by scaling back the LoRA layers + unscale_lora_layers(self.text_encoder, lora_scale) + + # Pad to longest + if prompt_attention_mask is not None: + prompt_attention_mask = prompt_attention_mask.to(device=prompt_embeds.device, dtype=prompt_embeds.dtype) + + if negative_prompt_embeds is not None: + if negative_prompt_attention_mask is not None: + negative_prompt_attention_mask = negative_prompt_attention_mask.to( + device=negative_prompt_embeds.device, dtype=negative_prompt_embeds.dtype + ) + max_tokens = max(negative_prompt_embeds.shape[1], prompt_embeds.shape[1]) + + prompt_embeds, prompt_attention_mask = self.pad_embedding( + prompt_embeds, max_tokens, attention_mask=prompt_attention_mask + ) + prompt_layers = [self.pad_embedding(layer, max_tokens)[0] for layer in prompt_layers] + + negative_prompt_embeds, negative_prompt_attention_mask = self.pad_embedding( + negative_prompt_embeds, max_tokens, attention_mask=negative_prompt_attention_mask + ) + negative_prompt_layers = [self.pad_embedding(layer, max_tokens)[0] for layer in negative_prompt_layers] + else: + max_tokens = prompt_embeds.shape[1] + prompt_embeds, prompt_attention_mask = self.pad_embedding( + prompt_embeds, max_tokens, attention_mask=prompt_attention_mask + ) + negative_prompt_layers = None + + dtype = self.text_encoder.dtype + text_ids = torch.zeros(prompt_embeds.shape[0], max_tokens, 3).to(device=device, dtype=dtype) + + return ( + prompt_embeds, + negative_prompt_embeds, + text_ids, + prompt_attention_mask, + negative_prompt_attention_mask, + prompt_layers, + negative_prompt_layers, + ) + + @property + def guidance_scale(self): + return self._guidance_scale + + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://huggingface.co/papers/2205.11487 . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + + @property + def joint_attention_kwargs(self): + return self._joint_attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def interrupt(self): + return self._interrupt + + @staticmethod + # Based on diffusers.pipelines.flux.pipeline_flux.FluxPipeline._unpack_latents + def _unpack_latents(latents, height, width, vae_scale_factor): + batch_size, num_patches, channels = latents.shape + + height = height // vae_scale_factor + width = width // vae_scale_factor + + latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2) + latents = latents.permute(0, 3, 1, 4, 2, 5) + + latents = latents.reshape(batch_size, channels // (2 * 2), height, width) + return latents + + @staticmethod + # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._prepare_latent_image_ids + def _prepare_latent_image_ids(batch_size, height, width, device, dtype): + latent_image_ids = torch.zeros(height, width, 3) + latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None] + latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :] + + latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape + + latent_image_ids = latent_image_ids.reshape( + latent_image_id_height * latent_image_id_width, latent_image_id_channels + ) + + return latent_image_ids.to(device=device, dtype=dtype) + + @staticmethod + def _unpack_latents_no_patch(latents, height, width, vae_scale_factor): + batch_size, num_patches, channels = latents.shape + + height = height // vae_scale_factor + width = width // vae_scale_factor + + latents = latents.view(batch_size, height, width, channels) + latents = latents.permute(0, 3, 1, 2) + + return latents + + @staticmethod + def _pack_latents_no_patch(latents, batch_size, num_channels_latents, height, width): + latents = latents.permute(0, 2, 3, 1) + latents = latents.reshape(batch_size, height * width, num_channels_latents) + return latents + + @staticmethod + # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._pack_latents + def _pack_latents(latents, batch_size, num_channels_latents, height, width): + latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 2, 4, 1, 3, 5) + latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4) + + return latents + + def prepare_latents( + self, + batch_size, + num_channels_latents, + height, + width, + dtype, + device, + generator, + latents=None, + do_patching=False, + ): + height = int(height) // self.vae_scale_factor + width = int(width) // self.vae_scale_factor + + shape = (batch_size, num_channels_latents, height, width) + + if latents is not None: + latent_image_ids = self._prepare_latent_image_ids(batch_size, height, width, device, dtype) + return latents.to(device=device, dtype=dtype), latent_image_ids + + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + if do_patching: + latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width) + latent_image_ids = self._prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype) + else: + latents = self._pack_latents_no_patch(latents, batch_size, num_channels_latents, height, width) + latent_image_ids = self._prepare_latent_image_ids(batch_size, height, width, device, dtype) + + return latents, latent_image_ids + + @staticmethod + def _prepare_attention_mask(attention_mask): + attention_matrix = torch.einsum("bi,bj->bij", attention_mask, attention_mask) + + # convert to 0 - keep, -inf ignore + attention_matrix = torch.where( + attention_matrix == 1, 0.0, -torch.inf + ) # Apply -inf to ignored tokens for nulling softmax score + return attention_matrix + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + image: Optional[PipelineImageInput] = None, + mask: Optional[PipelineMaskInput] = None, + height: int | None = None, + width: int | None = None, + num_inference_steps: int = 30, + timesteps: List[int] = None, + seed: int | None = None, + guidance_scale: float = 5, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + output_type: str = "pil", + return_dict: bool = True, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + max_sequence_length: int = 3000, + do_patching=False, + _auto_resize: bool = True, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + image (`PIL.Image.Image` or `torch.FloatTensor`, *optional*): + The image to guide the image generation. If not defined, the pipeline will generate an image from + scratch. + height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The height in pixels of the generated image. This is set to 1024 by default for the best results. + width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The width in pixels of the generated image. This is set to 1024 by default for the best results. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + seed (`int`, *optional*): + A seed used to make generation deterministic. + timesteps (`List[int]`, *optional*): + Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument + in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is + passed will be used. Must be in descending order. + guidance_scale (`float`, *optional*, defaults to 5.0): + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. + of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting + `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to + the text `prompt`, usually at the expense of lower image quality. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will ge generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead + of a plain tuple. + joint_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int` defaults to 3000): Maximum sequence length to use with the `prompt`. + do_patching (`bool`, *optional*, defaults to `False`): Whether to use patching. + Examples: + Returns: + [`~pipelines.flux.BriaFiboPipelineOutput`] or `tuple`: [`~pipelines.flux.BriaFiboPipelineOutput`] if + `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the + generated images. + """ + + if height is None or width is None: + if image is not None: + image_height, image_width = self.image_processor.get_default_height_width(image) + if _auto_resize: + image_width, image_height = min( + PREFERRED_RESOLUTION[1024 * 1024], + key=lambda size: abs(size[0] / size[1] - image_width / image_height), + ) + width, height = image_width, image_height + else: + raise ValueError("You must provide either an image or both height and width.") + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + seed=seed, + image=image, + mask=mask, + prompt=prompt, + height=height, + width=width, + prompt_embeds=prompt_embeds, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + max_sequence_length=max_sequence_length, + ) + + if mask is not None and image is not None: + image = paste_mask_on_image(mask, image) + + self._guidance_scale = guidance_scale + self._joint_attention_kwargs = joint_attention_kwargs + self._interrupt = False + + # 2. Define call parameters + + if prompt is not None and is_valid_edit_json(prompt): + prompt = json.dumps(prompt) + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + if generator is None and seed is not None: + generator = torch.Generator(device=device).manual_seed(seed) + lora_scale = ( + self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None + ) + + ( + prompt_embeds, + negative_prompt_embeds, + text_ids, + prompt_attention_mask, + negative_prompt_attention_mask, + prompt_layers, + negative_prompt_layers, + ) = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + guidance_scale=guidance_scale, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + device=device, + max_sequence_length=max_sequence_length, + num_images_per_prompt=num_images_per_prompt, + lora_scale=lora_scale, + ) + prompt_batch_size = prompt_embeds.shape[0] + + if guidance_scale > 1: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + prompt_layers = [ + torch.cat([negative_prompt_layers[i], prompt_layers[i]], dim=0) for i in range(len(prompt_layers)) + ] + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) + + total_num_layers_transformer = len(self.transformer.transformer_blocks) + len( + self.transformer.single_transformer_blocks + ) + if len(prompt_layers) >= total_num_layers_transformer: + # remove first layers + prompt_layers = prompt_layers[len(prompt_layers) - total_num_layers_transformer :] + else: + # duplicate last layer + prompt_layers = prompt_layers + [prompt_layers[-1]] * (total_num_layers_transformer - len(prompt_layers)) + + # Preprocess image + if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels): + image = self.image_processor.resize(image, height, width) + image = self.image_processor.preprocess(image, height, width) + + # 5. Prepare latent variables + num_channels_latents = self.transformer.config.in_channels + if do_patching: + num_channels_latents = int(num_channels_latents / 4) + + latents, latent_image_ids = self.prepare_latents( + prompt_batch_size, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + do_patching, + ) + + if image is not None: + image_latents, image_ids = self.prepare_image_latents( + image=image, + batch_size=batch_size * num_images_per_prompt, + num_channels_latents=num_channels_latents, + height=height, + width=width, + dtype=prompt_embeds.dtype, + device=device, + generator=generator, + ) + latent_image_ids = torch.cat([latent_image_ids, image_ids], dim=0) # dim 0 is sequence dimension + else: + image_latents = None + + latent_attention_mask = torch.ones( + [latents.shape[0], latents.shape[1]], dtype=latents.dtype, device=latents.device + ) + if guidance_scale > 1: + latent_attention_mask = latent_attention_mask.repeat(2, 1) + + if image_latents is None: + attention_mask = torch.cat([prompt_attention_mask, latent_attention_mask], dim=1) + else: + image_latent_attention_mask = torch.ones( + [image_latents.shape[0], image_latents.shape[1]], + dtype=image_latents.dtype, + device=image_latents.device, + ) + if guidance_scale > 1: + image_latent_attention_mask = image_latent_attention_mask.repeat(2, 1) + attention_mask = torch.cat( + [prompt_attention_mask, latent_attention_mask, image_latent_attention_mask], dim=1 + ) + + attention_mask = self.create_attention_matrix(attention_mask) # batch, seq => batch, seq, seq + attention_mask = attention_mask.unsqueeze(dim=1).to(dtype=self.transformer.dtype) # for head broadcasting + + if self._joint_attention_kwargs is None: + self._joint_attention_kwargs = {} + self._joint_attention_kwargs["attention_mask"] = attention_mask + + # Adapt scheduler to dynamic shifting (resolution dependent) + + if do_patching: + seq_len = (height // (self.vae_scale_factor * 2)) * (width // (self.vae_scale_factor * 2)) + else: + seq_len = (height // self.vae_scale_factor) * (width // self.vae_scale_factor) + + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) + + mu = calculate_shift( + seq_len, + self.scheduler.config.base_image_seq_len, + self.scheduler.config.max_image_seq_len, + self.scheduler.config.base_shift, + self.scheduler.config.max_shift, + ) + + # Init sigmas and timesteps according to shift size + # This changes the scheduler in-place according to the dynamic scheduling + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps=num_inference_steps, + device=device, + timesteps=None, + sigmas=sigmas, + mu=mu, + ) + + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + # Support old different diffusers versions + if len(latent_image_ids.shape) == 3: + latent_image_ids = latent_image_ids[0] + + if len(text_ids.shape) == 3: + text_ids = text_ids[0] + + # 6. Denoising loop + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + latent_model_input = latents + + if image_latents is not None: + latent_model_input = torch.cat([latent_model_input, image_latents], dim=1) + + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latent_model_input] * 2) if guidance_scale > 1 else latent_model_input + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latent_model_input.shape[0]).to( + device=latent_model_input.device, dtype=latent_model_input.dtype + ) + + # This is predicts "v" from flow-matching or eps from diffusion + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=prompt_embeds, + text_encoder_layers=prompt_layers, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + txt_ids=text_ids, + img_ids=latent_image_ids, + )[0] + + # perform guidance + if guidance_scale > 1: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents_dtype = latents.dtype + latents = self.scheduler.step(noise_pred[:, : latents.shape[1], ...], t, latents, return_dict=False)[0] + + if latents.dtype != latents_dtype: + if torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 + latents = latents.to(latents_dtype) + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + if output_type == "latent": + image = latents + + else: + if do_patching: + latents = self._unpack_latents(latents, height, width, self.vae_scale_factor) + else: + latents = self._unpack_latents_no_patch(latents, height, width, self.vae_scale_factor) + + latents = latents.unsqueeze(dim=2) + latents_device = latents[0].device + latents_dtype = latents[0].dtype + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents_device, latents_dtype) + ) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( + latents_device, latents_dtype + ) + latents_scaled = [latent / latents_std + latents_mean for latent in latents] + latents_scaled = torch.cat(latents_scaled, dim=0) + image = [] + for scaled_latent in latents_scaled: + curr_image = self.vae.decode(scaled_latent.unsqueeze(0), return_dict=False)[0] + curr_image = self.image_processor.postprocess(curr_image.squeeze(dim=2), output_type=output_type) + image.append(curr_image) + if len(image) == 1: + image = image[0] + else: + image = np.stack(image, axis=0) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return BriaFiboPipelineOutput(images=image) + + def prepare_image_latents( + self, + image: torch.Tensor, + batch_size: int, + num_channels_latents: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: torch.Generator | list[torch.Generator] | None = None, + ): + image = image.to(device=device, dtype=dtype) + + height = int(height) // self.vae_scale_factor + width = int(width) // self.vae_scale_factor + + # scaling + latents_mean = ( + torch.tensor(self.vae.config.latents_mean).view(1, self.vae.config.z_dim, 1, 1, 1).to(device, dtype) + ) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( + device, dtype + ) + + image_latents_cthw = self.vae.encode(image.unsqueeze(2)).latent_dist.mean + latents_scaled = [(latent - latents_mean) * latents_std for latent in image_latents_cthw] + image_latents_cthw = torch.concat(latents_scaled, dim=0) + image_latents_bchw = image_latents_cthw[:, :, 0, :, :] + + image_latent_height, image_latent_width = image_latents_bchw.shape[2:] + image_latents_bsd = self._pack_latents_no_patch( + latents=image_latents_bchw, + batch_size=batch_size, + num_channels_latents=num_channels_latents, + height=image_latent_height, + width=image_latent_width, + ) + # breakpoint() + image_ids = self._prepare_latent_image_ids( + batch_size=batch_size, height=image_latent_height, width=image_latent_width, device=device, dtype=dtype + ) + # image ids are the same as latent ids with the first dimension set to 1 instead of 0 + image_ids[..., 0] = 1 + return image_latents_bsd, image_ids + + def check_inputs( + self, + prompt, + seed, + image, + mask, + height, + width, + negative_prompt=None, + prompt_embeds=None, + negative_prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + max_sequence_length=None, + ): + if seed is not None and not isinstance(seed, int): + raise ValueError("Seed must be an integer") + if image is not None and not isinstance(image, (torch.Tensor, Image.Image, list)): + raise ValueError("Image must be a valid image") + if image is None and mask is not None: + raise ValueError("If mask is provided, image must also be provided") + + if mask is not None and not is_valid_mask(mask): + raise ValueError("Mask must be a valid mask") + + if mask is not None and image is not None and not (get_mask_size(mask) == get_image_size(image)): + raise ValueError("Mask and image must have the same size") + + if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0: + logger.warning( + f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and {width}. Dimensions will be resized accordingly" + ) + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and not is_valid_edit_json(prompt): + raise ValueError(f"`prompt` has to be a valid JSON string or dict but is {type(prompt)}") + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}." + ) + + if max_sequence_length is not None and max_sequence_length > 3000: + raise ValueError(f"`max_sequence_length` cannot be greater than 3000 but is {max_sequence_length}") + + def create_attention_matrix(self, attention_mask): + attention_matrix = torch.einsum("bi,bj->bij", attention_mask, attention_mask) + + # convert to 0 - keep, -inf ignore + attention_matrix = torch.where( + attention_matrix == 1, 0.0, -torch.inf + ) # Apply -inf to ignored tokens for nulling softmax score + return attention_matrix diff --git a/src/diffusers/pipelines/bria_fibo/pipeline_output.py b/src/diffusers/pipelines/bria_fibo/pipeline_output.py index f459185a2c7c..0c131db29d9f 100644 --- a/src/diffusers/pipelines/bria_fibo/pipeline_output.py +++ b/src/diffusers/pipelines/bria_fibo/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -13,9 +12,9 @@ class BriaFiboPipelineOutput(BaseOutput): Output class for BriaFibo pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) + images (`list[PIL.Image.Image]` or `np.ndarray`) List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image, np.ndarray] diff --git a/src/diffusers/pipelines/chroma/__init__.py b/src/diffusers/pipelines/chroma/__init__.py index d9238b735c41..6ecf024c5a97 100644 --- a/src/diffusers/pipelines/chroma/__init__.py +++ b/src/diffusers/pipelines/chroma/__init__.py @@ -23,7 +23,9 @@ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) else: _import_structure["pipeline_chroma"] = ["ChromaPipeline"] + _import_structure["pipeline_chroma_radiance"] = ["ChromaRadiancePipeline"] _import_structure["pipeline_chroma_img2img"] = ["ChromaImg2ImgPipeline"] + _import_structure["pipeline_chroma_inpainting"] = ["ChromaInpaintPipeline"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: if not (is_transformers_available() and is_torch_available()): @@ -33,6 +35,8 @@ else: from .pipeline_chroma import ChromaPipeline from .pipeline_chroma_img2img import ChromaImg2ImgPipeline + from .pipeline_chroma_radiance import ChromaRadiancePipeline + from .pipeline_chroma_inpainting import ChromaInpaintPipeline else: import sys diff --git a/src/diffusers/pipelines/chroma/pipeline_chroma.py b/src/diffusers/pipelines/chroma/pipeline_chroma.py index ed6c2c2105b6..bc782107022d 100644 --- a/src/diffusers/pipelines/chroma/pipeline_chroma.py +++ b/src/diffusers/pipelines/chroma/pipeline_chroma.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -91,10 +91,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -109,15 +109,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -208,11 +208,11 @@ def __init__( def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -264,24 +264,24 @@ def _get_t5_prompt_embeds( def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Union[str, List[str]] = None, - device: Optional[torch.device] = None, + prompt: str | list[str], + negative_prompt: str | list[str] = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, do_classifier_free_guidance: bool = True, max_sequence_length: int = 512, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). device: (`torch.device`): @@ -642,39 +642,39 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 35, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 5.0, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - negative_ip_adapter_image: Optional[PipelineImageInput] = None, - negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + negative_ip_adapter_image: PipelineImageInput | None = None, + negative_ip_adapter_image_embeds: list[torch.Tensor] | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is not greater than `1`). @@ -685,7 +685,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -697,7 +697,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -708,13 +708,13 @@ def __call__( Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, text embeddings will be generated from `prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not provided, embeddings are computed from the `ip_adapter_image` input argument. negative_ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + negative_ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not provided, embeddings are computed from the `ip_adapter_image` input argument. @@ -744,7 +744,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py b/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py index 470c746e4146..e1f6e2f8d8af 100644 --- a/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py +++ b/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -88,7 +88,7 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -103,10 +103,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -121,15 +121,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -222,11 +222,11 @@ def __init__( def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -291,24 +291,24 @@ def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Union[str, List[str]] = None, - device: Optional[torch.device] = None, + prompt: str | list[str], + negative_prompt: str | list[str] = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, do_classifier_free_guidance: bool = True, max_sequence_length: int = 512, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). device: (`torch.device`): @@ -701,41 +701,41 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 35, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 5.0, strength: float = 0.9, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - negative_ip_adapter_image: Optional[PipelineImageInput] = None, - negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.tensor] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + negative_ip_adapter_image: PipelineImageInput | None = None, + negative_ip_adapter_image_embeds: list[torch.Tensor] | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is not greater than `1`). @@ -746,7 +746,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 35): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -764,7 +764,7 @@ def __call__( A value of 1, therefore, essentially ignores image. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -775,13 +775,13 @@ def __call__( Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, text embeddings will be generated from `prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not provided, embeddings are computed from the `ip_adapter_image` input argument. negative_ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + negative_ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not provided, embeddings are computed from the `ip_adapter_image` input argument. @@ -811,7 +811,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/chroma/pipeline_chroma_inpainting.py b/src/diffusers/pipelines/chroma/pipeline_chroma_inpainting.py new file mode 100644 index 000000000000..52c2f7e51cf2 --- /dev/null +++ b/src/diffusers/pipelines/chroma/pipeline_chroma_inpainting.py @@ -0,0 +1,1184 @@ +""" +ChromaInpaintPipeline implements a text-guided image inpainting pipeline for the lodestones/Chroma1-HD model, based on +the ChromaPipeline from Hugging Face Diffusers:contentReference[oaicite:0]{index=0} and the Stable Diffusion inpainting +approach:contentReference[oaicite:1]{index=1}. +""" + +# Copyright 2025 Black Forest Labs and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import inspect +from typing import Any, Callable, Dict, List, Optional, Union + +import numpy as np +import PIL.Image +import torch +from transformers import ( + CLIPImageProcessor, + CLIPVisionModelWithProjection, + T5EncoderModel, + T5TokenizerFast, +) + +from ...image_processor import PipelineImageInput, VaeImageProcessor +from ...loaders import FluxIPAdapterMixin, FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin +from ...models.autoencoders import AutoencoderKL +from ...models.transformers import ChromaTransformer2DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import ( + USE_PEFT_BACKEND, + is_torch_xla_available, + logging, + replace_example_docstring, + scale_lora_layers, + unscale_lora_layers, +) +from ...utils.torch_utils import randn_tensor +from ..chroma.pipeline_output import ChromaPipelineOutput +from ..pipeline_utils import DiffusionPipeline + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import ChromaInpaintPipeline + >>> from diffusers.utils import load_image + + >>> pipe = ChromaInpaintPipeline.from_pretrained("lodestones/Chroma1-HD", torch_dtype=torch.bfloat16) + >>> pipe.to("cuda") + >>> prompt = "Face of a yellow cat, high resolution, sitting on a park bench" + >>> img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png" + >>> mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png" + >>> source = load_image(img_url) + >>> mask = load_image(mask_url) + >>> image = pipe(prompt=prompt, image=source, mask_image=mask).images[0] + >>> image.save("chroma_inpainting.png") + ``` +""" + + +# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents +def retrieve_latents( + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" +): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`list[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`list[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class ChromaInpaintPipeline( + DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin, FluxIPAdapterMixin +): + r""" + The Flux pipeline for image inpainting. + + Reference: https://blackforestlabs.ai/announcing-black-forest-labs/ + + Args: + transformer ([`ChromaTransformer2DModel`]): + Conditional Transformer (MMDiT) architecture to denoise the encoded image latents. + scheduler ([`DDIMScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + vae ([`AutoencoderKL`]): + Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. + text_encoder ([`CLIPTextModel`]): + [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically + the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant. + text_encoder_2 ([`T5EncoderModel`]): + [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically + the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant. + tokenizer (`CLIPTokenizer`): + Tokenizer of class + [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer). + tokenizer_2 (`T5TokenizerFast`): + Second Tokenizer of class + [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast). + """ + + model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->transformer->vae" + _optional_components = ["image_encoder", "feature_extractor"] + _callback_tensor_inputs = ["latents", "prompt_embeds"] + + def __init__( + self, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKL, + text_encoder: T5EncoderModel, + tokenizer: T5TokenizerFast, + transformer: ChromaTransformer2DModel, + image_encoder: CLIPVisionModelWithProjection = None, + feature_extractor: CLIPImageProcessor = None, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + scheduler=scheduler, + image_encoder=image_encoder, + feature_extractor=feature_extractor, + ) + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 + self.latent_channels = self.vae.config.latent_channels if getattr(self, "vae", None) else 16 + + # Flux latents are turned into 2x2 patches and packed. This means the latent width and height has to be divisible + # by the patch size. So the vae scale factor is multiplied by the patch size to account for this + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) + self.default_sample_size = 128 + + self.mask_processor = VaeImageProcessor( + vae_scale_factor=self.vae_scale_factor * 2, + vae_latent_channels=self.latent_channels, + do_normalize=False, + do_binarize=True, + do_convert_grayscale=True, + ) + + def _get_t5_prompt_embeds( + self, + prompt: Union[str, List[str], None] = None, + num_images_per_prompt: int = 1, + max_sequence_length: int = 512, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + if isinstance(self, TextualInversionLoaderMixin): + prompt = self.maybe_convert_prompt(prompt, self.tokenizer) + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + return_length=False, + return_overflowing_tokens=False, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + tokenizer_mask = text_inputs.attention_mask + + tokenizer_mask = tokenizer_mask.to(device) + + prompt_embeds = self.text_encoder( + text_input_ids.to(device), + output_hidden_states=False, + attention_mask=tokenizer_mask, + )[0] + + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + + seq_lengths = tokenizer_mask.sum(dim=1) + mask_indices = torch.arange(tokenizer_mask.size(1), device=device).unsqueeze(0).expand(batch_size, -1) + attention_mask = (mask_indices <= seq_lengths.unsqueeze(1)).to(dtype=dtype, device=device) + + _, seq_len, _ = prompt_embeds.shape + + # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + attention_mask = attention_mask.repeat(1, num_images_per_prompt) + attention_mask = attention_mask.view(batch_size * num_images_per_prompt, seq_len) + + return prompt_embeds, attention_mask + + def encode_prompt( + self, + prompt: Union[str, List[str]], + negative_prompt: Union[str, List[str], None] = None, + device: torch.device | None = None, + num_images_per_prompt: int = 1, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + do_classifier_free_guidance: bool = True, + max_sequence_length: int = 256, + lora_scale: bool | None = None, + ): + r""" + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + negative_prompt (`str` or `List[str]`, *optional*): + The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` + instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). + device: (`torch.device`): + torch device + num_images_per_prompt (`int`): + number of images that should be generated per prompt + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + lora_scale (`float`, *optional*): + A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded. + """ + device = device or self._execution_device + + # set lora scale so that monkey patched LoRA + # function of text encoder can correctly access it + if lora_scale is not None and isinstance(self, FluxLoraLoaderMixin): + self._lora_scale = lora_scale + + # dynamically adjust the LoRA scale + if self.text_encoder is not None and USE_PEFT_BACKEND: + scale_lora_layers(self.text_encoder, lora_scale) + + prompt = [prompt] if isinstance(prompt, str) else prompt + + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds, prompt_attention_mask = self._get_t5_prompt_embeds( + prompt=prompt, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + ) + + dtype = self.text_encoder.dtype if self.text_encoder is not None else self.transformer.dtype + text_ids = torch.zeros(prompt_embeds.shape[1], 3, device=device, dtype=dtype) + negative_text_ids = None + + if do_classifier_free_guidance: + if negative_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt = ( + batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + ) + + if prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + + negative_prompt_embeds, negative_prompt_attention_mask = self._get_t5_prompt_embeds( + prompt=negative_prompt, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + ) + + negative_text_ids = torch.zeros(negative_prompt_embeds.shape[1], 3, device=device, dtype=dtype) + + if self.text_encoder is not None: + if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND: + # Retrieve the original scale by scaling back the LoRA layers + unscale_lora_layers(self.text_encoder, lora_scale) + + return ( + prompt_embeds, + text_ids, + prompt_attention_mask, + negative_prompt_embeds, + negative_text_ids, + negative_prompt_attention_mask, + ) + + # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_image + def encode_image(self, image, device, num_images_per_prompt): + dtype = next(self.image_encoder.parameters()).dtype + + if not isinstance(image, torch.Tensor): + image = self.feature_extractor(image, return_tensors="pt").pixel_values + + image = image.to(device=device, dtype=dtype) + image_embeds = self.image_encoder(image).image_embeds + image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) + return image_embeds + + # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.prepare_ip_adapter_image_embeds + def prepare_ip_adapter_image_embeds( + self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt + ): + image_embeds = [] + if ip_adapter_image_embeds is None: + if not isinstance(ip_adapter_image, list): + ip_adapter_image = [ip_adapter_image] + + if len(ip_adapter_image) != self.transformer.encoder_hid_proj.num_ip_adapters: + raise ValueError( + f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {self.transformer.encoder_hid_proj.num_ip_adapters} IP Adapters." + ) + + for single_ip_adapter_image in ip_adapter_image: + single_image_embeds = self.encode_image(single_ip_adapter_image, device, 1) + image_embeds.append(single_image_embeds[None, :]) + else: + if not isinstance(ip_adapter_image_embeds, list): + ip_adapter_image_embeds = [ip_adapter_image_embeds] + + if len(ip_adapter_image_embeds) != self.transformer.encoder_hid_proj.num_ip_adapters: + raise ValueError( + f"`ip_adapter_image_embeds` must have same length as the number of IP Adapters. Got {len(ip_adapter_image_embeds)} image embeds and {self.transformer.encoder_hid_proj.num_ip_adapters} IP Adapters." + ) + + for single_image_embeds in ip_adapter_image_embeds: + image_embeds.append(single_image_embeds) + + ip_adapter_image_embeds = [] + for single_image_embeds in image_embeds: + single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) + single_image_embeds = single_image_embeds.to(device=device) + ip_adapter_image_embeds.append(single_image_embeds) + + return ip_adapter_image_embeds + + # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3_inpaint.StableDiffusion3InpaintPipeline._encode_vae_image + def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): + if isinstance(generator, list): + image_latents = [ + retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) + for i in range(image.shape[0]) + ] + image_latents = torch.cat(image_latents, dim=0) + else: + image_latents = retrieve_latents(self.vae.encode(image), generator=generator) + + image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor + + return image_latents + + # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3_img2img.StableDiffusion3Img2ImgPipeline.get_timesteps + def get_timesteps(self, num_inference_steps, strength, device): + # get the original timestep using init_timestep + init_timestep = min(num_inference_steps * strength, num_inference_steps) + + t_start = int(max(num_inference_steps - init_timestep, 0)) + timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :] + if hasattr(self.scheduler, "set_begin_index"): + self.scheduler.set_begin_index(t_start * self.scheduler.order) + + return timesteps, num_inference_steps - t_start + + def check_inputs( + self, + prompt, + image, + mask_image, + strength, + height, + width, + output_type, + negative_prompt=None, + prompt_embeds=None, + negative_prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + padding_mask_crop=None, + max_sequence_length=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + ): + if strength < 0 or strength > 1: + raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}") + + if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0: + logger.warning( + f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and {width}. Dimensions will be resized accordingly" + ) + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}." + ) + + if prompt_embeds is not None and prompt_attention_mask is None: + raise ValueError("Cannot provide `prompt_embeds` without also providing `prompt_attention_mask") + + if negative_prompt_embeds is not None and negative_prompt_attention_mask is None: + raise ValueError( + "Cannot provide `negative_prompt_embeds` without also providing `negative_prompt_attention_mask" + ) + + if padding_mask_crop is not None: + if not isinstance(image, PIL.Image.Image): + raise ValueError( + f"The image should be a PIL image when inpainting mask crop, but is of type {type(image)}." + ) + if not isinstance(mask_image, PIL.Image.Image): + raise ValueError( + f"The mask image should be a PIL image when inpainting mask crop, but is of type" + f" {type(mask_image)}." + ) + if output_type != "pil": + raise ValueError(f"The output type should be PIL when inpainting mask crop, but is {output_type}.") + + if max_sequence_length is not None and max_sequence_length > 512: + raise ValueError(f"`max_sequence_length` cannot be greater than 512 but is {max_sequence_length}") + + @staticmethod + # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._prepare_latent_image_ids + def _prepare_latent_image_ids(batch_size, height, width, device, dtype): + latent_image_ids = torch.zeros(height, width, 3) + latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None] + latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :] + + latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape + + latent_image_ids = latent_image_ids.reshape( + latent_image_id_height * latent_image_id_width, latent_image_id_channels + ) + + return latent_image_ids.to(device=device, dtype=dtype) + + @staticmethod + # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._pack_latents + def _pack_latents(latents, batch_size, num_channels_latents, height, width): + latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 2, 4, 1, 3, 5) + latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4) + + return latents + + @staticmethod + # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._unpack_latents + def _unpack_latents(latents, height, width, vae_scale_factor): + batch_size, num_patches, channels = latents.shape + + # VAE applies 8x compression on images but we must also account for packing which requires + # latent height and width to be divisible by 2. + height = 2 * (int(height) // (vae_scale_factor * 2)) + width = 2 * (int(width) // (vae_scale_factor * 2)) + + latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2) + latents = latents.permute(0, 3, 1, 4, 2, 5) + + latents = latents.reshape(batch_size, channels // (2 * 2), height, width) + + return latents + + def prepare_latents( + self, + image, + timestep, + batch_size, + num_channels_latents, + height, + width, + dtype, + device, + generator, + latents=None, + ): + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + # VAE applies 8x compression on images but we must also account for packing which requires + # latent height and width to be divisible by 2. + height = 2 * (int(height) // (self.vae_scale_factor * 2)) + width = 2 * (int(width) // (self.vae_scale_factor * 2)) + shape = (batch_size, num_channels_latents, height, width) + latent_image_ids = self._prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype) + + image = image.to(device=device, dtype=dtype) + if image.shape[1] != self.latent_channels: + image_latents = self._encode_vae_image(image=image, generator=generator) + else: + image_latents = image + + if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0: + # expand init_latents for batch_size + additional_image_per_prompt = batch_size // image_latents.shape[0] + image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0) + elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0: + raise ValueError( + f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts." + ) + else: + image_latents = torch.cat([image_latents], dim=0) + + if latents is None: + noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = self.scheduler.scale_noise(image_latents, timestep, noise) + else: + noise = latents.to(device) + latents = noise + + noise = self._pack_latents(noise, batch_size, num_channels_latents, height, width) + image_latents = self._pack_latents(image_latents, batch_size, num_channels_latents, height, width) + latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width) + return latents, noise, image_latents, latent_image_ids + + def prepare_mask_latents( + self, + mask, + masked_image, + batch_size, + num_channels_latents, + num_images_per_prompt, + height, + width, + dtype, + device, + generator, + ): + # VAE applies 8x compression on images but we must also account for packing which requires + # latent height and width to be divisible by 2. + height = 2 * (int(height) // (self.vae_scale_factor * 2)) + width = 2 * (int(width) // (self.vae_scale_factor * 2)) + # resize the mask to latents shape as we concatenate the mask to the latents + # we do that before converting to dtype to avoid breaking in case we're using cpu_offload + # and half precision + mask = torch.nn.functional.interpolate(mask, size=(height, width)) + mask = mask.to(device=device, dtype=dtype) + + batch_size = batch_size * num_images_per_prompt + + masked_image = masked_image.to(device=device, dtype=dtype) + + if masked_image.shape[1] == 16: + masked_image_latents = masked_image + else: + masked_image_latents = retrieve_latents(self.vae.encode(masked_image), generator=generator) + + masked_image_latents = ( + masked_image_latents - self.vae.config.shift_factor + ) * self.vae.config.scaling_factor + + # duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method + if mask.shape[0] < batch_size: + if not batch_size % mask.shape[0] == 0: + raise ValueError( + "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to" + f" a total batch size of {batch_size}, but {mask.shape[0]} masks were passed. Make sure the number" + " of masks that you pass is divisible by the total requested batch size." + ) + mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1) + if masked_image_latents.shape[0] < batch_size: + if not batch_size % masked_image_latents.shape[0] == 0: + raise ValueError( + "The passed images and the required batch size don't match. Images are supposed to be duplicated" + f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed." + " Make sure the number of images that you pass is divisible by the total requested batch size." + ) + masked_image_latents = masked_image_latents.repeat(batch_size // masked_image_latents.shape[0], 1, 1, 1) + + # aligning device to prevent device errors when concating it with the latent model input + masked_image_latents = masked_image_latents.to(device=device, dtype=dtype) + masked_image_latents = self._pack_latents( + masked_image_latents, + batch_size, + num_channels_latents, + height, + width, + ) + mask = self._pack_latents( + mask.repeat(1, num_channels_latents, 1, 1), + batch_size, + num_channels_latents, + height, + width, + ) + + return mask, masked_image_latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def joint_attention_kwargs(self): + return self._joint_attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def interrupt(self): + return self._interrupt + + def _prepare_attention_mask( + self, + batch_size, + sequence_length, + dtype, + attention_mask=None, + ): + if attention_mask is None: + return attention_mask + + # Extend the prompt attention mask to account for image tokens in the final sequence + attention_mask = torch.cat( + [attention_mask, torch.ones(batch_size, sequence_length, device=attention_mask.device, dtype=torch.bool)], + dim=1, + ) + attention_mask = attention_mask.to(dtype) + + return attention_mask + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1 + + @replace_example_docstring(EXAMPLE_DOC_STRING) + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]] = None, + negative_prompt: Union[str, List[str]] = None, + true_cfg_scale: float = 1.0, + image: PipelineImageInput = None, + mask_image: PipelineImageInput = None, + masked_image_latents: PipelineImageInput = None, + height: int | None = None, + width: int | None = None, + padding_mask_crop: int | None = None, + strength: float = 0.6, + num_inference_steps: int = 28, + sigmas: list[float] | None = None, + guidance_scale: float = 7.0, + num_images_per_prompt: Optional[int] = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + ip_adapter_image: Optional[PipelineImageInput] = None, + ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, + negative_ip_adapter_image: Optional[PipelineImageInput] = None, + negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + output_type: str = "pil", + return_dict: bool = True, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + max_sequence_length: int = 256, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + not greater than `1`). + height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The height in pixels of the generated image. This is set to 1024 by default for the best results. + width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The width in pixels of the generated image. This is set to 1024 by default for the best results. + num_inference_steps (`int`, *optional*, defaults to 35): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + guidance_scale (`float`, *optional*, defaults to 3.5): + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. + of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting + `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to + the text `prompt`, usually at the expense of lower image quality. + strength (`float, *optional*, defaults to 0.9): + Conceptually, indicates how much to transform the reference image. Must be between 0 and 1. image will + be used as a starting point, adding more noise to it the larger the strength. The number of denoising + steps depends on the amount of noise initially added. When strength is 1, added noise will be maximum + and the denoising process will run for the full number of iterations specified in num_inference_steps. + A value of 1, therefore, essentially ignores image. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. + ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of + IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not + provided, embeddings are computed from the `ip_adapter_image` input argument. + negative_ip_adapter_image: + (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. + negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of + IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not + provided, embeddings are computed from the `ip_adapter_image` input argument. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + prompt_attention_mask (torch.Tensor, *optional*): + Attention mask for the prompt embeddings. Used to mask out padding tokens in the prompt sequence. + Chroma requires a single padding token remain unmasked. Please refer to + https://huggingface.co/lodestones/Chroma#tldr-masking-t5-padding-tokens-enhanced-fidelity-and-increased-stability-during-training + negative_prompt_attention_mask (torch.Tensor, *optional*): + Attention mask for the negative prompt embeddings. Used to mask out padding tokens in the negative + prompt sequence. Chroma requires a single padding token remain unmasked. PLease refer to + https://huggingface.co/lodestones/Chroma#tldr-masking-t5-padding-tokens-enhanced-fidelity-and-increased-stability-during-training + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.chroma.ChromaPipelineOutput`] instead of a plain tuple. + joint_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`. + + Examples: + + Returns: + [`~pipelines.chroma.ChromaPipelineOutput`] or `tuple`: [`~pipelines.chroma.ChromaPipelineOutput`] if + `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the + generated images. + """ + + height = height or self.default_sample_size * self.vae_scale_factor + width = width or self.default_sample_size * self.vae_scale_factor + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt=prompt, + height=height, + width=width, + output_type=output_type, + strength=strength, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + max_sequence_length=max_sequence_length, + image=image, + mask_image=mask_image, + padding_mask_crop=padding_mask_crop, + ) + + self._guidance_scale = guidance_scale + self._joint_attention_kwargs = joint_attention_kwargs + self._current_timestep = None + self._interrupt = False + + # 2. Preprocess mask and image + if padding_mask_crop is not None: + crops_coords = self.mask_processor.get_crop_region(mask_image, width, height, pad=padding_mask_crop) + resize_mode = "fill" + else: + crops_coords = None + resize_mode = "default" + + original_image = image + init_image = self.image_processor.preprocess( + image, height=height, width=width, crops_coords=crops_coords, resize_mode=resize_mode + ) + init_image = init_image.to(dtype=torch.float32) + + # 3. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + lora_scale = ( + self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None + ) + + ( + prompt_embeds, + text_ids, + prompt_attention_mask, + negative_prompt_embeds, + negative_text_ids, + negative_prompt_attention_mask, + ) = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + do_classifier_free_guidance=self.do_classifier_free_guidance, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + lora_scale=lora_scale, + ) + + # 4. Prepare timesteps + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas + image_seq_len = (int(height) // self.vae_scale_factor // 2) * (int(width) // self.vae_scale_factor // 2) + mu = calculate_shift( + image_seq_len, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.5), + self.scheduler.config.get("max_shift", 1.15), + ) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + sigmas=sigmas, + mu=mu, + ) + timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) + + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + if num_inference_steps < 1: + raise ValueError( + f"After adjusting the num_inference_steps by strength parameter: {strength}, the number of pipeline" + f"steps is {num_inference_steps} which is < 1 and not appropriate for this pipeline." + ) + latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) + + # 5. Prepare latent variables + num_channels_latents = self.transformer.config.in_channels // 4 + num_channels_transformer = self.transformer.config.in_channels + + latents, noise, image_latents, latent_image_ids = self.prepare_latents( + init_image, + latent_timestep, + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + ) + + mask_condition = self.mask_processor.preprocess( + mask_image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords + ) + + if masked_image_latents is None: + masked_image = init_image * (mask_condition < 0.5) + else: + masked_image = masked_image_latents + + mask, masked_image_latents = self.prepare_mask_latents( + mask_condition, + masked_image, + batch_size, + num_channels_latents, + num_images_per_prompt, + height, + width, + prompt_embeds.dtype, + device, + generator, + ) + + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + # handle guidance + if self.transformer.config.guidance_embeds: + guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32) + guidance = guidance.expand(latents.shape[0]) + else: + guidance = None + + if (ip_adapter_image is not None or ip_adapter_image_embeds is not None) and ( + negative_ip_adapter_image is None and negative_ip_adapter_image_embeds is None + ): + negative_ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8) + elif (ip_adapter_image is None and ip_adapter_image_embeds is None) and ( + negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None + ): + ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8) + + if self.joint_attention_kwargs is None: + self._joint_attention_kwargs = {} + + image_embeds = None + negative_image_embeds = None + if ip_adapter_image is not None or ip_adapter_image_embeds is not None: + image_embeds = self.prepare_ip_adapter_image_embeds( + ip_adapter_image, + ip_adapter_image_embeds, + device, + batch_size * num_images_per_prompt, + ) + if negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None: + negative_image_embeds = self.prepare_ip_adapter_image_embeds( + negative_ip_adapter_image, + negative_ip_adapter_image_embeds, + device, + batch_size * num_images_per_prompt, + ) + + attention_mask = self._prepare_attention_mask( + batch_size=latents.shape[0], + sequence_length=image_seq_len, + dtype=latents.dtype, + attention_mask=prompt_attention_mask, + ) + negative_attention_mask = self._prepare_attention_mask( + batch_size=latents.shape[0], + sequence_length=image_seq_len, + dtype=latents.dtype, + attention_mask=negative_prompt_attention_mask, + ) + + # 6. Denoising loop + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + self._current_timestep = t + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latents.shape[0]).to(latents.dtype) + + if image_embeds is not None: + self._joint_attention_kwargs["ip_adapter_image_embeds"] = image_embeds + + noise_pred = self.transformer( + hidden_states=latents, + timestep=timestep / 1000, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, + img_ids=latent_image_ids, + attention_mask=attention_mask, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + )[0] + + if self.do_classifier_free_guidance: + if negative_image_embeds is not None: + self._joint_attention_kwargs["ip_adapter_image_embeds"] = negative_image_embeds + + noise_pred_uncond = self.transformer( + hidden_states=latents, + timestep=timestep / 1000, + encoder_hidden_states=negative_prompt_embeds, + txt_ids=negative_text_ids, + img_ids=latent_image_ids, + attention_mask=negative_attention_mask, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + )[0] + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents_dtype = latents.dtype + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + # for 64 channel transformer only. + init_latents_proper = image_latents + init_mask = mask + + if i < len(timesteps) - 1: + noise_timestep = timesteps[i + 1] + init_latents_proper = self.scheduler.scale_noise( + init_latents_proper, torch.tensor([noise_timestep]), noise + ) + + latents = (1 - init_mask) * init_latents_proper + init_mask * latents + + if latents.dtype != latents_dtype: + if torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 + latents = latents.to(latents_dtype) + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + self._current_timestep = None + + if output_type == "latent": + image = latents + else: + latents = self._unpack_latents(latents, height, width, self.vae_scale_factor) + latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor + image = self.vae.decode(latents, return_dict=False)[0] + image = self.image_processor.postprocess(image, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return ChromaPipelineOutput(images=image) diff --git a/src/diffusers/pipelines/chroma/pipeline_chroma_radiance.py b/src/diffusers/pipelines/chroma/pipeline_chroma_radiance.py new file mode 100644 index 000000000000..dba6d9619681 --- /dev/null +++ b/src/diffusers/pipelines/chroma/pipeline_chroma_radiance.py @@ -0,0 +1,883 @@ +# Copyright 2025 Black Forest Labs and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +from typing import Any, Callable, Dict, List, Optional, Union + +import numpy as np +import torch +from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection, T5EncoderModel, T5TokenizerFast + +from ...image_processor import PipelineImageInput, VaeImageProcessor +from ...loaders import FluxIPAdapterMixin, FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin +from ...models import ChromaRadianceTransformer2DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import ( + USE_PEFT_BACKEND, + is_torch_xla_available, + logging, + replace_example_docstring, + scale_lora_layers, + unscale_lora_layers, +) +from ...utils.torch_utils import randn_tensor +from ..pipeline_utils import DiffusionPipeline +from .pipeline_output import ChromaPipelineOutput + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import ChromaRadiancePipeline + + >>> model_id = "lodestones/Chroma1-Radiance" + >>> ckpt_path = "https://huggingface.co/lodestones/Chroma1-HD/blob/main/latest_x0.safetensors" + >>> transformer = ChromaTransformer2DModel.from_single_file(ckpt_path, torch_dtype=torch.bfloat16) + >>> pipe = ChromaRadiancePipeline.from_pretrained( + ... model_id, + ... transformer=transformer, + ... torch_dtype=torch.bfloat16, + ... ) + >>> pipe.enable_model_cpu_offload() + >>> prompt = [ + ... "A high-fashion close-up portrait of a blonde woman in clear sunglasses. The image uses a bold teal and red color split for dramatic lighting. The background is a simple teal-green. The photo is sharp and well-composed, and is designed for viewing with anaglyph 3D glasses for optimal effect. It looks professionally done." + ... ] + >>> negative_prompt = [ + ... "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors" + ... ] + >>> image = pipe(prompt, negative_prompt=negative_prompt).images[0] + >>> image.save("chroma.png") + ``` +""" + + +# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class ChromaRadiancePipeline( + DiffusionPipeline, + FluxLoraLoaderMixin, + FromSingleFileMixin, + TextualInversionLoaderMixin, + FluxIPAdapterMixin, +): + r""" + The Chroma Radiance pipeline for text-to-image generation. + + Reference: https://huggingface.co/lodestones/Chroma1-Radiance/ + + Args: + transformer ([`ChromaRadianceTransformer2DModel`]): + Conditional Transformer (MMDiT) architecture to denoise the encoded image latents. + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + text_encoder ([`T5EncoderModel`]): + [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically + the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant. + tokenizer (`T5TokenizerFast`): + Tokenizer of class + [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast). + """ + + model_cpu_offload_seq = "text_encoder->image_encoder->transformer" + _optional_components = ["image_encoder", "feature_extractor"] + _callback_tensor_inputs = ["latents", "prompt_embeds"] + + def __init__( + self, + scheduler: FlowMatchEulerDiscreteScheduler, + text_encoder: T5EncoderModel, + tokenizer: T5TokenizerFast, + transformer: ChromaRadianceTransformer2DModel, + image_encoder: CLIPVisionModelWithProjection = None, + feature_extractor: CLIPImageProcessor = None, + ): + super().__init__() + + self.register_modules( + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + scheduler=scheduler, + image_encoder=image_encoder, + feature_extractor=feature_extractor, + ) + self.image_processor = VaeImageProcessor() + self.default_sample_size = 1024 + + def _get_t5_prompt_embeds( + self, + prompt: Union[str, List[str]] = None, + num_images_per_prompt: int = 1, + max_sequence_length: int = 512, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + if isinstance(self, TextualInversionLoaderMixin): + prompt = self.maybe_convert_prompt(prompt, self.tokenizer) + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + return_length=False, + return_overflowing_tokens=False, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + tokenizer_mask = text_inputs.attention_mask + + tokenizer_mask_device = tokenizer_mask.to(device) + + # unlike FLUX, Chroma uses the attention mask when generating the T5 embedding + prompt_embeds = self.text_encoder( + text_input_ids.to(device), + output_hidden_states=False, + attention_mask=tokenizer_mask_device, + )[0] + + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + + # for the text tokens, chroma requires that all except the first padding token are masked out during the forward pass through the transformer + seq_lengths = tokenizer_mask_device.sum(dim=1) + mask_indices = torch.arange(tokenizer_mask_device.size(1), device=device).unsqueeze(0).expand(batch_size, -1) + attention_mask = (mask_indices <= seq_lengths.unsqueeze(1)).to(dtype=dtype, device=device) + + _, seq_len, _ = prompt_embeds.shape + + # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + attention_mask = attention_mask.repeat(1, num_images_per_prompt) + attention_mask = attention_mask.view(batch_size * num_images_per_prompt, seq_len) + + return prompt_embeds, attention_mask + + def encode_prompt( + self, + prompt: Union[str, List[str]], + negative_prompt: Union[str, List[str]] = None, + device: Optional[torch.device] = None, + num_images_per_prompt: int = 1, + prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_attention_mask: Optional[torch.Tensor] = None, + negative_prompt_attention_mask: Optional[torch.Tensor] = None, + do_classifier_free_guidance: bool = True, + max_sequence_length: int = 512, + lora_scale: Optional[float] = None, + ): + r""" + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + negative_prompt (`str` or `List[str]`, *optional*): + The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` + instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). + device: (`torch.device`): + torch device + num_images_per_prompt (`int`): + number of images that should be generated per prompt + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + lora_scale (`float`, *optional*): + A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded. + """ + device = device or self._execution_device + + # set lora scale so that monkey patched LoRA + # function of text encoder can correctly access it + if lora_scale is not None and isinstance(self, FluxLoraLoaderMixin): + self._lora_scale = lora_scale + + # dynamically adjust the LoRA scale + if self.text_encoder is not None and USE_PEFT_BACKEND: + scale_lora_layers(self.text_encoder, lora_scale) + + prompt = [prompt] if isinstance(prompt, str) else prompt + + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds, prompt_attention_mask = self._get_t5_prompt_embeds( + prompt=prompt, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + ) + + dtype = self.text_encoder.dtype if self.text_encoder is not None else self.transformer.dtype + text_ids = torch.zeros(prompt_embeds.shape[1], 3).to(device=device, dtype=dtype) + negative_text_ids = None + + if do_classifier_free_guidance: + if negative_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt = ( + batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + ) + + if prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + + negative_prompt_embeds, negative_prompt_attention_mask = self._get_t5_prompt_embeds( + prompt=negative_prompt, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + ) + + negative_text_ids = torch.zeros(negative_prompt_embeds.shape[1], 3).to(device=device, dtype=dtype) + + if self.text_encoder is not None: + if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND: + # Retrieve the original scale by scaling back the LoRA layers + unscale_lora_layers(self.text_encoder, lora_scale) + + return ( + prompt_embeds, + text_ids, + prompt_attention_mask, + negative_prompt_embeds, + negative_text_ids, + negative_prompt_attention_mask, + ) + + # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_image + def encode_image(self, image, device, num_images_per_prompt): + dtype = next(self.image_encoder.parameters()).dtype + + if not isinstance(image, torch.Tensor): + image = self.feature_extractor(image, return_tensors="pt").pixel_values + + image = image.to(device=device, dtype=dtype) + image_embeds = self.image_encoder(image).image_embeds + image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) + return image_embeds + + # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.prepare_ip_adapter_image_embeds + def prepare_ip_adapter_image_embeds( + self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt + ): + image_embeds = [] + if ip_adapter_image_embeds is None: + if not isinstance(ip_adapter_image, list): + ip_adapter_image = [ip_adapter_image] + + if len(ip_adapter_image) != self.transformer.encoder_hid_proj.num_ip_adapters: + raise ValueError( + f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {self.transformer.encoder_hid_proj.num_ip_adapters} IP Adapters." + ) + + for single_ip_adapter_image in ip_adapter_image: + single_image_embeds = self.encode_image(single_ip_adapter_image, device, 1) + image_embeds.append(single_image_embeds[None, :]) + else: + if not isinstance(ip_adapter_image_embeds, list): + ip_adapter_image_embeds = [ip_adapter_image_embeds] + + if len(ip_adapter_image_embeds) != self.transformer.encoder_hid_proj.num_ip_adapters: + raise ValueError( + f"`ip_adapter_image_embeds` must have same length as the number of IP Adapters. Got {len(ip_adapter_image_embeds)} image embeds and {self.transformer.encoder_hid_proj.num_ip_adapters} IP Adapters." + ) + + for single_image_embeds in ip_adapter_image_embeds: + image_embeds.append(single_image_embeds) + + ip_adapter_image_embeds = [] + for single_image_embeds in image_embeds: + single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) + single_image_embeds = single_image_embeds.to(device=device) + ip_adapter_image_embeds.append(single_image_embeds) + + return ip_adapter_image_embeds + + def check_inputs( + self, + prompt, + height, + width, + negative_prompt=None, + prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_embeds=None, + negative_prompt_attention_mask=None, + callback_on_step_end_tensor_inputs=None, + max_sequence_length=None, + ): + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if prompt_embeds is not None and prompt_attention_mask is None: + raise ValueError("Cannot provide `prompt_embeds` without also providing `prompt_attention_mask") + + if negative_prompt_embeds is not None and negative_prompt_attention_mask is None: + raise ValueError( + "Cannot provide `negative_prompt_embeds` without also providing `negative_prompt_attention_mask" + ) + + if max_sequence_length is not None and max_sequence_length > 512: + raise ValueError(f"`max_sequence_length` cannot be greater than 512 but is {max_sequence_length}") + + @staticmethod + def _prepare_latent_image_ids(batch_size, height, width, device, dtype): + latent_image_ids = torch.zeros(height, width, 3) + latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None] + latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :] + + latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape + + latent_image_ids = latent_image_ids.reshape( + latent_image_id_height * latent_image_id_width, latent_image_id_channels + ) + + return latent_image_ids.to(device=device, dtype=dtype) + + def prepare_latents( + self, + batch_size, + num_channels_latents, + height, + width, + dtype, + device, + generator, + latents=None, + patch_size=2, + ): + shape = (batch_size, num_channels_latents, height, width) + + if latents is not None: + latent_image_ids = self._prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype) + return latents.to(device=device, dtype=dtype), latent_image_ids + + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + # latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width) + + latent_image_ids = self._prepare_latent_image_ids( + batch_size, height // patch_size, width // patch_size, device, dtype + ) + + return latents, latent_image_ids + + def _prepare_attention_mask( + self, + batch_size, + sequence_length, + dtype, + attention_mask=None, + ): + if attention_mask is None: + return attention_mask + + # Extend the prompt attention mask to account for image tokens in the final sequence + attention_mask = torch.cat( + [attention_mask, torch.ones(batch_size, sequence_length, device=attention_mask.device, dtype=torch.bool)], + dim=1, + ) + + return attention_mask + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def joint_attention_kwargs(self): + return self._joint_attention_kwargs + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + negative_prompt: Union[str, List[str]] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 35, + sigmas: Optional[List[float]] = None, + guidance_scale: float = 5.0, + num_images_per_prompt: Optional[int] = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + ip_adapter_image: Optional[PipelineImageInput] = None, + ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, + negative_ip_adapter_image: Optional[PipelineImageInput] = None, + negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_attention_mask: Optional[torch.Tensor] = None, + negative_prompt_attention_mask: Optional[torch.Tensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + max_sequence_length: int = 512, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + not greater than `1`). + height (`int`, *optional*, defaults to self.unet.config.sample_size): + The height in pixels of the generated image. This is set to 1024 by default for the best results. + width (`int`, *optional*, defaults to self.unet.config.sample_size): + The width in pixels of the generated image. This is set to 1024 by default for the best results. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + guidance_scale (`float`, *optional*, defaults to 3.5): + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. + of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting + `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to + the text `prompt`, usually at the expense of lower image quality. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. + ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of + IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not + provided, embeddings are computed from the `ip_adapter_image` input argument. + negative_ip_adapter_image: + (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. + negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of + IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not + provided, embeddings are computed from the `ip_adapter_image` input argument. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + prompt_attention_mask (torch.Tensor, *optional*): + Attention mask for the prompt embeddings. Used to mask out padding tokens in the prompt sequence. + Chroma requires a single padding token remain unmasked. Please refer to + https://huggingface.co/lodestones/Chroma#tldr-masking-t5-padding-tokens-enhanced-fidelity-and-increased-stability-during-training + negative_prompt_attention_mask (torch.Tensor, *optional*): + Attention mask for the negative prompt embeddings. Used to mask out padding tokens in the negative + prompt sequence. Chroma requires a single padding token remain unmasked. PLease refer to + https://huggingface.co/lodestones/Chroma#tldr-masking-t5-padding-tokens-enhanced-fidelity-and-increased-stability-during-training + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.flux.ChromaPipelineOutput`] instead of a plain tuple. + joint_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`. + + Examples: + + Returns: + [`~pipelines.chroma.ChromaPipelineOutput`] or `tuple`: [`~pipelines.chroma.ChromaPipelineOutput`] if + `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the + generated images. + """ + + height = height or self.default_sample_size + width = width or self.default_sample_size + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + height, + width, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_embeds=negative_prompt_embeds, + negative_prompt_attention_mask=negative_prompt_attention_mask, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + max_sequence_length=max_sequence_length, + ) + + self._guidance_scale = guidance_scale + self._joint_attention_kwargs = joint_attention_kwargs + self._current_timestep = None + self._interrupt = False + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + + lora_scale = ( + self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None + ) + ( + prompt_embeds, + text_ids, + prompt_attention_mask, + negative_prompt_embeds, + negative_text_ids, + negative_prompt_attention_mask, + ) = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + do_classifier_free_guidance=self.do_classifier_free_guidance, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + lora_scale=lora_scale, + ) + + # 4. Prepare latent variables + num_channels_latents = self.transformer.config.in_channels + latents, latent_image_ids = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + patch_size=self.transformer.config.patch_size, + ) + + # 5. Prepare timesteps + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas + image_seq_len = latent_image_ids.shape[0] + mu = calculate_shift( + image_seq_len, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.5), + self.scheduler.config.get("max_shift", 1.15), + ) + + attention_mask = self._prepare_attention_mask( + batch_size=latents.shape[0], + sequence_length=image_seq_len, + dtype=latents.dtype, + attention_mask=prompt_attention_mask, + ) + negative_attention_mask = self._prepare_attention_mask( + batch_size=latents.shape[0], + sequence_length=image_seq_len, + dtype=latents.dtype, + attention_mask=negative_prompt_attention_mask, + ) + + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + sigmas=sigmas, + mu=mu, + ) + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + if (ip_adapter_image is not None or ip_adapter_image_embeds is not None) and ( + negative_ip_adapter_image is None and negative_ip_adapter_image_embeds is None + ): + negative_ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8) + negative_ip_adapter_image = [negative_ip_adapter_image] * self.transformer.encoder_hid_proj.num_ip_adapters + + elif (ip_adapter_image is None and ip_adapter_image_embeds is None) and ( + negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None + ): + ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8) + ip_adapter_image = [ip_adapter_image] * self.transformer.encoder_hid_proj.num_ip_adapters + + if self.joint_attention_kwargs is None: + self._joint_attention_kwargs = {} + + image_embeds = None + negative_image_embeds = None + if ip_adapter_image is not None or ip_adapter_image_embeds is not None: + image_embeds = self.prepare_ip_adapter_image_embeds( + ip_adapter_image, + ip_adapter_image_embeds, + device, + batch_size * num_images_per_prompt, + ) + if negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None: + negative_image_embeds = self.prepare_ip_adapter_image_embeds( + negative_ip_adapter_image, + negative_ip_adapter_image_embeds, + device, + batch_size * num_images_per_prompt, + ) + + # 6. Denoising loop + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + self._current_timestep = t + if image_embeds is not None: + self._joint_attention_kwargs["ip_adapter_image_embeds"] = image_embeds + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latents.shape[0]).to(latents.dtype) + + noise_pred = self.transformer( + hidden_states=latents, + timestep=timestep / 1000, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, + img_ids=latent_image_ids, + attention_mask=attention_mask, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + )[0] + + if self.do_classifier_free_guidance: + if negative_image_embeds is not None: + self._joint_attention_kwargs["ip_adapter_image_embeds"] = negative_image_embeds + neg_noise_pred = self.transformer( + hidden_states=latents, + timestep=timestep / 1000, + encoder_hidden_states=negative_prompt_embeds, + txt_ids=negative_text_ids, + img_ids=latent_image_ids, + attention_mask=negative_attention_mask, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + )[0] + noise_pred = neg_noise_pred + guidance_scale * (noise_pred - neg_noise_pred) + + # compute the previous noisy sample x_t -> x_t-1 + latents_dtype = latents.dtype + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + if latents.dtype != latents_dtype: + if torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 + latents = latents.to(latents_dtype) + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + self._current_timestep = None + + # 7. + + if output_type == "latent": + image = latents + else: + image = self.image_processor.postprocess(latents, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return ChromaPipelineOutput(images=image) diff --git a/src/diffusers/pipelines/chroma/pipeline_output.py b/src/diffusers/pipelines/chroma/pipeline_output.py index 951d132dba2e..229b0fe42b90 100644 --- a/src/diffusers/pipelines/chroma/pipeline_output.py +++ b/src/diffusers/pipelines/chroma/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -13,9 +12,9 @@ class ChromaPipelineOutput(BaseOutput): Output class for Stable Diffusion pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image] | np.ndarray diff --git a/src/diffusers/pipelines/chronoedit/pipeline_chronoedit.py b/src/diffusers/pipelines/chronoedit/pipeline_chronoedit.py index 79f6580fbed6..1e0cc0ea5c2a 100644 --- a/src/diffusers/pipelines/chronoedit/pipeline_chronoedit.py +++ b/src/diffusers/pipelines/chronoedit/pipeline_chronoedit.py @@ -13,7 +13,7 @@ # limitations under the License. import html -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import PIL import regex as re @@ -113,7 +113,7 @@ def prompt_clean(text): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -185,11 +185,11 @@ def __init__( # Copied from diffusers.pipelines.wan.pipeline_wan_i2v.WanImageToVideoPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -228,7 +228,7 @@ def _get_t5_prompt_embeds( def encode_image( self, image: PipelineImageInput, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): device = device or self._execution_device image = self.image_processor(images=image, return_tensors="pt").to(device) @@ -238,23 +238,23 @@ def encode_image( # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -381,11 +381,11 @@ def prepare_latents( height: int = 480, width: int = 832, num_frames: int = 81, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 latent_height = height // self.vae_scale_factor_spatial latent_width = width // self.vae_scale_factor_spatial @@ -469,26 +469,24 @@ def attention_kwargs(self): def __call__( self, image: PipelineImageInput, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, height: int = 480, width: int = 832, num_frames: int = 81, num_inference_steps: int = 50, guidance_scale: float = 5.0, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - image_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "np", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + image_embeds: torch.Tensor | None = None, + output_type: str | None = "np", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int, None], PipelineCallback | MultiPipelineCallbacks] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, enable_temporal_reasoning: bool = False, num_temporal_reasoning_steps: int = 0, @@ -499,10 +497,10 @@ def __call__( Args: image (`PipelineImageInput`): The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -523,7 +521,7 @@ def __call__( usually at the expense of lower image quality. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/chronoedit/pipeline_output.py b/src/diffusers/pipelines/chronoedit/pipeline_output.py index b1df5b9de35d..6247ce9f3a0c 100644 --- a/src/diffusers/pipelines/chronoedit/pipeline_output.py +++ b/src/diffusers/pipelines/chronoedit/pipeline_output.py @@ -11,7 +11,7 @@ class ChronoEditPipelineOutput(BaseOutput): Output class for ChronoEdit pipelines. Args: - frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape `(batch_size, num_frames, channels, height, width)`. diff --git a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox.py b/src/diffusers/pipelines/cogvideo/pipeline_cogvideox.py index 4ac33b24bbe1..b883e10a6732 100644 --- a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox.py +++ b/src/diffusers/pipelines/cogvideo/pipeline_cogvideox.py @@ -15,7 +15,7 @@ import inspect import math -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from transformers import T5EncoderModel, T5Tokenizer @@ -87,10 +87,10 @@ def get_resize_crop_region_for_grid(src, tgt_width, tgt_height): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -105,15 +105,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -182,7 +182,7 @@ def __init__( text_encoder: T5EncoderModel, vae: AutoencoderKLCogVideoX, transformer: CogVideoXTransformer3DModel, - scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler], + scheduler: CogVideoXDDIMScheduler | CogVideoXDPMScheduler, ): super().__init__() @@ -201,11 +201,11 @@ def __init__( def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -243,23 +243,23 @@ def _get_t5_prompt_embeds( def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -444,7 +444,7 @@ def _prepare_rotary_positional_embeddings( width: int, num_frames: int, device: torch.device, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: grid_height = height // (self.vae_scale_factor_spatial * self.transformer.config.patch_size) grid_width = width // (self.vae_scale_factor_spatial * self.transformer.config.patch_size) @@ -506,38 +506,36 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Optional[Union[str, List[str]]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_frames: Optional[int] = None, + prompt: str | list[str] | None = None, + negative_prompt: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, + num_frames: int | None = None, num_inference_steps: int = 50, - timesteps: Optional[List[int]] = None, + timesteps: list[int] | None = None, guidance_scale: float = 6, use_dynamic_cfg: bool = False, num_videos_per_prompt: int = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, output_type: str = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 226, - ) -> Union[CogVideoXPipelineOutput, Tuple]: + ) -> CogVideoXPipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -553,7 +551,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. @@ -565,7 +563,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of videos to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -594,7 +592,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -664,7 +662,13 @@ def __call__( prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, timesteps + ) self._num_timesteps = len(timesteps) # 5. Prepare latents diff --git a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py b/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py index c1335839f848..de5b969a9adc 100644 --- a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py +++ b/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py @@ -15,7 +15,7 @@ import inspect import math -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from PIL import Image @@ -94,10 +94,10 @@ def get_resize_crop_region_for_grid(src, tgt_width, tgt_height): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -112,15 +112,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -209,11 +209,11 @@ def __init__( # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -252,23 +252,23 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -360,8 +360,8 @@ def prepare_latents( # Adapted from https://github.com/aigc-apps/CogVideoX-Fun/blob/2a93e5c14e02b2b5921d533fd59fc8c0ed69fb24/cogvideox/pipeline/pipeline_cogvideox_control.py#L366 def prepare_control_latents( - self, mask: Optional[torch.Tensor] = None, masked_image: Optional[torch.Tensor] = None - ) -> Tuple[torch.Tensor, torch.Tensor]: + self, mask: torch.Tensor | None = None, masked_image: torch.Tensor | None = None + ) -> tuple[torch.Tensor, torch.Tensor]: if mask is not None: masks = [] for i in range(mask.size(0)): @@ -490,7 +490,7 @@ def _prepare_rotary_positional_embeddings( width: int, num_frames: int, device: torch.device, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: grid_height = height // (self.vae_scale_factor_spatial * self.transformer.config.patch_size) grid_width = width // (self.vae_scale_factor_spatial * self.transformer.config.patch_size) @@ -552,43 +552,41 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Optional[Union[str, List[str]]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - control_video: Optional[List[Image.Image]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] | None = None, + negative_prompt: str | list[str] | None = None, + control_video: list[Image.Image] | None = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: Optional[List[int]] = None, + timesteps: list[int] | None = None, guidance_scale: float = 6, use_dynamic_cfg: bool = False, num_videos_per_prompt: int = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - control_video_latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + control_video_latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, output_type: str = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 226, - ) -> Union[CogVideoXPipelineOutput, Tuple]: + ) -> CogVideoXPipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - control_video (`List[PIL.Image.Image]`): + control_video (`list[PIL.Image.Image]`): The control video to condition the generation on. Must be a list of images/frames of the video. If not provided, `control_video_latents` must be provided. height (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial): @@ -598,7 +596,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. @@ -610,7 +608,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of videos to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -642,7 +640,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -717,7 +715,13 @@ def __call__( prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, timesteps + ) self._num_timesteps = len(timesteps) # 5. Prepare latents diff --git a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py b/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py index c523c9adec98..9687d63bc7bf 100644 --- a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py +++ b/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py @@ -15,7 +15,7 @@ import inspect import math -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import PIL import torch @@ -86,10 +86,10 @@ def get_resize_crop_region_for_grid(src, tgt_width, tgt_height): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -104,15 +104,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -145,7 +145,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -195,7 +195,7 @@ def __init__( text_encoder: T5EncoderModel, vae: AutoencoderKLCogVideoX, transformer: CogVideoXTransformer3DModel, - scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler], + scheduler: CogVideoXDDIMScheduler | CogVideoXDPMScheduler, ): super().__init__() @@ -219,11 +219,11 @@ def __init__( # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -262,23 +262,23 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -349,10 +349,10 @@ def prepare_latents( num_frames: int = 13, height: int = 60, width: int = 90, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[torch.Generator] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, ): if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( @@ -470,7 +470,7 @@ def check_inputs( and not isinstance(image, list) ): raise ValueError( - "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is" + "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]` but is" f" {type(image)}" ) @@ -537,7 +537,7 @@ def _prepare_rotary_positional_embeddings( width: int, num_frames: int, device: torch.device, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: grid_height = height // (self.vae_scale_factor_spatial * self.transformer.config.patch_size) grid_width = width // (self.vae_scale_factor_spatial * self.transformer.config.patch_size) @@ -600,40 +600,38 @@ def interrupt(self): def __call__( self, image: PipelineImageInput, - prompt: Optional[Union[str, List[str]]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] | None = None, + negative_prompt: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, num_frames: int = 49, num_inference_steps: int = 50, - timesteps: Optional[List[int]] = None, + timesteps: list[int] | None = None, guidance_scale: float = 6, use_dynamic_cfg: bool = False, num_videos_per_prompt: int = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, output_type: str = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 226, - ) -> Union[CogVideoXPipelineOutput, Tuple]: + ) -> CogVideoXPipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: image (`PipelineImageInput`): The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -649,7 +647,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. @@ -661,7 +659,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of videos to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -690,7 +688,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -762,7 +760,13 @@ def __call__( prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, timesteps + ) self._num_timesteps = len(timesteps) # 5. Prepare latents diff --git a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py b/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py index 897dc6d1b70a..e3ce8292fad6 100644 --- a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py +++ b/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py @@ -15,7 +15,7 @@ import inspect import math -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from PIL import Image @@ -95,10 +95,10 @@ def get_resize_crop_region_for_grid(src, tgt_width, tgt_height): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -113,15 +113,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -154,7 +154,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -204,7 +204,7 @@ def __init__( text_encoder: T5EncoderModel, vae: AutoencoderKLCogVideoX, transformer: CogVideoXTransformer3DModel, - scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler], + scheduler: CogVideoXDDIMScheduler | CogVideoXDPMScheduler, ): super().__init__() @@ -225,11 +225,11 @@ def __init__( # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -268,23 +268,23 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -349,16 +349,16 @@ def encode_prompt( def prepare_latents( self, - video: Optional[torch.Tensor] = None, + video: torch.Tensor | None = None, batch_size: int = 1, num_channels_latents: int = 16, height: int = 60, width: int = 90, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[torch.Generator] = None, - latents: Optional[torch.Tensor] = None, - timestep: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, + timestep: torch.Tensor | None = None, ): if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( @@ -514,7 +514,7 @@ def _prepare_rotary_positional_embeddings( width: int, num_frames: int, device: torch.device, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: grid_height = height // (self.vae_scale_factor_spatial * self.transformer.config.patch_size) grid_width = width // (self.vae_scale_factor_spatial * self.transformer.config.patch_size) @@ -576,41 +576,39 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - video: List[Image.Image] = None, - prompt: Optional[Union[str, List[str]]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + video: list[Image.Image] = None, + prompt: str | list[str] | None = None, + negative_prompt: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: Optional[List[int]] = None, + timesteps: list[int] | None = None, strength: float = 0.8, guidance_scale: float = 6, use_dynamic_cfg: bool = False, num_videos_per_prompt: int = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, output_type: str = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 226, - ) -> Union[CogVideoXPipelineOutput, Tuple]: + ) -> CogVideoXPipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: - video (`List[PIL.Image.Image]`): + video (`list[PIL.Image.Image]`): The input video to condition the generation on. Must be a list of images/frames of the video. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -621,7 +619,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. @@ -635,7 +633,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of videos to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -664,7 +662,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -737,7 +735,13 @@ def __call__( prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, timesteps + ) timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, timesteps, strength, device) latent_timestep = timesteps[:1].repeat(batch_size * num_videos_per_prompt) self._num_timesteps = len(timesteps) diff --git a/src/diffusers/pipelines/cogvideo/pipeline_output.py b/src/diffusers/pipelines/cogvideo/pipeline_output.py index 3de030dd6928..56ff50132231 100644 --- a/src/diffusers/pipelines/cogvideo/pipeline_output.py +++ b/src/diffusers/pipelines/cogvideo/pipeline_output.py @@ -11,8 +11,8 @@ class CogVideoXPipelineOutput(BaseOutput): Output class for CogVideo pipelines. Args: - frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): - List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): + list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape `(batch_size, num_frames, channels, height, width)`. """ diff --git a/src/diffusers/pipelines/cogview3/pipeline_cogview3plus.py b/src/diffusers/pipelines/cogview3/pipeline_cogview3plus.py index 304a5c5ad00b..8880e3a0d1e2 100644 --- a/src/diffusers/pipelines/cogview3/pipeline_cogview3plus.py +++ b/src/diffusers/pipelines/cogview3/pipeline_cogview3plus.py @@ -14,7 +14,7 @@ # limitations under the License. import inspect -from typing import Callable, Dict, List, Optional, Tuple, Union +from typing import Callable import torch from transformers import T5EncoderModel, T5Tokenizer @@ -58,10 +58,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -76,15 +76,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -153,7 +153,7 @@ def __init__( text_encoder: T5EncoderModel, vae: AutoencoderKL, transformer: CogView3PlusTransformer2DModel, - scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler], + scheduler: CogVideoXDDIMScheduler | CogVideoXDPMScheduler, ): super().__init__() @@ -167,11 +167,11 @@ def __init__( # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._get_t5_prompt_embeds with num_videos_per_prompt->num_images_per_prompt def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -209,23 +209,23 @@ def _get_t5_prompt_embeds( def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 224, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -408,36 +408,34 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Optional[Union[str, List[str]]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] | None = None, + negative_prompt: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: Optional[List[int]] = None, + timesteps: list[int] | None = None, guidance_scale: float = 5.0, num_images_per_prompt: int = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - original_size: Optional[Tuple[int, int]] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + original_size: tuple[int, int] | None = None, + crops_coords_top_left: tuple[int, int] = (0, 0), output_type: str = "pil", return_dict: bool = True, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 224, - ) -> Union[CogView3PipelineOutput, Tuple]: + ) -> CogView3PipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -448,7 +446,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to `50`): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. @@ -460,7 +458,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to `1`): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -474,12 +472,12 @@ def __call__( Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of @@ -499,7 +497,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -566,7 +564,13 @@ def __call__( prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, timesteps + ) self._num_timesteps = len(timesteps) # 5. Prepare latents. diff --git a/src/diffusers/pipelines/cogview3/pipeline_output.py b/src/diffusers/pipelines/cogview3/pipeline_output.py index 3891dd51e691..6c89e117b74c 100644 --- a/src/diffusers/pipelines/cogview3/pipeline_output.py +++ b/src/diffusers/pipelines/cogview3/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -13,9 +12,9 @@ class CogView3PipelineOutput(BaseOutput): Output class for CogView3 pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image] | np.ndarray diff --git a/src/diffusers/pipelines/cogview4/pipeline_cogview4.py b/src/diffusers/pipelines/cogview4/pipeline_cogview4.py index 22510f5d9d50..329b76d11e0d 100644 --- a/src/diffusers/pipelines/cogview4/pipeline_cogview4.py +++ b/src/diffusers/pipelines/cogview4/pipeline_cogview4.py @@ -14,7 +14,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import torch @@ -69,10 +69,10 @@ def calculate_shift( def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -87,15 +87,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) @@ -177,10 +177,10 @@ def __init__( def _get_glm_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, max_sequence_length: int = 1024, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -220,23 +220,23 @@ def _get_glm_embeds( def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, max_sequence_length: int = 1024, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -403,37 +403,35 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Optional[Union[str, List[str]]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] | None = None, + negative_prompt: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, guidance_scale: float = 5.0, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - original_size: Optional[Tuple[int, int]] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + original_size: tuple[int, int] | None = None, + crops_coords_top_left: tuple[int, int] = (0, 0), output_type: str = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 1024, - ) -> Union[CogView4PipelineOutput, Tuple]: + ) -> CogView4PipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -444,11 +442,11 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to `50`): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -460,7 +458,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to `1`): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -474,12 +472,12 @@ def __call__( Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of @@ -499,7 +497,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -599,8 +597,12 @@ def __call__( self.scheduler.config.get("base_shift", 0.25), self.scheduler.config.get("max_shift", 0.75), ) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas, mu=mu + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas, mu=mu ) self._num_timesteps = len(timesteps) diff --git a/src/diffusers/pipelines/cogview4/pipeline_cogview4_control.py b/src/diffusers/pipelines/cogview4/pipeline_cogview4_control.py index e26b7ba415de..6282bf4cd7a4 100644 --- a/src/diffusers/pipelines/cogview4/pipeline_cogview4_control.py +++ b/src/diffusers/pipelines/cogview4/pipeline_cogview4_control.py @@ -14,7 +14,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import torch @@ -71,10 +71,10 @@ def calculate_shift( # Copied from diffusers.pipelines.cogview4.pipeline_cogview4.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -89,15 +89,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) @@ -180,10 +180,10 @@ def __init__( # Copied from diffusers.pipelines.cogview4.pipeline_cogview4.CogView4Pipeline._get_glm_embeds def _get_glm_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, max_sequence_length: int = 1024, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -224,23 +224,23 @@ def _get_glm_embeds( # Copied from diffusers.pipelines.cogview4.pipeline_cogview4.CogView4Pipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, max_sequence_length: int = 1024, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -435,38 +435,36 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Optional[Union[str, List[str]]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] | None = None, + negative_prompt: str | list[str] | None = None, control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, guidance_scale: float = 5.0, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - original_size: Optional[Tuple[int, int]] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + original_size: tuple[int, int] | None = None, + crops_coords_top_left: tuple[int, int] = (0, 0), output_type: str = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 1024, - ) -> Union[CogView4PipelineOutput, Tuple]: + ) -> CogView4PipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -477,11 +475,11 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to `50`): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -493,7 +491,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to `1`): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -507,12 +505,12 @@ def __call__( Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of @@ -532,7 +530,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -649,8 +647,12 @@ def __call__( self.scheduler.config.get("base_shift", 0.25), self.scheduler.config.get("max_shift", 0.75), ) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas, mu=mu + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas, mu=mu ) self._num_timesteps = len(timesteps) # Denoising loop diff --git a/src/diffusers/pipelines/cogview4/pipeline_output.py b/src/diffusers/pipelines/cogview4/pipeline_output.py index 4efec1310845..997444c6c009 100644 --- a/src/diffusers/pipelines/cogview4/pipeline_output.py +++ b/src/diffusers/pipelines/cogview4/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -13,9 +12,9 @@ class CogView4PipelineOutput(BaseOutput): Output class for CogView3 pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image] | np.ndarray diff --git a/src/diffusers/pipelines/consisid/consisid_utils.py b/src/diffusers/pipelines/consisid/consisid_utils.py index 521d4d787e54..c1646e15efbc 100644 --- a/src/diffusers/pipelines/consisid/consisid_utils.py +++ b/src/diffusers/pipelines/consisid/consisid_utils.py @@ -135,7 +135,7 @@ def process_face_embeddings( is_align_face: Boolean flag indicating whether face alignment should be performed. Returns: - Tuple: + tuple: - id_cond: Concatenated tensor of Ante face embedding and CLIP vision embedding - id_vit_hidden: Hidden state of the CLIP vision model, a list of tensors. - return_face_features_image_2: Processed face features image after normalization and parsing. @@ -245,7 +245,7 @@ def process_face_embeddings_infer( is_align_face: Boolean flag indicating whether face alignment should be performed (default: True). Returns: - Tuple: + tuple: - id_cond: Concatenated tensor of Ante face embedding and CLIP vision embedding. - id_vit_hidden: Hidden state of the CLIP vision model, a list of tensors. - image: Processed face image after feature extraction and alignment. diff --git a/src/diffusers/pipelines/consisid/pipeline_consisid.py b/src/diffusers/pipelines/consisid/pipeline_consisid.py index 3e6c149d7f80..20b779bf5aaa 100644 --- a/src/diffusers/pipelines/consisid/pipeline_consisid.py +++ b/src/diffusers/pipelines/consisid/pipeline_consisid.py @@ -14,7 +14,7 @@ import inspect import math -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL @@ -102,7 +102,7 @@ def draw_kps(image_pil, kps, color_list=[(255, 0, 0), (0, 255, 0), (0, 0, 255), Parameters: - image_pil (PIL.Image): Input image as a PIL object. - kps (list of tuples): A list of keypoints where each keypoint is a tuple of (x, y) coordinates. - - color_list (list of tuples, optional): List of colors (in RGB format) for each keypoint. Default is a set of five + - color_list (list of tuples, optional): list of colors (in RGB format) for each keypoint. Default is a set of five colors. Returns: @@ -176,10 +176,10 @@ def get_resize_crop_region_for_grid(src, tgt_width, tgt_height): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -194,15 +194,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -235,7 +235,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -311,11 +311,11 @@ def __init__( # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -354,23 +354,23 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -441,11 +441,11 @@ def prepare_latents( num_frames: int = 13, height: int = 60, width: int = 90, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[torch.Generator] = None, - latents: Optional[torch.Tensor] = None, - kps_cond: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, + kps_cond: torch.Tensor | None = None, ): if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( @@ -572,7 +572,7 @@ def check_inputs( and not isinstance(image, list) ): raise ValueError( - "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is" + "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]` but is" f" {type(image)}" ) @@ -623,7 +623,7 @@ def _prepare_rotary_positional_embeddings( width: int, num_frames: int, device: torch.device, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: grid_height = height // (self.vae_scale_factor_spatial * self.transformer.config.patch_size) grid_width = width // (self.vae_scale_factor_spatial * self.transformer.config.patch_size) base_size_width = self.transformer.config.sample_width // self.transformer.config.patch_size @@ -663,8 +663,8 @@ def interrupt(self): def __call__( self, image: PipelineImageInput, - prompt: Optional[Union[str, List[str]]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] | None = None, + negative_prompt: str | list[str] | None = None, height: int = 480, width: int = 720, num_frames: int = 49, @@ -673,32 +673,30 @@ def __call__( use_dynamic_cfg: bool = False, num_videos_per_prompt: int = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, output_type: str = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 226, - id_vit_hidden: Optional[torch.Tensor] = None, - id_cond: Optional[torch.Tensor] = None, - kps_cond: Optional[torch.Tensor] = None, - ) -> Union[ConsisIDPipelineOutput, Tuple]: + id_vit_hidden: torch.Tensor | None = None, + id_cond: torch.Tensor | None = None, + kps_cond: torch.Tensor | None = None, + ) -> ConsisIDPipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: image (`PipelineImageInput`): The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -727,7 +725,7 @@ def __call__( more faithful image generation, while later steps reduce it for more diverse and natural results. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of videos to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -756,22 +754,22 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. max_sequence_length (`int`, defaults to `226`): Maximum sequence length in encoded prompt. Must be consistent with `self.transformer.config.max_text_seq_length` otherwise may lead to poor results. - id_vit_hidden (`Optional[torch.Tensor]`, *optional*): + id_vit_hidden (`torch.Tensor | None`, *optional*): The tensor representing the hidden features extracted from the face model, which are used to condition the local facial extractor. This is crucial for the model to obtain high-frequency information of the face. If not provided, the local facial extractor will not run normally. - id_cond (`Optional[torch.Tensor]`, *optional*): + id_cond (`torch.Tensor | None`, *optional*): The tensor representing the hidden features extracted from the clip model, which are used to condition the local facial extractor. This is crucial for the model to edit facial features If not provided, the local facial extractor will not run normally. - kps_cond (`Optional[torch.Tensor]`, *optional*): + kps_cond (`torch.Tensor | None`, *optional*): A tensor that determines whether the global facial extractor use keypoint information for conditioning. If provided, this tensor controls whether facial keypoints such as eyes, nose, and mouth landmarks are used during the generation process. This helps ensure the model retains more facial low-frequency diff --git a/src/diffusers/pipelines/consisid/pipeline_output.py b/src/diffusers/pipelines/consisid/pipeline_output.py index dd4a63aa50b9..83a5be8d230b 100644 --- a/src/diffusers/pipelines/consisid/pipeline_output.py +++ b/src/diffusers/pipelines/consisid/pipeline_output.py @@ -11,8 +11,8 @@ class ConsisIDPipelineOutput(BaseOutput): Output class for ConsisID pipelines. Args: - frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): - List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): + list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape `(batch_size, num_frames, channels, height, width)`. """ diff --git a/src/diffusers/pipelines/consistency_models/pipeline_consistency_models.py b/src/diffusers/pipelines/consistency_models/pipeline_consistency_models.py index 1fbdeb1f2741..85e59adc39a4 100644 --- a/src/diffusers/pipelines/consistency_models/pipeline_consistency_models.py +++ b/src/diffusers/pipelines/consistency_models/pipeline_consistency_models.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Callable, List, Optional, Union +from typing import Callable import torch @@ -173,27 +173,27 @@ def check_inputs(self, num_inference_steps, timesteps, latents, batch_size, img_ def __call__( self, batch_size: int = 1, - class_labels: Optional[Union[torch.Tensor, List[int], int]] = None, + class_labels: torch.Tensor | list[int] | int | None = None, num_inference_steps: int = 1, - timesteps: List[int] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + timesteps: list[int] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, ): r""" Args: batch_size (`int`, *optional*, defaults to 1): The number of images to generate. - class_labels (`torch.Tensor` or `List[int]` or `int`, *optional*): + class_labels (`torch.Tensor` or `list[int]` or `int`, *optional*): Optional class labels for conditioning class-conditional consistency models. Not used if the model is not class-conditional. num_inference_steps (`int`, *optional*, defaults to 1): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps` timesteps are used. Must be in descending order. generator (`torch.Generator`, *optional*): diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet.py index fe0e69314cca..86fa135abff4 100644 --- a/src/diffusers/pipelines/controlnet/pipeline_controlnet.py +++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet.py @@ -14,7 +14,7 @@ import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -102,10 +102,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -120,15 +120,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -189,7 +189,7 @@ class StableDiffusionControlNetPipeline( A `CLIPTokenizer` to tokenize text. unet ([`UNet2DConditionModel`]): A `UNet2DConditionModel` to denoise the encoded image latents. - controlnet ([`ControlNetModel`] or `List[ControlNetModel]`): + controlnet ([`ControlNetModel`] or `list[ControlNetModel]`): Provides additional conditioning to the `unet` during the denoising process. If you set multiple ControlNets as a list, the outputs from each ControlNet are added together to create one combined additional conditioning. @@ -215,7 +215,7 @@ def __init__( text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, unet: UNet2DConditionModel, - controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel], + controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel, scheduler: KarrasDiffusionSchedulers, safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, @@ -269,9 +269,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -302,16 +302,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -319,7 +319,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -418,7 +418,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -908,45 +908,43 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, + cross_attention_kwargs: dict[str, Any] | None = None, + controlnet_conditioning_scale: float | list[float] = 1.0, guess_mode: bool = False, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: - `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,: + `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`): The ControlNet input condition to provide guidance to the `unet` for generation. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or @@ -962,18 +960,18 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -981,7 +979,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -995,7 +993,7 @@ def __call__( Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -1014,16 +1012,16 @@ def __call__( cross_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set the corresponding scale as a list. guess_mode (`bool`, *optional*, defaults to `False`): The ControlNet encoder tries to recognize the content of the input image even if you remove all prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended. - control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the ControlNet starts applying. - control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the ControlNet stops applying. clip_skip (`int`, *optional*): Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that @@ -1033,7 +1031,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1195,8 +1193,12 @@ def __call__( assert False # 5. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) self._num_timesteps = len(timesteps) diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_blip_diffusion.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_blip_diffusion.py index e0f1879405aa..4ca92b906842 100644 --- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_blip_diffusion.py +++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_blip_diffusion.py @@ -12,8 +12,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Optional, Union - import PIL.Image import torch from transformers import CLIPTokenizer @@ -126,8 +124,8 @@ def __init__( controlnet: ControlNetModel, image_processor: BlipImageProcessor, ctx_begin_pos: int = 2, - mean: List[float] = None, - std: List[float] = None, + mean: list[float] = None, + std: list[float] = None, ): super().__init__() @@ -241,36 +239,36 @@ def prepare_control_image( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: List[str], + prompt: list[str], reference_image: PIL.Image.Image, condtioning_image: PIL.Image.Image, - source_subject_category: List[str], - target_subject_category: List[str], - latents: Optional[torch.Tensor] = None, + source_subject_category: list[str], + target_subject_category: list[str], + latents: torch.Tensor | None = None, guidance_scale: float = 7.5, height: int = 512, width: int = 512, num_inference_steps: int = 50, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - neg_prompt: Optional[str] = "", + generator: torch.Generator | list[torch.Generator] | None = None, + neg_prompt: str | None = "", prompt_strength: float = 1.0, prompt_reps: int = 20, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, ): """ Function invoked when calling the pipeline for generation. Args: - prompt (`List[str]`): + prompt (`list[str]`): The prompt or prompts to guide the image generation. reference_image (`PIL.Image.Image`): The reference image to condition the generation on. condtioning_image (`PIL.Image.Image`): The conditioning canny edge image to condition the generation on. - source_subject_category (`List[str]`): + source_subject_category (`list[str]`): The source subject category. - target_subject_category (`List[str]`): + target_subject_category (`list[str]`): The target subject category. latents (`torch.Tensor`, *optional*): Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image @@ -291,7 +289,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. neg_prompt (`str`, *optional*, defaults to ""): diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py index 12cc6f630d80..f0cfabc66f25 100644 --- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py +++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -101,7 +101,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -167,7 +167,7 @@ class StableDiffusionControlNetImg2ImgPipeline( A `CLIPTokenizer` to tokenize text. unet ([`UNet2DConditionModel`]): A `UNet2DConditionModel` to denoise the encoded image latents. - controlnet ([`ControlNetModel`] or `List[ControlNetModel]`): + controlnet ([`ControlNetModel`] or `list[ControlNetModel]`): Provides additional conditioning to the `unet` during the denoising process. If you set multiple ControlNets as a list, the outputs from each ControlNet are added together to create one combined additional conditioning. @@ -193,7 +193,7 @@ def __init__( text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, unet: UNet2DConditionModel, - controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel], + controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel, scheduler: KarrasDiffusionSchedulers, safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, @@ -247,9 +247,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -280,16 +280,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -297,7 +297,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -396,7 +396,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -906,49 +906,47 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, image: PipelineImageInput = None, control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, strength: float = 0.8, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 0.8, + cross_attention_kwargs: dict[str, Any] | None = None, + controlnet_conditioning_scale: float | list[float] = 0.8, guess_mode: bool = False, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: - `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,: + `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`): The initial image to be used as the starting point for the image generation process. Can also accept image latents as `image`, and if passing latents directly they are not encoded again. - control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: - `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,: + `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`): The ControlNet input condition to provide guidance to the `unet` for generation. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or @@ -971,7 +969,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -979,7 +977,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -993,7 +991,7 @@ def __call__( Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -1006,16 +1004,16 @@ def __call__( cross_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set the corresponding scale as a list. guess_mode (`bool`, *optional*, defaults to `False`): The ControlNet encoder tries to recognize the content of the input image even if you remove all prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended. - control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the ControlNet starts applying. - control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the ControlNet stops applying. clip_skip (`int`, *optional*): Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that @@ -1025,7 +1023,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py index 6de8e5747b02..d34278d0086b 100644 --- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py +++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py @@ -15,7 +15,7 @@ # This model implementation is heavily inspired by https://github.com/haofanwang/ControlNet-for-Diffusers/ import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -113,7 +113,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -163,7 +163,7 @@ class StableDiffusionControlNetInpaintPipeline( A `CLIPTokenizer` to tokenize text. unet ([`UNet2DConditionModel`]): A `UNet2DConditionModel` to denoise the encoded image latents. - controlnet ([`ControlNetModel`] or `List[ControlNetModel]`): + controlnet ([`ControlNetModel`] or `list[ControlNetModel]`): Provides additional conditioning to the `unet` during the denoising process. If you set multiple ControlNets as a list, the outputs from each ControlNet are added together to create one combined additional conditioning. @@ -196,7 +196,7 @@ def __init__( text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, unet: UNet2DConditionModel, - controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel], + controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel, scheduler: KarrasDiffusionSchedulers, safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, @@ -253,9 +253,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -286,16 +286,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -303,7 +303,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -402,7 +402,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -995,62 +995,60 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, image: PipelineImageInput = None, mask_image: PipelineImageInput = None, control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, - padding_mask_crop: Optional[int] = None, + height: int | None = None, + width: int | None = None, + padding_mask_crop: int | None = None, strength: float = 1.0, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 0.5, + cross_attention_kwargs: dict[str, Any] | None = None, + controlnet_conditioning_scale: float | list[float] = 0.5, guess_mode: bool = False, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, - `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, + `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, NumPy array or tensor representing an image batch to be used as the starting point. For both NumPy array and PyTorch tensor, the expected value range is between `[0, 1]`. If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a NumPy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)`. It can also accept image latents as `image`, but if passing latents directly it is not encoded again. - mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, - `List[PIL.Image.Image]`, or `List[np.ndarray]`): + mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, + `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, NumPy array or tensor representing an image batch to mask `image`. White pixels in the mask are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a single channel (luminance) before use. If it's a NumPy array or PyTorch tensor, it should contain one color channel (L) instead of 3, so the expected shape for PyTorch tensor would be `(B, 1, H, W)`, `(B, H, W)`, `(1, H, W)`, `(H, W)`. And for NumPy array, it would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W, 1)`, or `(H, W)`. - control_image (`torch.Tensor`, `PIL.Image.Image`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, - `List[List[torch.Tensor]]`, or `List[List[PIL.Image.Image]]`): + control_image (`torch.Tensor`, `PIL.Image.Image`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, + `list[list[torch.Tensor]]`, or `list[list[PIL.Image.Image]]`): The ControlNet input condition to provide guidance to the `unet` for generation. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or @@ -1080,7 +1078,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -1088,7 +1086,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -1102,7 +1100,7 @@ def __call__( Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -1115,16 +1113,16 @@ def __call__( cross_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 0.5): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 0.5): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set the corresponding scale as a list. guess_mode (`bool`, *optional*, defaults to `False`): The ControlNet encoder tries to recognize the content of the input image even if you remove all prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended. - control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the ControlNet starts applying. - control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the ControlNet stops applying. clip_skip (`int`, *optional*): Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that @@ -1134,7 +1132,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py index fb09d04832f3..942bcb49083e 100644 --- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py +++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -71,7 +71,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -244,13 +244,13 @@ def __init__( tokenizer: CLIPTokenizer, tokenizer_2: CLIPTokenizer, unet: UNet2DConditionModel, - controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel], + controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel, scheduler: KarrasDiffusionSchedulers, requires_aesthetics_score: bool = False, force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, - feature_extractor: Optional[CLIPImageProcessor] = None, - image_encoder: Optional[CLIPVisionModelWithProjection] = None, + add_watermarker: bool | None = None, + feature_extractor: CLIPImageProcessor | None = None, + image_encoder: CLIPVisionModelWithProjection | None = None, ): super().__init__() @@ -291,26 +291,26 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, + prompt_2: str | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders device: (`torch.device`): @@ -319,11 +319,11 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -441,7 +441,7 @@ def encode_prompt( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) - uncond_tokens: List[str] + uncond_tokens: list[str] if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -1166,62 +1166,57 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, image: PipelineImageInput = None, mask_image: PipelineImageInput = None, - control_image: Union[ - PipelineImageInput, - List[PipelineImageInput], - ] = None, - height: Optional[int] = None, - width: Optional[int] = None, - padding_mask_crop: Optional[int] = None, + control_image: PipelineImageInput | list[PipelineImageInput] = None, + height: int | None = None, + width: int | None = None, + padding_mask_crop: int | None = None, strength: float = 0.9999, num_inference_steps: int = 50, - denoising_start: Optional[float] = None, - denoising_end: Optional[float] = None, + denoising_start: float | None = None, + denoising_end: float | None = None, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, + cross_attention_kwargs: dict[str, Any] | None = None, + controlnet_conditioning_scale: float | list[float] = 1.0, guess_mode: bool = False, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, guidance_rescale: float = 0.0, - original_size: Tuple[int, int] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Tuple[int, int] = None, + original_size: tuple[int, int] = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] = None, aesthetic_score: float = 6.0, negative_aesthetic_score: float = 2.5, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders image (`PIL.Image.Image`): @@ -1275,11 +1270,11 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -1290,7 +1285,7 @@ def __call__( weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -1324,17 +1319,17 @@ def __call__( A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(width, height)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(width, height)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). @@ -1354,7 +1349,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py index 0e2a1441f8f6..89e8d2e54123 100644 --- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py +++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py @@ -14,7 +14,7 @@ import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -120,10 +120,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -138,15 +138,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -212,7 +212,7 @@ class StableDiffusionXLControlNetPipeline( A `CLIPTokenizer` to tokenize text. unet ([`UNet2DConditionModel`]): A `UNet2DConditionModel` to denoise the encoded image latents. - controlnet ([`ControlNetModel`] or `List[ControlNetModel]`): + controlnet ([`ControlNetModel`] or `list[ControlNetModel]`): Provides additional conditioning to the `unet` during the denoising process. If you set multiple ControlNets as a list, the outputs from each ControlNet are added together to create one combined additional conditioning. @@ -257,10 +257,10 @@ def __init__( tokenizer: CLIPTokenizer, tokenizer_2: CLIPTokenizer, unet: UNet2DConditionModel, - controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel], + controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel, scheduler: KarrasDiffusionSchedulers, force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, + add_watermarker: bool | None = None, feature_extractor: CLIPImageProcessor = None, image_encoder: CLIPVisionModelWithProjection = None, ): @@ -299,26 +299,26 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, + prompt_2: str | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders device: (`torch.device`): @@ -327,11 +327,11 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -449,7 +449,7 @@ def encode_prompt( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) - uncond_tokens: List[str] + uncond_tokens: list[str] if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -998,59 +998,57 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_end: Optional[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, + denoising_end: float | None = None, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, + cross_attention_kwargs: dict[str, Any] | None = None, + controlnet_conditioning_scale: float | list[float] = 1.0, guess_mode: bool = False, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, - original_size: Tuple[int, int] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Tuple[int, int] = None, - negative_original_size: Optional[Tuple[int, int]] = None, - negative_crops_coords_top_left: Tuple[int, int] = (0, 0), - negative_target_size: Optional[Tuple[int, int]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, + original_size: tuple[int, int] = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] = None, + negative_original_size: tuple[int, int] | None = None, + negative_crops_coords_top_left: tuple[int, int] = (0, 0), + negative_target_size: tuple[int, int] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: - `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,: + `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`): The ControlNet input condition to provide guidance to the `unet` for generation. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or @@ -1068,11 +1066,11 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -1086,10 +1084,10 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 5.0): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. This is sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders. num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -1097,7 +1095,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -1118,7 +1116,7 @@ def __call__( weighting). If not provided, pooled `negative_prompt_embeds` are generated from `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -1131,42 +1129,42 @@ def __call__( cross_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set the corresponding scale as a list. guess_mode (`bool`, *optional*, defaults to `False`): The ControlNet encoder tries to recognize the content of the input image even if you remove all prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended. - control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the ControlNet starts applying. - control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the ControlNet stops applying. - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a specific image resolution. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a target image resolution. It should be as same as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more @@ -1179,7 +1177,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1344,8 +1342,12 @@ def __call__( assert False # 5. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) self._num_timesteps = len(timesteps) diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py index 94c4c394465b..8c39856ec7cb 100644 --- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py +++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py @@ -14,7 +14,7 @@ import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -84,7 +84,6 @@ >>> from diffusers import ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline, AutoencoderKL >>> from diffusers.utils import load_image - >>> depth_estimator = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to("cuda") >>> feature_extractor = DPTImageProcessor.from_pretrained("Intel/dpt-hybrid-midas") >>> controlnet = ControlNetModel.from_pretrained( @@ -148,7 +147,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -200,7 +199,7 @@ class StableDiffusionXLControlNetImg2ImgPipeline( Second Tokenizer of class [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents. - controlnet ([`ControlNetModel`] or `List[ControlNetModel]`): + controlnet ([`ControlNetModel`] or `list[ControlNetModel]`): Provides additional conditioning to the unet during the denoising process. If you set multiple ControlNets as a list, the outputs from each ControlNet are added together to create one combined additional conditioning. @@ -249,11 +248,11 @@ def __init__( tokenizer: CLIPTokenizer, tokenizer_2: CLIPTokenizer, unet: UNet2DConditionModel, - controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel], + controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel, scheduler: KarrasDiffusionSchedulers, requires_aesthetics_score: bool = False, force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, + add_watermarker: bool | None = None, feature_extractor: CLIPImageProcessor = None, image_encoder: CLIPVisionModelWithProjection = None, ): @@ -293,26 +292,26 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, + prompt_2: str | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders device: (`torch.device`): @@ -321,11 +320,11 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -443,7 +442,7 @@ def encode_prompt( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) - uncond_tokens: List[str] + uncond_tokens: list[str] if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -1078,65 +1077,63 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, image: PipelineImageInput = None, control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, strength: float = 0.8, num_inference_steps: int = 50, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 0.8, + cross_attention_kwargs: dict[str, Any] | None = None, + controlnet_conditioning_scale: float | list[float] = 0.8, guess_mode: bool = False, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, - original_size: Tuple[int, int] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Tuple[int, int] = None, - negative_original_size: Optional[Tuple[int, int]] = None, - negative_crops_coords_top_left: Tuple[int, int] = (0, 0), - negative_target_size: Optional[Tuple[int, int]] = None, + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, + original_size: tuple[int, int] = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] = None, + negative_original_size: tuple[int, int] | None = None, + negative_crops_coords_top_left: tuple[int, int] = (0, 0), + negative_target_size: tuple[int, int] | None = None, aesthetic_score: float = 6.0, negative_aesthetic_score: float = 2.5, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: - `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,: + `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`): The initial image will be used as the starting point for the image generation process. Can also accept image latents as `image`, if passing latents directly, it will not be encoded again. - control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: - `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,: + `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`): The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height @@ -1166,11 +1163,11 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -1178,7 +1175,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -1200,7 +1197,7 @@ def __call__( weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -1215,42 +1212,42 @@ def __call__( A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original unet. If multiple ControlNets are specified in init, you can set the corresponding scale as a list. guess_mode (`bool`, *optional*, defaults to `False`): In this mode, the ControlNet encoder will try best to recognize the content of the input image even if you remove all prompts. The `guidance_scale` between 3.0 and 5.0 is recommended. - control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the controlnet starts applying. - control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the controlnet stops applying. - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a specific image resolution. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a target image resolution. It should be as same as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more @@ -1271,7 +1268,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py index e234015f8616..4b7ca284d636 100644 --- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py +++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -76,7 +76,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -231,15 +231,16 @@ def __init__( tokenizer: CLIPTokenizer, tokenizer_2: CLIPTokenizer, unet: UNet2DConditionModel, - controlnet: Union[ - ControlNetUnionModel, List[ControlNetUnionModel], Tuple[ControlNetUnionModel], MultiControlNetUnionModel - ], + controlnet: ControlNetUnionModel + | list[ControlNetUnionModel] + | tuple[ControlNetUnionModel] + | MultiControlNetUnionModel, scheduler: KarrasDiffusionSchedulers, requires_aesthetics_score: bool = False, force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, - feature_extractor: Optional[CLIPImageProcessor] = None, - image_encoder: Optional[CLIPVisionModelWithProjection] = None, + add_watermarker: bool | None = None, + feature_extractor: CLIPImageProcessor | None = None, + image_encoder: CLIPVisionModelWithProjection | None = None, ): super().__init__() @@ -280,26 +281,26 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, + prompt_2: str | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders device: (`torch.device`): @@ -308,11 +309,11 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -430,7 +431,7 @@ def encode_prompt( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) - uncond_tokens: List[str] + uncond_tokens: list[str] if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -1145,60 +1146,58 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, image: PipelineImageInput = None, mask_image: PipelineImageInput = None, - control_image: Union[PipelineImageInput, List[PipelineImageInput]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - padding_mask_crop: Optional[int] = None, + control_image: PipelineImageInput | list[PipelineImageInput] = None, + height: int | None = None, + width: int | None = None, + padding_mask_crop: int | None = None, strength: float = 0.9999, num_inference_steps: int = 50, - denoising_start: Optional[float] = None, - denoising_end: Optional[float] = None, + denoising_start: float | None = None, + denoising_end: float | None = None, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, + cross_attention_kwargs: dict[str, Any] | None = None, + controlnet_conditioning_scale: float | list[float] = 1.0, guess_mode: bool = False, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, - control_mode: Optional[Union[int, List[int], List[List[int]]]] = None, + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, + control_mode: int | list[int] | list[list[int]] | None = None, guidance_rescale: float = 0.0, - original_size: Tuple[int, int] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Tuple[int, int] = None, + original_size: tuple[int, int] = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] = None, aesthetic_score: float = 6.0, negative_aesthetic_score: float = 2.5, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders image (`PIL.Image.Image`): @@ -1209,7 +1208,7 @@ def __call__( repainted, while black pixels will be preserved. If `mask_image` is a PIL image, it will be converted to a single channel (luminance) before use. If it's a tensor, it should contain one color channel (L) instead of 3, so the expected shape would be `(B, H, W, 1)`. - control_image (`PipelineImageInput` or `List[PipelineImageInput]`, *optional*): + control_image (`PipelineImageInput` or `list[PipelineImageInput]`, *optional*): The ControlNet input condition to provide guidance to the `unet` for generation. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or @@ -1259,11 +1258,11 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -1274,7 +1273,7 @@ def __call__( weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -1308,33 +1307,33 @@ def __call__( A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set the corresponding scale as a list. guess_mode (`bool`, *optional*, defaults to `False`): The ControlNet encoder tries to recognize the content of the input image even if you remove all prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended. - control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the ControlNet starts applying. - control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the ControlNet stops applying. - control_mode (`int` or `List[int]` or `List[List[int]], *optional*): + control_mode (`int` or `list[int]` or `list[list[int]], *optional*): The control condition types for the ControlNet. See the ControlNet's model card forinformation on the available control modes. If multiple ControlNets are specified in `init`, control_mode should be a list where each ControlNet should have its corresponding control mode list. Should reflect the order of conditions in control_image. - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(width, height)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(width, height)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). @@ -1354,7 +1353,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl.py index 40cc76cf70d8..87057c2392df 100644 --- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl.py +++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl.py @@ -14,7 +14,7 @@ import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -115,10 +115,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -133,15 +133,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -246,12 +246,13 @@ def __init__( tokenizer: CLIPTokenizer, tokenizer_2: CLIPTokenizer, unet: UNet2DConditionModel, - controlnet: Union[ - ControlNetUnionModel, List[ControlNetUnionModel], Tuple[ControlNetUnionModel], MultiControlNetUnionModel - ], + controlnet: ControlNetUnionModel + | list[ControlNetUnionModel] + | tuple[ControlNetUnionModel] + | MultiControlNetUnionModel, scheduler: KarrasDiffusionSchedulers, force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, + add_watermarker: bool | None = None, feature_extractor: CLIPImageProcessor = None, image_encoder: CLIPVisionModelWithProjection = None, ): @@ -290,26 +291,26 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, + prompt_2: str | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders device: (`torch.device`): @@ -318,11 +319,11 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -440,7 +441,7 @@ def encode_prompt( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) - uncond_tokens: List[str] + uncond_tokens: list[str] if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -974,58 +975,56 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - control_image: Union[PipelineImageInput, List[PipelineImageInput]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + control_image: PipelineImageInput | list[PipelineImageInput] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_end: Optional[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, + denoising_end: float | None = None, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, + cross_attention_kwargs: dict[str, Any] | None = None, + controlnet_conditioning_scale: float | list[float] = 1.0, guess_mode: bool = False, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, - control_mode: Optional[Union[int, List[int], List[List[int]]]] = None, - original_size: Tuple[int, int] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Tuple[int, int] = None, - negative_original_size: Optional[Tuple[int, int]] = None, - negative_crops_coords_top_left: Tuple[int, int] = (0, 0), - negative_target_size: Optional[Tuple[int, int]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, + control_mode: int | list[int] | list[list[int]] | None = None, + original_size: tuple[int, int] = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] = None, + negative_original_size: tuple[int, int] | None = None, + negative_crops_coords_top_left: tuple[int, int] = (0, 0), + negative_target_size: tuple[int, int] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders. - control_image (`PipelineImageInput` or `List[PipelineImageInput]`, *optional*): + control_image (`PipelineImageInput` or `list[PipelineImageInput]`, *optional*): The ControlNet input condition to provide guidance to the `unet` for generation. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or @@ -1043,11 +1042,11 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -1061,10 +1060,10 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 5.0): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. This is sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders. num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -1072,7 +1071,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -1093,7 +1092,7 @@ def __call__( weighting). If not provided, pooled `negative_prompt_embeds` are generated from `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -1106,47 +1105,47 @@ def __call__( cross_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set the corresponding scale as a list. guess_mode (`bool`, *optional*, defaults to `False`): The ControlNet encoder tries to recognize the content of the input image even if you remove all prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended. - control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the ControlNet starts applying. - control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the ControlNet stops applying. - control_mode (`int` or `List[int]` or `List[List[int]], *optional*): + control_mode (`int` or `list[int]` or `list[list[int]], *optional*): The control condition types for the ControlNet. See the ControlNet's model card forinformation on the available control modes. If multiple ControlNets are specified in `init`, control_mode should be a list where each ControlNet should have its corresponding control mode list. Should reflect the order of conditions in control_image. - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a specific image resolution. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a target image resolution. It should be as same as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more @@ -1159,7 +1158,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1339,8 +1338,12 @@ def __call__( height, width = control_image[0][0].shape[-2:] # 5. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) self._num_timesteps = len(timesteps) diff --git a/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl_img2img.py b/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl_img2img.py index 4d0093132b9c..2fb87e57ea8a 100644 --- a/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl_img2img.py +++ b/src/diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl_img2img.py @@ -14,7 +14,7 @@ import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -173,7 +173,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -263,13 +263,14 @@ def __init__( tokenizer: CLIPTokenizer, tokenizer_2: CLIPTokenizer, unet: UNet2DConditionModel, - controlnet: Union[ - ControlNetUnionModel, List[ControlNetUnionModel], Tuple[ControlNetUnionModel], MultiControlNetUnionModel - ], + controlnet: ControlNetUnionModel + | list[ControlNetUnionModel] + | tuple[ControlNetUnionModel] + | MultiControlNetUnionModel, scheduler: KarrasDiffusionSchedulers, requires_aesthetics_score: bool = False, force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, + add_watermarker: bool | None = None, feature_extractor: CLIPImageProcessor = None, image_encoder: CLIPVisionModelWithProjection = None, ): @@ -309,26 +310,26 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, + prompt_2: str | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders device: (`torch.device`): @@ -337,11 +338,11 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -459,7 +460,7 @@ def encode_prompt( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) - uncond_tokens: List[str] + uncond_tokens: list[str] if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -1066,65 +1067,63 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, image: PipelineImageInput = None, - control_image: Union[PipelineImageInput, List[PipelineImageInput]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + control_image: PipelineImageInput | list[PipelineImageInput] = None, + height: int | None = None, + width: int | None = None, strength: float = 0.8, num_inference_steps: int = 50, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 0.8, + cross_attention_kwargs: dict[str, Any] | None = None, + controlnet_conditioning_scale: float | list[float] = 0.8, guess_mode: bool = False, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, - control_mode: Optional[Union[int, List[int], List[List[int]]]] = None, - original_size: Tuple[int, int] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Tuple[int, int] = None, - negative_original_size: Optional[Tuple[int, int]] = None, - negative_crops_coords_top_left: Tuple[int, int] = (0, 0), - negative_target_size: Optional[Tuple[int, int]] = None, + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, + control_mode: int | list[int] | list[list[int]] | None = None, + original_size: tuple[int, int] = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] = None, + negative_original_size: tuple[int, int] | None = None, + negative_crops_coords_top_left: tuple[int, int] = (0, 0), + negative_target_size: tuple[int, int] | None = None, aesthetic_score: float = 6.0, negative_aesthetic_score: float = 2.5, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: - `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,: + `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`): The initial image will be used as the starting point for the image generation process. Can also accept image latents as `image`, if passing latents directly, it will not be encoded again. - control_image (`PipelineImageInput` or `List[PipelineImageInput]`, *optional*): + control_image (`PipelineImageInput` or `list[PipelineImageInput]`, *optional*): The ControlNet input condition to provide guidance to the `unet` for generation. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or @@ -1154,11 +1153,11 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -1166,7 +1165,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -1188,7 +1187,7 @@ def __call__( weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -1203,47 +1202,47 @@ def __call__( A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set the corresponding scale as a list. guess_mode (`bool`, *optional*, defaults to `False`): In this mode, the ControlNet encoder will try best to recognize the content of the input image even if you remove all prompts. The `guidance_scale` between 3.0 and 5.0 is recommended. - control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the ControlNet starts applying. - control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the ControlNet stops applying. - control_mode (`int` or `List[int]` or `List[List[int]], *optional*): + control_mode (`int` or `list[int]` or `list[list[int]], *optional*): The control condition types for the ControlNet. See the ControlNet's model card forinformation on the available control modes. If multiple ControlNets are specified in `init`, control_mode should be a list where each ControlNet should have its corresponding control mode list. Should reflect the order of conditions in control_image - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a specific image resolution. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a target image resolution. It should be as same as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more @@ -1264,7 +1263,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py b/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py index d4c6f336dfef..a26b9068afd1 100644 --- a/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py +++ b/src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py @@ -14,7 +14,6 @@ import warnings from functools import partial -from typing import Dict, List, Optional, Union import jax import jax.numpy as jnp @@ -148,9 +147,7 @@ def __init__( tokenizer: CLIPTokenizer, unet: FlaxUNet2DConditionModel, controlnet: FlaxControlNetModel, - scheduler: Union[ - FlaxDDIMScheduler, FlaxPNDMScheduler, FlaxLMSDiscreteScheduler, FlaxDPMSolverMultistepScheduler - ], + scheduler: FlaxDDIMScheduler | FlaxPNDMScheduler | FlaxLMSDiscreteScheduler | FlaxDPMSolverMultistepScheduler, safety_checker: FlaxStableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, dtype: jnp.dtype = jnp.float32, @@ -180,7 +177,7 @@ def __init__( ) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - def prepare_text_inputs(self, prompt: Union[str, List[str]]): + def prepare_text_inputs(self, prompt: str | list[str]): if not isinstance(prompt, (str, list)): raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") @@ -194,7 +191,7 @@ def prepare_text_inputs(self, prompt: Union[str, List[str]]): return text_input.input_ids - def prepare_image_inputs(self, image: Union[Image.Image, List[Image.Image]]): + def prepare_image_inputs(self, image: Image.Image | list[Image.Image]): if not isinstance(image, (Image.Image, list)): raise ValueError(f"image has to be of type `PIL.Image.Image` or list but is {type(image)}") @@ -243,12 +240,12 @@ def _generate( self, prompt_ids: jnp.ndarray, image: jnp.ndarray, - params: Union[Dict, FrozenDict], + params: dict | FrozenDict, prng_seed: jax.Array, num_inference_steps: int, guidance_scale: float, - latents: Optional[jnp.ndarray] = None, - neg_prompt_ids: Optional[jnp.ndarray] = None, + latents: jnp.ndarray | None = None, + neg_prompt_ids: jnp.ndarray | None = None, controlnet_conditioning_scale: float = 1.0, ): height, width = image.shape[-2:] @@ -353,13 +350,13 @@ def __call__( self, prompt_ids: jnp.ndarray, image: jnp.ndarray, - params: Union[Dict, FrozenDict], + params: dict | FrozenDict, prng_seed: jax.Array, num_inference_steps: int = 50, - guidance_scale: Union[float, jnp.ndarray] = 7.5, + guidance_scale: float | jnp.ndarray = 7.5, latents: jnp.ndarray = None, neg_prompt_ids: jnp.ndarray = None, - controlnet_conditioning_scale: Union[float, jnp.ndarray] = 1.0, + controlnet_conditioning_scale: float | jnp.ndarray = 1.0, return_dict: bool = True, jit: bool = False, ): diff --git a/src/diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py b/src/diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py index 29a7d6147638..8882b561f0a1 100644 --- a/src/diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py +++ b/src/diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Callable, Dict, List, Optional, Tuple, Union +from typing import Callable import numpy as np import torch @@ -176,10 +176,10 @@ class HunyuanDiTControlNetPipeline(DiffusionPipeline): vae ([`AutoencoderKL`]): Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. We use `sdxl-vae-fp16-fix`. - text_encoder (Optional[`~transformers.BertModel`, `~transformers.CLIPTextModel`]): + text_encoder (`~transformers.BertModel`, `~transformers.CLIPTextModel` | None): Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). HunyuanDiT uses a fine-tuned [bilingual CLIP]. - tokenizer (Optional[`~transformers.BertTokenizer`, `~transformers.CLIPTokenizer`]): + tokenizer (`~transformers.BertTokenizer`, `~transformers.CLIPTokenizer` | None): A `BertTokenizer` or `CLIPTokenizer` to tokenize text. transformer ([`HunyuanDiT2DModel`]): The HunyuanDiT model designed by Tencent Hunyuan. @@ -189,7 +189,7 @@ class HunyuanDiTControlNetPipeline(DiffusionPipeline): The tokenizer for the mT5 embedder. scheduler ([`DDPMScheduler`]): A scheduler to be used in combination with HunyuanDiT to denoise the encoded image latents. - controlnet ([`HunyuanDiT2DControlNetModel`] or `List[HunyuanDiT2DControlNetModel]` or [`HunyuanDiT2DControlNetModel`]): + controlnet ([`HunyuanDiT2DControlNetModel`] or `list[HunyuanDiT2DControlNetModel]` or [`HunyuanDiT2DControlNetModel`]): Provides additional conditioning to the `unet` during the denoising process. If you set multiple ControlNets as a list, the outputs from each ControlNet are added together to create one combined additional conditioning. @@ -222,14 +222,12 @@ def __init__( scheduler: DDPMScheduler, safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, - controlnet: Union[ - HunyuanDiT2DControlNetModel, - List[HunyuanDiT2DControlNetModel], - Tuple[HunyuanDiT2DControlNetModel], - HunyuanDiT2DMultiControlNetModel, - ], - text_encoder_2: Optional[T5EncoderModel] = None, - tokenizer_2: Optional[T5Tokenizer] = None, + controlnet: HunyuanDiT2DControlNetModel + | list[HunyuanDiT2DControlNetModel] + | tuple[HunyuanDiT2DControlNetModel] + | HunyuanDiT2DMultiControlNetModel, + text_encoder_2: T5EncoderModel | None = None, + tokenizer_2: T5Tokenizer | None = None, requires_safety_checker: bool = True, ): super().__init__() @@ -282,19 +280,19 @@ def encode_prompt( dtype: torch.dtype = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - max_sequence_length: Optional[int] = None, + negative_prompt: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + max_sequence_length: int | None = None, text_encoder_index: int = 0, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -304,7 +302,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -393,7 +391,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -635,43 +633,41 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: Optional[int] = 50, - guidance_scale: Optional[float] = 5.0, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, + num_inference_steps: int | None = 50, + guidance_scale: float | None = 5.0, control_image: PipelineImageInput = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - eta: Optional[float] = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_2: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_2: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - prompt_attention_mask_2: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask_2: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + controlnet_conditioning_scale: float | list[float] = 1.0, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + eta: float | None = 0.0, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_2: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_2: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + prompt_attention_mask_2: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask_2: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], guidance_rescale: float = 0.0, - original_size: Optional[Tuple[int, int]] = (1024, 1024), - target_size: Optional[Tuple[int, int]] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), + original_size: tuple[int, int] | None = (1024, 1024), + target_size: tuple[int, int] | None = None, + crops_coords_top_left: tuple[int, int] = (0, 0), use_resolution_binning: bool = True, ): r""" The call function to the pipeline for generation with HunyuanDiT. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`): The height in pixels of the generated image. @@ -683,23 +679,23 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the ControlNet starts applying. - control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the ControlNet stops applying. - control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: - `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,: + `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`): The ControlNet input condition to provide guidance to the `unet` for generation. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or width are passed, `image` is resized accordingly. If multiple ControlNets are specified in `init`, images must be passed as a list such that each element of the list can be correctly batched for input to a single ControlNet. - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set the corresponding scale as a list. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -707,7 +703,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. prompt_embeds (`torch.Tensor`, *optional*): @@ -735,19 +731,19 @@ def __call__( return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a plain tuple. - callback_on_step_end (`Callable[[int, int, Dict], None]`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): + callback_on_step_end (`Callable[[int, int], None]`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): A callback function or a list of callback functions to be called at the end of each denoising step. - callback_on_step_end_tensor_inputs (`List[str]`, *optional*): + callback_on_step_end_tensor_inputs (`list[str]`, *optional*): A list of tensor inputs that should be passed to the callback function. If not defined, all tensor inputs will be passed. guidance_rescale (`float`, *optional*, defaults to 0.0): Rescale the noise_cfg according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). See Section 3.4 - original_size (`Tuple[int, int]`, *optional*, defaults to `(1024, 1024)`): + original_size (`tuple[int, int]`, *optional*, defaults to `(1024, 1024)`): The original size of the image. Used to calculate the time ids. - target_size (`Tuple[int, int]`, *optional*): + target_size (`tuple[int, int]`, *optional*): The target size of the image. Used to calculate the time ids. - crops_coords_top_left (`Tuple[int, int]`, *optional*, defaults to `(0, 0)`): + crops_coords_top_left (`tuple[int, int]`, *optional*, defaults to `(0, 0)`): The top left coordinates of the crop. Used to calculate the time ids. use_resolution_binning (`bool`, *optional*, defaults to `True`): Whether to use resolution binning or not. If `True`, the input resolution will be mapped to the closest diff --git a/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py b/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py index d605eac1f2b1..a787a34bdc01 100644 --- a/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py +++ b/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from transformers import ( @@ -83,10 +83,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -101,15 +101,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -174,7 +174,7 @@ class StableDiffusion3ControlNetPipeline( tokenizer_3 (`T5TokenizerFast`): Tokenizer of class [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - controlnet ([`SD3ControlNetModel`] or `List[SD3ControlNetModel]` or [`SD3MultiControlNetModel`]): + controlnet ([`SD3ControlNetModel`] or `list[SD3ControlNetModel]` or [`SD3MultiControlNetModel`]): Provides additional conditioning to the `unet` during the denoising process. If you set multiple ControlNets as a list, the outputs from each ControlNet are added together to create one combined additional conditioning. @@ -199,11 +199,12 @@ def __init__( tokenizer_2: CLIPTokenizer, text_encoder_3: T5EncoderModel, tokenizer_3: T5TokenizerFast, - controlnet: Union[ - SD3ControlNetModel, List[SD3ControlNetModel], Tuple[SD3ControlNetModel], SD3MultiControlNetModel - ], - image_encoder: Optional[SiglipVisionModel] = None, - feature_extractor: Optional[SiglipImageProcessor] = None, + controlnet: SD3ControlNetModel + | list[SD3ControlNetModel] + | tuple[SD3ControlNetModel] + | SD3MultiControlNetModel, + image_encoder: SiglipVisionModel | None = None, + feature_extractor: SiglipImageProcessor | None = None, ): super().__init__() if isinstance(controlnet, (list, tuple)): @@ -250,11 +251,11 @@ def __init__( # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 256, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -307,10 +308,10 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - clip_skip: Optional[int] = None, + device: torch.device | None = None, + clip_skip: int | None = None, clip_model_index: int = 0, ): device = device or self._execution_device @@ -363,32 +364,32 @@ def _get_clip_prompt_embeds( # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], - prompt_3: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + prompt_2: str | list[str], + prompt_3: str | list[str], + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt_3: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_3: str | list[str] | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + clip_skip: int | None = None, max_sequence_length: int = 256, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in all text-encoders - prompt_3 (`str` or `List[str]`, *optional*): + prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is used in all text-encoders device: (`torch.device`): @@ -397,14 +398,14 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders. - negative_prompt_3 (`str` or `List[str]`, *optional*): + negative_prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `negative_prompt` is used in all the text-encoders. prompt_embeds (`torch.FloatTensor`, *optional*): @@ -760,9 +761,9 @@ def encode_image(self, image: PipelineImageInput, device: torch.device) -> torch # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.prepare_ip_adapter_image_embeds def prepare_ip_adapter_image_embeds( self, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: torch.Tensor | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, ) -> torch.Tensor: @@ -819,50 +820,50 @@ def enable_sequential_cpu_offload(self, *args, **kwargs): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - prompt_3: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + prompt_3: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 28, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 7.0, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, control_image: PipelineImageInput = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, - controlnet_pooled_projections: Optional[torch.FloatTensor] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt_3: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + controlnet_conditioning_scale: float | list[float] = 1.0, + controlnet_pooled_projections: torch.FloatTensor | None = None, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_3: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 256, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead - prompt_3 (`str` or `List[str]`, *optional*): + prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is will be used instead height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): @@ -872,7 +873,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -882,37 +883,37 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the ControlNet starts applying. - control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the ControlNet stops applying. - control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: - `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,: + `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`): The ControlNet input condition to provide guidance to the `unet` for generation. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or width are passed, `image` is resized accordingly. If multiple ControlNets are specified in `init`, images must be passed as a list such that each element of the list can be correctly batched for input to a single ControlNet. - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set the corresponding scale as a list. controlnet_pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected from the embeddings of controlnet input conditions. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used instead - negative_prompt_3 (`str` or `List[str]`, *optional*): + negative_prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `negative_prompt` is used instead num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -954,7 +955,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1098,7 +1099,13 @@ def __call__( assert False # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, sigmas=sigmas) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, sigmas=sigmas + ) num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) self._num_timesteps = len(timesteps) diff --git a/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py b/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py index 9d0158c6b654..96f53b16cbe8 100644 --- a/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py +++ b/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py @@ -13,8 +13,10 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable +import numpy as np +import PIL.Image import torch from transformers import ( CLIPTextModelWithProjection, @@ -39,7 +41,7 @@ scale_lora_layers, unscale_lora_layers, ) -from ...utils.torch_utils import randn_tensor +from ...utils.torch_utils import is_compiled_module, randn_tensor from ..pipeline_utils import DiffusionPipeline from ..stable_diffusion_3.pipeline_output import StableDiffusion3PipelineOutput @@ -104,10 +106,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -122,15 +124,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -195,7 +197,7 @@ class StableDiffusion3ControlNetInpaintingPipeline( tokenizer_3 (`T5TokenizerFast`): Tokenizer of class [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). - controlnet ([`SD3ControlNetModel`] or `List[SD3ControlNetModel]` or [`SD3MultiControlNetModel`]): + controlnet ([`SD3ControlNetModel`] or `list[SD3ControlNetModel]` or [`SD3MultiControlNetModel`]): Provides additional conditioning to the `transformer` during the denoising process. If you set multiple ControlNets as a list, the outputs from each ControlNet are added together to create one combined additional conditioning. @@ -220,13 +222,16 @@ def __init__( tokenizer_2: CLIPTokenizer, text_encoder_3: T5EncoderModel, tokenizer_3: T5TokenizerFast, - controlnet: Union[ - SD3ControlNetModel, List[SD3ControlNetModel], Tuple[SD3ControlNetModel], SD3MultiControlNetModel - ], + controlnet: SD3ControlNetModel + | list[SD3ControlNetModel] + | tuple[SD3ControlNetModel] + | SD3MultiControlNetModel, image_encoder: SiglipModel = None, - feature_extractor: Optional[SiglipImageProcessor] = None, + feature_extractor: SiglipImageProcessor | None = None, ): super().__init__() + if isinstance(controlnet, (list, tuple)): + controlnet = SD3MultiControlNetModel(controlnet) self.register_modules( vae=vae, @@ -268,11 +273,11 @@ def __init__( # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 256, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -325,10 +330,10 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - clip_skip: Optional[int] = None, + device: torch.device | None = None, + clip_skip: int | None = None, clip_model_index: int = 0, ): device = device or self._execution_device @@ -381,32 +386,32 @@ def _get_clip_prompt_embeds( # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], - prompt_3: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + prompt_2: str | list[str], + prompt_3: str | list[str], + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt_3: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_3: str | list[str] | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + clip_skip: int | None = None, max_sequence_length: int = 256, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in all text-encoders - prompt_3 (`str` or `List[str]`, *optional*): + prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is used in all text-encoders device: (`torch.device`): @@ -415,14 +420,14 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders. - negative_prompt_3 (`str` or `List[str]`, *optional*): + negative_prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `negative_prompt` is used in all the text-encoders. prompt_embeds (`torch.FloatTensor`, *optional*): @@ -572,14 +577,52 @@ def encode_prompt( return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.check_inputs + # Copied from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl.StableDiffusionXLControlNetPipeline.check_image + def check_image(self, image, prompt, prompt_embeds): + image_is_pil = isinstance(image, PIL.Image.Image) + image_is_tensor = isinstance(image, torch.Tensor) + image_is_np = isinstance(image, np.ndarray) + image_is_pil_list = isinstance(image, list) and isinstance(image[0], PIL.Image.Image) + image_is_tensor_list = isinstance(image, list) and isinstance(image[0], torch.Tensor) + image_is_np_list = isinstance(image, list) and isinstance(image[0], np.ndarray) + + if ( + not image_is_pil + and not image_is_tensor + and not image_is_np + and not image_is_pil_list + and not image_is_tensor_list + and not image_is_np_list + ): + raise TypeError( + f"image must be passed and be one of PIL image, numpy array, torch tensor, list of PIL images, list of numpy arrays or list of torch tensors, but is {type(image)}" + ) + + if image_is_pil: + image_batch_size = 1 + else: + image_batch_size = len(image) + + if prompt is not None and isinstance(prompt, str): + prompt_batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + prompt_batch_size = len(prompt) + elif prompt_embeds is not None: + prompt_batch_size = prompt_embeds.shape[0] + + if image_batch_size != 1 and image_batch_size != prompt_batch_size: + raise ValueError( + f"If image batch size is not 1, image batch size must be same as prompt batch size. image batch size: {image_batch_size}, prompt batch size: {prompt_batch_size}" + ) + def check_inputs( self, + height, + width, + image, prompt, prompt_2, prompt_3, - height, - width, negative_prompt=None, negative_prompt_2=None, negative_prompt_3=None, @@ -587,6 +630,11 @@ def check_inputs( negative_prompt_embeds=None, pooled_prompt_embeds=None, negative_pooled_prompt_embeds=None, + ip_adapter_image=None, + ip_adapter_image_embeds=None, + controlnet_conditioning_scale=1.0, + control_guidance_start=0.0, + control_guidance_end=1.0, callback_on_step_end_tensor_inputs=None, max_sequence_length=None, ): @@ -669,6 +717,76 @@ def check_inputs( if max_sequence_length is not None and max_sequence_length > 512: raise ValueError(f"`max_sequence_length` cannot be greater than 512 but is {max_sequence_length}") + # `prompt` needs more sophisticated handling when there are multiple + # conditionings. + if isinstance(self.controlnet, SD3MultiControlNetModel): + if isinstance(prompt, list) and len(prompt) > 1: + logger.warning( + f"You have {len(self.controlnet.nets)} ControlNets and you have passed {len(prompt)}" + " prompts. The conditionings will be fixed across the prompts." + ) + + # Check `image` + controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet + + if isinstance(controlnet, SD3ControlNetModel): + self.check_image(image, prompt, prompt_embeds) + elif isinstance(controlnet, SD3MultiControlNetModel): + if not isinstance(image, list): + raise TypeError("For multiple controlnets: `image` must be type `list`") + elif len(image) != len(self.controlnet.nets): + raise ValueError( + f"For multiple controlnets: `image` must have the same length as the number of controlnets, but got {len(image)} images and {len(self.controlnet.nets)} ControlNets." + ) + for image_ in image: + self.check_image(image_, prompt, prompt_embeds) + + # Check `controlnet_conditioning_scale` + if isinstance(controlnet, SD3MultiControlNetModel): + if isinstance(controlnet_conditioning_scale, list) and len(controlnet_conditioning_scale) != len( + self.controlnet.nets + ): + raise ValueError( + "For multiple controlnets: When `controlnet_conditioning_scale` is specified as `list`, it must have" + " the same length as the number of controlnets" + ) + + if len(control_guidance_start) != len(control_guidance_end): + raise ValueError( + f"`control_guidance_start` has {len(control_guidance_start)} elements, but `control_guidance_end` has {len(control_guidance_end)} elements. Make sure to provide the same number of elements to each list." + ) + + if isinstance(controlnet, SD3MultiControlNetModel): + if len(control_guidance_start) != len(self.controlnet.nets): + raise ValueError( + f"`control_guidance_start`: {control_guidance_start} has {len(control_guidance_start)} elements but there are {len(self.controlnet.nets)} controlnets available. Make sure to provide {len(self.controlnet.nets)}." + ) + + for start, end in zip(control_guidance_start, control_guidance_end): + if start >= end: + raise ValueError( + f"control_guidance_start: {start} cannot be larger or equal to control guidance end: {end}." + ) + if start < 0.0: + raise ValueError(f"control_guidance_start: {start} can't be smaller than 0.") + if end > 1.0: + raise ValueError(f"control_guidance_end: {end} can't be larger than 1.0.") + + if ip_adapter_image is not None and ip_adapter_image_embeds is not None: + raise ValueError( + "Provide either `ip_adapter_image` or `ip_adapter_image_embeds`. Cannot leave both `ip_adapter_image` and `ip_adapter_image_embeds` defined." + ) + + if ip_adapter_image_embeds is not None: + if not isinstance(ip_adapter_image_embeds, list): + raise ValueError( + f"`ip_adapter_image_embeds` has to be of type `list` but is {type(ip_adapter_image_embeds)}" + ) + elif ip_adapter_image_embeds[0].ndim not in [3, 4]: + raise ValueError( + f"`ip_adapter_image_embeds` has to be a list of 3D or 4D tensors but is {ip_adapter_image_embeds[0].ndim}D" + ) + # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.prepare_latents def prepare_latents( self, @@ -810,9 +928,9 @@ def encode_image(self, image: PipelineImageInput, device: torch.device) -> torch # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.prepare_ip_adapter_image_embeds def prepare_ip_adapter_image_embeds( self, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: torch.Tensor | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, ) -> torch.Tensor: @@ -869,51 +987,51 @@ def enable_sequential_cpu_offload(self, *args, **kwargs): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - prompt_3: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + prompt_3: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 28, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 7.0, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, control_image: PipelineImageInput = None, control_mask: PipelineImageInput = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, - controlnet_pooled_projections: Optional[torch.FloatTensor] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt_3: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + controlnet_conditioning_scale: float | list[float] = 1.0, + controlnet_pooled_projections: torch.FloatTensor | None = None, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_3: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 256, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead - prompt_3 (`str` or `List[str]`, *optional*): + prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is will be used instead height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): @@ -923,7 +1041,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -933,41 +1051,41 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the ControlNet starts applying. - control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the ControlNet stops applying. - control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`): + control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be inpainted (which parts of the image to be masked out with `control_mask` and repainted according to `prompt`). For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)`. If it is a numpy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)`. - control_mask (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`): + control_mask (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one color channel (L) instead of 3, so the expected shape for pytorch tensor would be `(B, 1, H, W)`. And for numpy array would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W, 1)`, or `(H, W)`. - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set the corresponding scale as a list. controlnet_pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected from the embeddings of controlnet input conditions. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used instead - negative_prompt_3 (`str` or `List[str]`, *optional*): + negative_prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `negative_prompt` is used instead num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -1009,7 +1127,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1040,11 +1158,12 @@ def __call__( # 1. Check inputs. Raise error if not correct self.check_inputs( + height, + width, + control_image, prompt, prompt_2, prompt_3, - height, - width, negative_prompt=negative_prompt, negative_prompt_2=negative_prompt_2, negative_prompt_3=negative_prompt_3, @@ -1052,6 +1171,11 @@ def __call__( negative_prompt_embeds=negative_prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + ip_adapter_image=ip_adapter_image, + ip_adapter_image_embeds=ip_adapter_image_embeds, + controlnet_conditioning_scale=controlnet_conditioning_scale, + control_guidance_start=control_guidance_start, + control_guidance_end=control_guidance_end, callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, max_sequence_length=max_sequence_length, ) @@ -1119,9 +1243,26 @@ def __call__( width = latent_width * self.vae_scale_factor elif isinstance(self.controlnet, SD3MultiControlNetModel): - raise NotImplementedError("MultiControlNetModel is not supported for SD3ControlNetInpaintingPipeline.") + control_images = [] + + for control_image_ in control_image: + control_image_ = self.prepare_image_with_mask( + image=control_image_, + mask=control_mask, + width=width, + height=height, + batch_size=batch_size * num_images_per_prompt, + num_images_per_prompt=num_images_per_prompt, + device=device, + dtype=dtype, + do_classifier_free_guidance=self.do_classifier_free_guidance, + guess_mode=False, + ) + control_images.append(control_image_) + + control_image = control_images else: - assert False + assert ValueError("Controlnet not found. Please check the controlnet model.") if controlnet_pooled_projections is None: controlnet_pooled_projections = torch.zeros_like(pooled_prompt_embeds) @@ -1129,7 +1270,13 @@ def __call__( controlnet_pooled_projections = controlnet_pooled_projections or pooled_prompt_embeds # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, sigmas=sigmas) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, sigmas=sigmas + ) num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) self._num_timesteps = len(timesteps) diff --git a/src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py b/src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py index 3682ddc91156..9c81eb57e6c5 100644 --- a/src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py +++ b/src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -150,7 +150,7 @@ def __init__( vae: AutoencoderKL, text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, - unet: Union[UNet2DConditionModel, UNetControlNetXSModel], + unet: UNet2DConditionModel | UNetControlNetXSModel, controlnet: ControlNetXSAdapter, scheduler: KarrasDiffusionSchedulers, safety_checker: StableDiffusionSafetyChecker, @@ -203,9 +203,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -236,16 +236,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -253,7 +253,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -352,7 +352,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -644,39 +644,37 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, + cross_attention_kwargs: dict[str, Any] | None = None, + controlnet_conditioning_scale: float | list[float] = 1.0, control_guidance_start: float = 0.0, control_guidance_end: float = 1.0, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: - `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,: + `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`): The ControlNet input condition to provide guidance to the `unet` for generation. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or @@ -693,7 +691,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -701,7 +699,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -722,13 +720,13 @@ def __call__( cross_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set the corresponding scale as a list. - control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the ControlNet starts applying. - control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the ControlNet stops applying. clip_skip (`int`, *optional*): Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that @@ -738,7 +736,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs_sd_xl.py b/src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs_sd_xl.py index 7bf610f3a0ba..e7a5862b5b7a 100644 --- a/src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs_sd_xl.py +++ b/src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs_sd_xl.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -182,11 +182,11 @@ def __init__( text_encoder_2: CLIPTextModelWithProjection, tokenizer: CLIPTokenizer, tokenizer_2: CLIPTokenizer, - unet: Union[UNet2DConditionModel, UNetControlNetXSModel], + unet: UNet2DConditionModel | UNetControlNetXSModel, controlnet: ControlNetXSAdapter, scheduler: KarrasDiffusionSchedulers, force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, + add_watermarker: bool | None = None, feature_extractor: CLIPImageProcessor = None, ): super().__init__() @@ -223,26 +223,26 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, + prompt_2: str | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders device: (`torch.device`): @@ -251,11 +251,11 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -373,7 +373,7 @@ def encode_prompt( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) - uncond_tokens: List[str] + uncond_tokens: list[str] if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -718,52 +718,50 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, + cross_attention_kwargs: dict[str, Any] | None = None, + controlnet_conditioning_scale: float | list[float] = 1.0, control_guidance_start: float = 0.0, control_guidance_end: float = 1.0, - original_size: Tuple[int, int] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Tuple[int, int] = None, - negative_original_size: Optional[Tuple[int, int]] = None, - negative_crops_coords_top_left: Tuple[int, int] = (0, 0), - negative_target_size: Optional[Tuple[int, int]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + original_size: tuple[int, int] = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] = None, + negative_original_size: tuple[int, int] | None = None, + negative_crops_coords_top_left: tuple[int, int] = (0, 0), + negative_target_size: tuple[int, int] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: - `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,: + `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`): The ControlNet input condition to provide guidance to the `unet` for generation. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or @@ -784,10 +782,10 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 5.0): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. This is sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders. num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -795,7 +793,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -823,38 +821,38 @@ def __call__( cross_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. control_guidance_start (`float`, *optional*, defaults to 0.0): The percentage of total steps at which the ControlNet starts applying. control_guidance_end (`float`, *optional*, defaults to 1.0): The percentage of total steps at which the ControlNet stops applying. - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(width, height)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(width, height)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a specific image resolution. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a target image resolution. It should be as same as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more @@ -867,7 +865,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/cosmos/__init__.py b/src/diffusers/pipelines/cosmos/__init__.py index 944f16553173..5fc66cdf84b6 100644 --- a/src/diffusers/pipelines/cosmos/__init__.py +++ b/src/diffusers/pipelines/cosmos/__init__.py @@ -25,6 +25,9 @@ _import_structure["pipeline_cosmos2_5_predict"] = [ "Cosmos2_5_PredictBasePipeline", ] + _import_structure["pipeline_cosmos2_5_transfer"] = [ + "Cosmos2_5_TransferPipeline", + ] _import_structure["pipeline_cosmos2_text2image"] = ["Cosmos2TextToImagePipeline"] _import_structure["pipeline_cosmos2_video2world"] = ["Cosmos2VideoToWorldPipeline"] _import_structure["pipeline_cosmos_text2world"] = ["CosmosTextToWorldPipeline"] @@ -41,6 +44,7 @@ from .pipeline_cosmos2_5_predict import ( Cosmos2_5_PredictBasePipeline, ) + from .pipeline_cosmos2_5_transfer import Cosmos2_5_TransferPipeline from .pipeline_cosmos2_text2image import Cosmos2TextToImagePipeline from .pipeline_cosmos2_video2world import Cosmos2VideoToWorldPipeline from .pipeline_cosmos_text2world import CosmosTextToWorldPipeline diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py index 372684e0b521..cdea71a5ab93 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import numpy as np import torch @@ -52,10 +52,19 @@ def __init__(self, *args, **kwargs): logger = logging.get_logger(__name__) # pylint: disable=invalid-name +DEFAULT_NEGATIVE_PROMPT = ( + "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, " + "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, " + "underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, " + "jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, " + "fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. " + "Overall, the video is of poor quality." +) + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -76,7 +85,7 @@ def retrieve_latents( >>> model_id = "nvidia/Cosmos-Predict2.5-2B" >>> pipe = Cosmos2_5_PredictBasePipeline.from_pretrained( - ... model_id, revision="diffusers/base/pre-trianed", torch_dtype=torch.bfloat16 + ... model_id, revision="diffusers/base/post-trained", torch_dtype=torch.bfloat16 ... ) >>> pipe = pipe.to("cuda") @@ -237,10 +246,10 @@ def __init__( def _get_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -278,6 +287,9 @@ def _get_prompt_embeds( truncation=True, padding="max_length", ) + input_ids = ( + input_ids["input_ids"] if not isinstance(input_ids, list) and "input_ids" in input_ids else input_ids + ) input_ids = torch.LongTensor(input_ids) input_ids_batch.append(input_ids) @@ -304,23 +316,23 @@ def _get_prompt_embeds( # Modified from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -359,7 +371,7 @@ def encode_prompt( prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) if do_classifier_free_guidance and negative_prompt_embeds is None: - negative_prompt = negative_prompt or "" + negative_prompt = negative_prompt if negative_prompt is not None else DEFAULT_NEGATIVE_PROMPT negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt if prompt is not None and type(prompt) is not type(negative_prompt): @@ -389,7 +401,7 @@ def encode_prompt( # diffusers.pipelines.cosmos.pipeline_cosmos2_video2world.Cosmos2TextToImagePipeline.prepare_latents def prepare_latents( self, - video: Optional[torch.Tensor], + video: torch.Tensor | None, batch_size: int, num_channels_latents: int = 16, height: int = 704, @@ -397,10 +409,10 @@ def prepare_latents( num_frames_in: int = 93, num_frames_out: int = 93, do_classifier_free_guidance: bool = True, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( @@ -528,27 +540,26 @@ def interrupt(self): def __call__( self, image: PipelineImageInput | None = None, - video: List[PipelineImageInput] | None = None, - prompt: Union[str, List[str]] | None = None, - negative_prompt: Optional[Union[str, List[str]]] = None, + video: list[PipelineImageInput] | None = None, + prompt: str | list[str] | None = None, + negative_prompt: str | list[str] | None = None, height: int = 704, width: int = 1280, num_frames: int = 93, num_inference_steps: int = 36, guidance_scale: float = 7.0, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int, None], PipelineCallback | MultiPipelineCallbacks] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, conditional_frame_timestep: float = 0.1, + num_latent_conditional_frames: int = 2, ): r""" The call function to the pipeline for generation. Supports three modes: @@ -565,9 +576,9 @@ def __call__( Args: image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, *optional*): Optional single image for Image2World conditioning. Must be `None` when `video` is provided. - video (`List[PIL.Image.Image]`, `np.ndarray`, `torch.Tensor`, *optional*): + video (`list[PIL.Image.Image]`, `np.ndarray`, `torch.Tensor`, *optional*): Optional input video for Video2World conditioning. Must be `None` when `image` is provided. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide generation. Required unless `prompt_embeds` is supplied. height (`int`, defaults to `704`): The height in pixels of the generated image. @@ -585,7 +596,7 @@ def __call__( `guidance_scale > 1`. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -614,6 +625,10 @@ def __call__( max_sequence_length (`int`, defaults to `512`): The maximum number of tokens in the prompt. If the prompt exceeds this length, it will be truncated. If the prompt is shorter than this length, it will be padded. + num_latent_conditional_frames (`int`, defaults to `2`): + Number of latent conditional frames to use for Video2World conditioning. The number of pixel frames + extracted from the input video is calculated as `4 * (num_latent_conditional_frames - 1) + 1`. Set to 1 + for Image2World-like behavior (single frame conditioning). Examples: @@ -692,19 +707,38 @@ def __call__( video = torch.zeros(batch_size, num_frames, 3, height, width, dtype=torch.uint8) num_frames_in = 0 else: - num_frames_in = len(video) - if batch_size != 1: raise ValueError(f"batch_size must be 1 for video input (given {batch_size})") + if num_latent_conditional_frames not in [1, 2]: + raise ValueError( + f"num_latent_conditional_frames must be 1 or 2, but got {num_latent_conditional_frames}" + ) + + frames_to_extract = 4 * (num_latent_conditional_frames - 1) + 1 + + total_input_frames = len(video) + + if total_input_frames < frames_to_extract: + raise ValueError( + f"Input video has only {total_input_frames} frames but Video2World requires at least " + f"{frames_to_extract} frames for conditioning." + ) + + num_frames_in = frames_to_extract + assert video is not None video = self.video_processor.preprocess_video(video, height, width) - # pad with last frame (for video2world) + # For Video2World: extract last frames_to_extract frames from input, then pad + if image is None and num_frames_in > 0 and num_frames_in < video.shape[2]: + video = video[:, :, -num_frames_in:, :, :] + num_frames_out = num_frames + if video.shape[2] < num_frames_out: - n_pad_frames = num_frames_out - num_frames_in - last_frame = video[0, :, -1:, :, :] # [C, T==1, H, W] + n_pad_frames = num_frames_out - video.shape[2] + last_frame = video[:, :, -1:, :, :] # [B, C, T==1, H, W] pad_frames = last_frame.repeat(1, 1, n_pad_frames, 1, 1) # [B, C, T, H, W] video = torch.cat((video, pad_frames), dim=2) diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_transfer.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_transfer.py new file mode 100644 index 000000000000..bbe38c44355e --- /dev/null +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_transfer.py @@ -0,0 +1,1012 @@ +# Copyright 2025 The NVIDIA Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Callable, Dict, List, Optional, Union + +import numpy as np +import PIL.Image +import torch +from transformers import AutoTokenizer, Qwen2_5_VLForConditionalGeneration + +from ...callbacks import MultiPipelineCallbacks, PipelineCallback +from ...image_processor import PipelineImageInput +from ...models import AutoencoderKLWan, CosmosControlNetModel, CosmosTransformer3DModel +from ...schedulers import UniPCMultistepScheduler +from ...utils import is_cosmos_guardrail_available, is_torch_xla_available, logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor +from ..pipeline_utils import DiffusionPipeline +from .pipeline_output import CosmosPipelineOutput + + +if is_cosmos_guardrail_available(): + from cosmos_guardrail import CosmosSafetyChecker +else: + + class CosmosSafetyChecker: + def __init__(self, *args, **kwargs): + raise ImportError( + "`cosmos_guardrail` is not installed. Please install it to use the safety checker for Cosmos: `pip install cosmos_guardrail`." + ) + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def _maybe_pad_or_trim_video(video: torch.Tensor, num_frames: int): + n_pad_frames = num_frames - video.shape[2] + if n_pad_frames > 0: + last_frame = video[:, :, -1:, :, :] + video = torch.cat((video, last_frame.repeat(1, 1, n_pad_frames, 1, 1)), dim=2) + elif num_frames < video.shape[2]: + video = video[:, :, :num_frames, :, :] + return video + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents +def retrieve_latents( + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" +): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +DEFAULT_NEGATIVE_PROMPT = "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of poor quality." + +EXAMPLE_DOC_STRING = """ + Examples: + ```python + >>> import cv2 + >>> import numpy as np + >>> import torch + >>> from diffusers import Cosmos2_5_TransferPipeline, AutoModel + >>> from diffusers.utils import export_to_video, load_video + + >>> model_id = "nvidia/Cosmos-Transfer2.5-2B" + >>> # Load a Transfer2.5 controlnet variant (edge, depth, seg, or blur) + >>> controlnet = AutoModel.from_pretrained(model_id, revision="diffusers/controlnet/general/edge") + >>> pipe = Cosmos2_5_TransferPipeline.from_pretrained( + ... model_id, controlnet=controlnet, revision="diffusers/general", torch_dtype=torch.bfloat16 + ... ) + >>> pipe = pipe.to("cuda") + + >>> # Video2World with edge control: Generate video guided by edge maps extracted from input video. + >>> prompt = ( + ... "The video is a demonstration of robotic manipulation, likely in a laboratory or testing environment. It" + ... "features two robotic arms interacting with a piece of blue fabric. The setting is a room with a beige" + ... "couch in the background, providing a neutral backdrop for the robotic activity. The robotic arms are" + ... "positioned on either side of the fabric, which is placed on a yellow cushion. The left robotic arm is" + ... "white with a black gripper, while the right arm is black with a more complex, articulated gripper. At the" + ... "beginning, the fabric is laid out on the cushion. The left robotic arm approaches the fabric, its gripper" + ... "opening and closing as it positions itself. The right arm remains stationary initially, poised to assist." + ... "As the video progresses, the left arm grips the fabric, lifting it slightly off the cushion. The right arm" + ... "then moves in, its gripper adjusting to grasp the opposite side of the fabric. Both arms work in" + ... "coordination, lifting and holding the fabric between them. The fabric is manipulated with precision," + ... "showcasing the dexterity and control of the robotic arms. The camera remains static throughout, focusing" + ... "on the interaction between the robotic arms and the fabric, allowing viewers to observe the detailed" + ... "movements and coordination involved in the task." + ... ) + >>> negative_prompt = ( + ... "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, " + ... "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, " + ... "underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky " + ... "movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, " + ... "fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. " + ... "Overall, the video is of poor quality." + ... ) + >>> input_video = load_video( + ... "https://github.com/nvidia-cosmos/cosmos-transfer2.5/raw/refs/heads/main/assets/robot_example/robot_input.mp4" + ... ) + >>> num_frames = 93 + + >>> # Extract edge maps from the input video using Canny edge detection + >>> edge_maps = [ + ... cv2.Canny(cv2.cvtColor(np.array(frame.convert("RGB")), cv2.COLOR_RGB2BGR), 100, 200) + ... for frame in input_video[:num_frames] + ... ] + >>> edge_maps = np.stack(edge_maps)[None] # (T, H, W) -> (1, T, H, W) + >>> controls = torch.from_numpy(edge_maps).expand(3, -1, -1, -1) # (1, T, H, W) -> (3, T, H, W) + >>> controls = [Image.fromarray(x.numpy()) for x in controls.permute(1, 2, 3, 0)] + >>> export_to_video(controls, "edge_controlled_video_edge.mp4", fps=30) + + >>> # Transfer inference with controls. + >>> video = pipe( + ... controls=controls, + ... controls_conditioning_scale=1.0, + ... prompt=prompt, + ... negative_prompt=negative_prompt, + ... num_frames=num_frames, + ... ).frames[0] + >>> export_to_video(video, "edge_controlled_video.mp4", fps=30) + ``` +""" + + +class Cosmos2_5_TransferPipeline(DiffusionPipeline): + r""" + Pipeline for Cosmos Transfer2.5, supporting auto-regressive inference. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods + implemented for all pipelines (downloading, saving, running on a particular device, etc.). + + Args: + text_encoder ([`Qwen2_5_VLForConditionalGeneration`]): + Frozen text-encoder. Cosmos Transfer2.5 uses the [Qwen2.5 + VL](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct) encoder. + tokenizer (`AutoTokenizer`): + Tokenizer associated with the Qwen2.5 VL encoder. + transformer ([`CosmosTransformer3DModel`]): + Conditional Transformer to denoise the encoded image latents. + scheduler ([`UniPCMultistepScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + vae ([`AutoencoderKLWan`]): + Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. + controlnet ([`CosmosControlNetModel`]): + ControlNet used to condition generation on control inputs. + """ + + model_cpu_offload_seq = "text_encoder->transformer->controlnet->vae" + _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] + # We mark safety_checker as optional here to get around some test failures, but it is not really optional + _optional_components = ["safety_checker"] + _exclude_from_cpu_offload = ["safety_checker"] + + def __init__( + self, + text_encoder: Qwen2_5_VLForConditionalGeneration, + tokenizer: AutoTokenizer, + transformer: CosmosTransformer3DModel, + vae: AutoencoderKLWan, + scheduler: UniPCMultistepScheduler, + controlnet: CosmosControlNetModel, + safety_checker: Optional[CosmosSafetyChecker] = None, + ): + super().__init__() + + if safety_checker is None: + safety_checker = CosmosSafetyChecker() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + controlnet=controlnet, + scheduler=scheduler, + safety_checker=safety_checker, + ) + + self.vae_scale_factor_temporal = 2 ** sum(self.vae.temperal_downsample) if getattr(self, "vae", None) else 4 + self.vae_scale_factor_spatial = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8 + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + + latents_mean = ( + torch.tensor(self.vae.config.latents_mean).view(1, self.vae.config.z_dim, 1, 1, 1).float() + if getattr(self.vae.config, "latents_mean", None) is not None + else None + ) + latents_std = ( + torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).float() + if getattr(self.vae.config, "latents_std", None) is not None + else None + ) + self.latents_mean = latents_mean + self.latents_std = latents_std + + if self.latents_mean is None or self.latents_std is None: + raise ValueError("VAE configuration must define both `latents_mean` and `latents_std`.") + + def _get_prompt_embeds( + self, + prompt: Union[str, List[str]] = None, + max_sequence_length: int = 512, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + prompt = [prompt] if isinstance(prompt, str) else prompt + + input_ids_batch = [] + + for sample_idx in range(len(prompt)): + conversations = [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a helpful assistant who will provide prompts to an image generator.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": prompt[sample_idx], + } + ], + }, + ] + input_ids = self.tokenizer.apply_chat_template( + conversations, + tokenize=True, + add_generation_prompt=False, + add_vision_id=False, + max_length=max_sequence_length, + truncation=True, + padding="max_length", + ) + input_ids = ( + input_ids["input_ids"] if not isinstance(input_ids, list) and "input_ids" in input_ids else input_ids + ) + input_ids = torch.LongTensor(input_ids) + input_ids_batch.append(input_ids) + + input_ids_batch = torch.stack(input_ids_batch, dim=0) + + outputs = self.text_encoder( + input_ids_batch.to(device), + output_hidden_states=True, + ) + hidden_states = outputs.hidden_states + + normalized_hidden_states = [] + for layer_idx in range(1, len(hidden_states)): + normalized_state = (hidden_states[layer_idx] - hidden_states[layer_idx].mean(dim=-1, keepdim=True)) / ( + hidden_states[layer_idx].std(dim=-1, keepdim=True) + 1e-8 + ) + normalized_hidden_states.append(normalized_state) + + prompt_embeds = torch.cat(normalized_hidden_states, dim=-1) + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + + return prompt_embeds + + # Modified from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline.encode_prompt + def encode_prompt( + self, + prompt: Union[str, List[str]], + negative_prompt: Optional[Union[str, List[str]]] = None, + do_classifier_free_guidance: bool = True, + num_videos_per_prompt: int = 1, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + max_sequence_length: int = 512, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): + Whether to use classifier free guidance or not. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos that should be generated per prompt. torch device to place the resulting embeddings on + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + device: (`torch.device`, *optional*): + torch device + dtype: (`torch.dtype`, *optional*): + torch dtype + """ + device = device or self._execution_device + + prompt = [prompt] if isinstance(prompt, str) else prompt + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds = self._get_prompt_embeds( + prompt=prompt, max_sequence_length=max_sequence_length, device=device, dtype=dtype + ) + + # duplicate text embeddings for each generation per prompt, using mps friendly method + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) + + if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + + if prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + + negative_prompt_embeds = self._get_prompt_embeds( + prompt=negative_prompt, max_sequence_length=max_sequence_length, device=device, dtype=dtype + ) + + # duplicate text embeddings for each generation per prompt, using mps friendly method + _, seq_len, _ = negative_prompt_embeds.shape + negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_videos_per_prompt, 1) + negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) + + return prompt_embeds, negative_prompt_embeds + + # Modified from diffusers.pipelines.cosmos.pipeline_cosmos2_video2world.Cosmos2VideoToWorldPipeline.prepare_latents and + # diffusers.pipelines.cosmos.pipeline_cosmos2_video2world.Cosmos2TextToImagePipeline.prepare_latents + def prepare_latents( + self, + video: Optional[torch.Tensor], + batch_size: int, + num_channels_latents: int = 16, + height: int = 704, + width: int = 1280, + num_frames_in: int = 93, + num_frames_out: int = 93, + do_classifier_free_guidance: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + num_cond_latent_frames: int = 0, + ) -> torch.Tensor: + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + B = batch_size + C = num_channels_latents + T = (num_frames_out - 1) // self.vae_scale_factor_temporal + 1 + H = height // self.vae_scale_factor_spatial + W = width // self.vae_scale_factor_spatial + shape = (B, C, T, H, W) + + if latents is not None: + if latents.shape[1:] != shape[1:]: + raise ValueError(f"Unexpected `latents` shape, got {latents.shape}, expected {shape}.") + latents = latents.to(device=device, dtype=dtype) + else: + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + + if num_frames_in == 0: + cond_mask = torch.zeros((B, 1, T, H, W), dtype=latents.dtype, device=latents.device) + cond_indicator = torch.zeros((B, 1, T, 1, 1), dtype=latents.dtype, device=latents.device) + + cond_latents = torch.zeros_like(latents) + + return ( + latents, + cond_latents, + cond_mask, + cond_indicator, + ) + else: + if video is None: + raise ValueError("`video` must be provided when `num_frames_in` is greater than 0.") + video = video.to(device=device, dtype=self.vae.dtype) + if isinstance(generator, list): + cond_latents = [ + retrieve_latents(self.vae.encode(video[i].unsqueeze(0)), generator=generator[i]) + for i in range(batch_size) + ] + else: + cond_latents = [retrieve_latents(self.vae.encode(vid.unsqueeze(0)), generator) for vid in video] + + cond_latents = torch.cat(cond_latents, dim=0).to(dtype) + + latents_mean = self.latents_mean.to(device=device, dtype=dtype) + latents_std = self.latents_std.to(device=device, dtype=dtype) + cond_latents = (cond_latents - latents_mean) / latents_std + + padding_shape = (B, 1, T, H, W) + ones_padding = latents.new_ones(padding_shape) + zeros_padding = latents.new_zeros(padding_shape) + + cond_indicator = latents.new_zeros(B, 1, latents.size(2), 1, 1) + cond_indicator[:, :, 0:num_cond_latent_frames, :, :] = 1.0 + cond_mask = cond_indicator * ones_padding + (1 - cond_indicator) * zeros_padding + + return ( + latents, + cond_latents, + cond_mask, + cond_indicator, + ) + + # Modified from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline.check_inputs + def check_inputs( + self, + prompt, + height, + width, + prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + num_ar_conditional_frames=None, + num_ar_latent_conditional_frames=None, + num_frames_per_chunk=None, + num_frames=None, + conditional_frame_timestep=0.1, + ): + if width <= 0 or height <= 0 or height % 16 != 0 or width % 16 != 0: + raise ValueError( + f"`height` and `width` have to be divisible by 16 (& positive) but are {height} and {width}." + ) + + if num_frames is not None and num_frames <= 0: + raise ValueError(f"`num_frames` has to be a positive integer when provided but is {num_frames}.") + + if conditional_frame_timestep < 0 or conditional_frame_timestep > 1: + raise ValueError( + "`conditional_frame_timestep` has to be a float in the [0, 1] interval but is " + f"{conditional_frame_timestep}." + ) + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if num_ar_latent_conditional_frames is not None and num_ar_conditional_frames is not None: + raise ValueError( + "Provide only one of `num_ar_conditional_frames` or `num_ar_latent_conditional_frames`, not both." + ) + if num_ar_latent_conditional_frames is None and num_ar_conditional_frames is None: + raise ValueError("Provide either `num_ar_conditional_frames` or `num_ar_latent_conditional_frames`.") + if num_ar_latent_conditional_frames is not None and num_ar_latent_conditional_frames < 0: + raise ValueError("`num_ar_latent_conditional_frames` must be >= 0.") + if num_ar_conditional_frames is not None and num_ar_conditional_frames < 0: + raise ValueError("`num_ar_conditional_frames` must be >= 0.") + + if num_ar_latent_conditional_frames is not None: + num_ar_conditional_frames = max( + 0, (num_ar_latent_conditional_frames - 1) * self.vae_scale_factor_temporal + 1 + ) + + min_chunk_len = self.vae_scale_factor_temporal + 1 + if num_frames_per_chunk < min_chunk_len: + logger.warning(f"{num_frames_per_chunk=} must be larger than {min_chunk_len=}, setting to min_chunk_len") + num_frames_per_chunk = min_chunk_len + + max_frames_by_rope = None + if getattr(self.transformer.config, "max_size", None) is not None: + max_frames_by_rope = max( + size // patch + for size, patch in zip(self.transformer.config.max_size, self.transformer.config.patch_size) + ) + if num_frames_per_chunk > max_frames_by_rope: + raise ValueError( + f"{num_frames_per_chunk=} is too large for RoPE setting ({max_frames_by_rope=}). " + "Please reduce `num_frames_per_chunk`." + ) + + if num_ar_conditional_frames >= num_frames_per_chunk: + raise ValueError( + f"{num_ar_conditional_frames=} must be smaller than {num_frames_per_chunk=} for chunked generation." + ) + + return num_frames_per_chunk + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + controls: PipelineImageInput | List[PipelineImageInput], + controls_conditioning_scale: Union[float, List[float]] = 1.0, + prompt: Union[str, List[str]] | None = None, + negative_prompt: Union[str, List[str]] = DEFAULT_NEGATIVE_PROMPT, + height: int = 704, + width: Optional[int] = None, + num_frames: Optional[int] = None, + num_frames_per_chunk: int = 93, + num_inference_steps: int = 36, + guidance_scale: float = 3.0, + num_videos_per_prompt: int = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + callback_on_step_end: Optional[ + Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] + ] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + max_sequence_length: int = 512, + conditional_frame_timestep: float = 0.1, + num_ar_conditional_frames: Optional[int] = 1, + num_ar_latent_conditional_frames: Optional[int] = None, + ): + r""" + `controls` drive the conditioning through ControlNet. Controls are assumed to be pre-processed, e.g. edge maps + are pre-computed. + + Setting `num_frames` will restrict the total number of frames output, if not provided or assigned to None + (default) then the number of output frames will match the input `controls`. + + Auto-regressive inference is supported and thus a sliding window of `num_frames_per_chunk` frames are used per + denoising loop. In addition, when auto-regressive inference is performed, the previous + `num_ar_latent_conditional_frames` or `num_ar_conditional_frames` are used to condition the following denoising + inference loops. + + Args: + controls (`PipelineImageInput`, `List[PipelineImageInput]`): + Control image or video input used by the ControlNet. + controls_conditioning_scale (`float` or `List[float]`, *optional*, defaults to `1.0`): + The scale factor(s) for the ControlNet outputs. A single float is broadcast to all control blocks. + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide generation. Required unless `prompt_embeds` is supplied. + height (`int`, defaults to `704`): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. If not provided, this will be determined based on the + aspect ratio of the input and the provided height. + num_frames (`int`, *optional*): + Number of output frames. Defaults to `None` to output the same number of frames as the input + `controls`. + num_inference_steps (`int`, defaults to `36`): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + guidance_scale (`float`, defaults to `3.0`): + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. + of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting + `guidance_scale > 1`. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make + generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs. Can be used to + tweak the same generation with different prompts. If not provided, a latents tensor is generated by + sampling using the supplied random `generator`. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not + provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generated image. Choose between `PIL.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`CosmosPipelineOutput`] instead of a plain tuple. + callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): + A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of + each denoising step during the inference. with the following arguments: `callback_on_step_end(self: + DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a + list of all tensors as specified by `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int`, defaults to `512`): + The maximum number of tokens in the prompt. If the prompt exceeds this length, it will be truncated. If + the prompt is shorter than this length, it will be padded. + num_ar_conditional_frames (`int`, *optional*, defaults to `1`): + Number of frames to condition on subsequent inference loops in auto-regressive inference, i.e. for the + second chunk and onwards. Only used if `num_ar_latent_conditional_frames` is `None`. + + This is only used when auto-regressive inference is performed, i.e. when the number of frames in + controls is > num_frames_per_chunk + num_ar_latent_conditional_frames (`int`, *optional*): + Number of latent frames to condition on subsequent inference loops in auto-regressive inference, i.e. + for the second chunk and onwards. Only used if `num_ar_conditional_frames` is `None`. + + This is only used when auto-regressive inference is performed, i.e. when the number of frames in + controls is > num_frames_per_chunk + Examples: + + Returns: + [`~CosmosPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`CosmosPipelineOutput`] is returned, otherwise a `tuple` is returned where + the first element is a list with the generated images and the second element is a list of `bool`s + indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content. + """ + if self.safety_checker is None: + raise ValueError( + f"You have disabled the safety checker for {self.__class__}. This is in violation of the " + "[NVIDIA Open Model License Agreement](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license). " + f"Please ensure that you are compliant with the license agreement." + ) + + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + if width is None: + frame = controls[0] if isinstance(controls, list) else controls + if isinstance(frame, list): + frame = frame[0] + if isinstance(frame, (torch.Tensor, np.ndarray)): + if frame.ndim == 5: + frame = frame[0, 0] + elif frame.ndim == 4: + frame = frame[0] + + if isinstance(frame, PIL.Image.Image): + width = int((height + 16) * (frame.width / frame.height)) + else: + if frame.ndim != 3: + raise ValueError("`controls` must contain 3D frames in CHW format.") + width = int((height + 16) * (frame.shape[2] / frame.shape[1])) # NOTE: assuming C H W + + num_frames_per_chunk = self.check_inputs( + prompt, + height, + width, + prompt_embeds, + callback_on_step_end_tensor_inputs, + num_ar_conditional_frames, + num_ar_latent_conditional_frames, + num_frames_per_chunk, + num_frames, + conditional_frame_timestep, + ) + + if num_ar_latent_conditional_frames is not None: + num_cond_latent_frames = num_ar_latent_conditional_frames + num_ar_conditional_frames = max(0, (num_cond_latent_frames - 1) * self.vae_scale_factor_temporal + 1) + else: + num_cond_latent_frames = max(0, (num_ar_conditional_frames - 1) // self.vae_scale_factor_temporal + 1) + + self._guidance_scale = guidance_scale + self._current_timestep = None + self._interrupt = False + + device = self._execution_device + + if self.safety_checker is not None: + self.safety_checker.to(device) + if prompt is not None: + prompt_list = [prompt] if isinstance(prompt, str) else prompt + for p in prompt_list: + if not self.safety_checker.check_text_safety(p): + raise ValueError( + f"Cosmos Guardrail detected unsafe text in the prompt: {p}. Please ensure that the " + f"prompt abides by the NVIDIA Open Model License Agreement." + ) + + # Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + # Encode input prompt + ( + prompt_embeds, + negative_prompt_embeds, + ) = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + num_videos_per_prompt=num_videos_per_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + device=device, + max_sequence_length=max_sequence_length, + ) + + vae_dtype = self.vae.dtype + transformer_dtype = self.transformer.dtype + + if getattr(self.transformer.config, "img_context_dim_in", None): + img_context = torch.zeros( + batch_size, + self.transformer.config.img_context_num_tokens, + self.transformer.config.img_context_dim_in, + device=prompt_embeds.device, + dtype=transformer_dtype, + ) + + if num_videos_per_prompt > 1: + img_context = img_context.repeat_interleave(num_videos_per_prompt, dim=0) + + encoder_hidden_states = (prompt_embeds, img_context) + neg_encoder_hidden_states = (negative_prompt_embeds, img_context) + else: + encoder_hidden_states = prompt_embeds + neg_encoder_hidden_states = negative_prompt_embeds + + control_video = self.video_processor.preprocess_video(controls, height, width) + if control_video.shape[0] != batch_size: + if control_video.shape[0] == 1: + control_video = control_video.repeat(batch_size, 1, 1, 1, 1) + else: + raise ValueError( + f"Expected controls batch size {batch_size} to match prompt batch size, but got {control_video.shape[0]}." + ) + + num_frames_out = control_video.shape[2] + if num_frames is not None: + num_frames_out = min(num_frames_out, num_frames) + + control_video = _maybe_pad_or_trim_video(control_video, num_frames_out) + + # chunk information + num_latent_frames_per_chunk = (num_frames_per_chunk - 1) // self.vae_scale_factor_temporal + 1 + chunk_stride = num_frames_per_chunk - num_ar_conditional_frames + chunk_idxs = [ + (start_idx, min(start_idx + num_frames_per_chunk, num_frames_out)) + for start_idx in range(0, num_frames_out - num_ar_conditional_frames, chunk_stride) + ] + + video_chunks = [] + latents_mean = self.latents_mean.to(dtype=vae_dtype, device=device) + latents_std = self.latents_std.to(dtype=vae_dtype, device=device) + + def decode_latents(latents): + latents = latents * latents_std + latents_mean + video = self.vae.decode(latents.to(dtype=self.vae.dtype, device=device), return_dict=False)[0] + return video + + latents_arg = latents + initial_num_cond_latent_frames = 0 + latent_chunks = [] + num_chunks = len(chunk_idxs) + total_steps = num_inference_steps * num_chunks + with self.progress_bar(total=total_steps) as progress_bar: + for chunk_idx, (start_idx, end_idx) in enumerate(chunk_idxs): + if chunk_idx == 0: + prev_output = torch.zeros((batch_size, num_frames_per_chunk, 3, height, width), dtype=vae_dtype) + prev_output = self.video_processor.preprocess_video(prev_output, height, width) + else: + prev_output = video_chunks[-1].clone() + if num_ar_conditional_frames > 0: + prev_output[:, :, :num_ar_conditional_frames] = prev_output[:, :, -num_ar_conditional_frames:] + prev_output[:, :, num_ar_conditional_frames:] = -1 # -1 == 0 in processed video space + else: + prev_output.fill_(-1) + + chunk_video = prev_output.to(device=device, dtype=vae_dtype) + chunk_video = _maybe_pad_or_trim_video(chunk_video, num_frames_per_chunk) + latents, cond_latent, cond_mask, cond_indicator = self.prepare_latents( + video=chunk_video, + batch_size=batch_size * num_videos_per_prompt, + num_channels_latents=self.transformer.config.in_channels - 1, + height=height, + width=width, + num_frames_in=chunk_video.shape[2], + num_frames_out=num_frames_per_chunk, + do_classifier_free_guidance=self.do_classifier_free_guidance, + dtype=torch.float32, + device=device, + generator=generator, + num_cond_latent_frames=initial_num_cond_latent_frames + if chunk_idx == 0 + else num_cond_latent_frames, + latents=latents_arg, + ) + cond_mask = cond_mask.to(transformer_dtype) + cond_timestep = torch.ones_like(cond_indicator) * conditional_frame_timestep + padding_mask = latents.new_zeros(1, 1, height, width, dtype=transformer_dtype) + + chunk_control_video = control_video[:, :, start_idx:end_idx, ...].to( + device=device, dtype=self.vae.dtype + ) + chunk_control_video = _maybe_pad_or_trim_video(chunk_control_video, num_frames_per_chunk) + if isinstance(generator, list): + controls_latents = [ + retrieve_latents(self.vae.encode(chunk_control_video[i].unsqueeze(0)), generator=generator[i]) + for i in range(chunk_control_video.shape[0]) + ] + else: + controls_latents = [ + retrieve_latents(self.vae.encode(vid.unsqueeze(0)), generator=generator) + for vid in chunk_control_video + ] + controls_latents = torch.cat(controls_latents, dim=0).to(transformer_dtype) + + controls_latents = (controls_latents - latents_mean) / latents_std + + # Denoising loop + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + self._num_timesteps = len(timesteps) + + gt_velocity = (latents - cond_latent) * cond_mask + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + self._current_timestep = t.cpu().item() + + # NOTE: assumes sigma(t) \in [0, 1] + sigma_t = ( + torch.tensor(self.scheduler.sigmas[i].item()) + .unsqueeze(0) + .to(device=device, dtype=transformer_dtype) + ) + + in_latents = cond_mask * cond_latent + (1 - cond_mask) * latents + in_latents = in_latents.to(transformer_dtype) + in_timestep = cond_indicator * cond_timestep + (1 - cond_indicator) * sigma_t + control_output = self.controlnet( + controls_latents=controls_latents, + latents=in_latents, + timestep=in_timestep, + encoder_hidden_states=encoder_hidden_states, + condition_mask=cond_mask, + conditioning_scale=controls_conditioning_scale, + padding_mask=padding_mask, + return_dict=False, + ) + control_blocks = control_output[0] + + noise_pred = self.transformer( + hidden_states=in_latents, + timestep=in_timestep, + encoder_hidden_states=encoder_hidden_states, + block_controlnet_hidden_states=control_blocks, + condition_mask=cond_mask, + padding_mask=padding_mask, + return_dict=False, + )[0] + noise_pred = gt_velocity + noise_pred * (1 - cond_mask) + + if self.do_classifier_free_guidance: + control_output = self.controlnet( + controls_latents=controls_latents, + latents=in_latents, + timestep=in_timestep, + encoder_hidden_states=neg_encoder_hidden_states, # NOTE: negative prompt + condition_mask=cond_mask, + conditioning_scale=controls_conditioning_scale, + padding_mask=padding_mask, + return_dict=False, + ) + control_blocks = control_output[0] + + noise_pred_neg = self.transformer( + hidden_states=in_latents, + timestep=in_timestep, + encoder_hidden_states=neg_encoder_hidden_states, # NOTE: negative prompt + block_controlnet_hidden_states=control_blocks, + condition_mask=cond_mask, + padding_mask=padding_mask, + return_dict=False, + )[0] + # NOTE: replace velocity (noise_pred_neg) with gt_velocity for conditioning inputs only + noise_pred_neg = gt_velocity + noise_pred_neg * (1 - cond_mask) + noise_pred = noise_pred + self.guidance_scale * (noise_pred - noise_pred_neg) + + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + # call the callback, if provided + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + + if i == total_steps - 1 or ((i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + video_chunks.append(decode_latents(latents).detach().cpu()) + latent_chunks.append(latents.detach().cpu()) + + self._current_timestep = None + + if not output_type == "latent": + video_chunks = [ + chunk[:, :, num_ar_conditional_frames:, ...] if chunk_idx != 0 else chunk + for chunk_idx, chunk in enumerate(video_chunks) + ] + video = torch.cat(video_chunks, dim=2) + video = video[:, :, :num_frames_out, ...] + + assert self.safety_checker is not None + self.safety_checker.to(device) + video = self.video_processor.postprocess_video(video, output_type="np") + video = (video * 255).astype(np.uint8) + video_batch = [] + for vid in video: + vid = self.safety_checker.check_video_safety(vid) + if vid is None: + video_batch.append(np.zeros_like(video[0])) + else: + video_batch.append(vid) + video = np.stack(video_batch).astype(np.float32) / 255.0 * 2 - 1 + video = torch.from_numpy(video).permute(0, 4, 1, 2, 3) + video = self.video_processor.postprocess_video(video, output_type=output_type) + else: + latent_T = (num_frames_out - 1) // self.vae_scale_factor_temporal + 1 + latent_chunks = [ + chunk[:, :, num_cond_latent_frames:, ...] if chunk_idx != 0 else chunk + for chunk_idx, chunk in enumerate(latent_chunks) + ] + video = torch.cat(latent_chunks, dim=2) + video = video[:, :, :latent_T, ...] + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (video,) + + return CosmosPipelineOutput(frames=video) diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos2_text2image.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos2_text2image.py index 66490c2be159..f24e19eea0d4 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos2_text2image.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos2_text2image.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import numpy as np import torch @@ -49,6 +49,14 @@ def __init__(self, *args, **kwargs): logger = logging.get_logger(__name__) # pylint: disable=invalid-name +DEFAULT_NEGATIVE_PROMPT = ( + "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, " + "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, " + "underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, " + "jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, " + "fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. " + "Overall, the video is of poor quality." +) EXAMPLE_DOC_STRING = """ Examples: @@ -75,10 +83,10 @@ def __init__(self, *args, **kwargs): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -93,15 +101,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -202,10 +210,10 @@ def __init__( # Copied from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -245,23 +253,23 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline.encode_prompt with num_videos_per_prompt->num_images_per_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -300,7 +308,7 @@ def encode_prompt( prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) if do_classifier_free_guidance and negative_prompt_embeds is None: - negative_prompt = negative_prompt or "" + negative_prompt = negative_prompt if negative_prompt is not None else DEFAULT_NEGATIVE_PROMPT negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt if prompt is not None and type(prompt) is not type(negative_prompt): @@ -333,10 +341,10 @@ def prepare_latents( height: int = 768, width: int = 1360, num_frames: int = 1, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: if latents is not None: return latents.to(device=device, dtype=dtype) * self.scheduler.config.sigma_max @@ -410,30 +418,28 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] | None = None, height: int = 768, width: int = 1360, num_inference_steps: int = 35, guidance_scale: float = 7.0, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. height (`int`, defaults to `768`): @@ -450,7 +456,7 @@ def __call__( `guidance_scale > 1`. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -472,7 +478,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py index 23a74ad00f93..bdb13af06637 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import numpy as np import torch @@ -50,6 +50,14 @@ def __init__(self, *args, **kwargs): logger = logging.get_logger(__name__) # pylint: disable=invalid-name +DEFAULT_NEGATIVE_PROMPT = ( + "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, " + "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, " + "underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, " + "jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, " + "fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. " + "Overall, the video is of poor quality." +) EXAMPLE_DOC_STRING = """ Examples: @@ -80,10 +88,10 @@ def __init__(self, *args, **kwargs): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -98,15 +106,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -139,7 +147,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -221,10 +229,10 @@ def __init__( # Copied from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -264,23 +272,23 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -319,7 +327,7 @@ def encode_prompt( prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) if do_classifier_free_guidance and negative_prompt_embeds is None: - negative_prompt = negative_prompt or "" + negative_prompt = negative_prompt if negative_prompt is not None else DEFAULT_NEGATIVE_PROMPT negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt if prompt is not None and type(prompt) is not type(negative_prompt): @@ -354,10 +362,10 @@ def prepare_latents( width: int = 1280, num_frames: int = 93, do_classifier_free_guidance: bool = True, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( @@ -479,26 +487,24 @@ def interrupt(self): def __call__( self, image: PipelineImageInput = None, - video: List[PipelineImageInput] = None, - prompt: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, + video: list[PipelineImageInput] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] | None = None, height: int = 704, width: int = 1280, num_frames: int = 93, num_inference_steps: int = 35, guidance_scale: float = 7.0, fps: int = 16, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, sigma_conditioning: float = 0.0001, ): @@ -508,9 +514,9 @@ def __call__( Args: image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, *optional*): The image to be used as a conditioning input for the video generation. - video (`List[PIL.Image.Image]`, `np.ndarray`, `torch.Tensor`, *optional*): + video (`list[PIL.Image.Image]`, `np.ndarray`, `torch.Tensor`, *optional*): The video to be used as a conditioning input for the video generation. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. height (`int`, defaults to `704`): @@ -531,7 +537,7 @@ def __call__( The frames per second of the generated video. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -553,7 +559,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py index f0aa1ecf0e0f..e144d62d5933 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import numpy as np import torch @@ -49,6 +49,14 @@ def __init__(self, *args, **kwargs): logger = logging.get_logger(__name__) # pylint: disable=invalid-name +DEFAULT_NEGATIVE_PROMPT = ( + "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, " + "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, " + "underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, " + "jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, " + "fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. " + "Overall, the video is of poor quality." +) EXAMPLE_DOC_STRING = """ Examples: @@ -72,10 +80,10 @@ def __init__(self, *args, **kwargs): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -90,15 +98,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -188,10 +196,10 @@ def __init__( def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -230,23 +238,23 @@ def _get_t5_prompt_embeds( def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -285,7 +293,7 @@ def encode_prompt( prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) if do_classifier_free_guidance and negative_prompt_embeds is None: - negative_prompt = negative_prompt or "" + negative_prompt = negative_prompt if negative_prompt is not None else DEFAULT_NEGATIVE_PROMPT negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt if prompt is not None and type(prompt) is not type(negative_prompt): @@ -318,10 +326,10 @@ def prepare_latents( height: int = 704, width: int = 1280, num_frames: int = 121, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: if latents is not None: return latents.to(device=device, dtype=dtype) * self.scheduler.config.sigma_max @@ -394,32 +402,30 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] | None = None, height: int = 704, width: int = 1280, num_frames: int = 121, num_inference_steps: int = 36, guidance_scale: float = 7.0, fps: int = 30, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. height (`int`, defaults to `720`): @@ -440,7 +446,7 @@ def __call__( The frames per second of the generated video. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -462,7 +468,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py index cd5a734cc311..377c3c05d284 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import numpy as np import torch @@ -50,6 +50,14 @@ def __init__(self, *args, **kwargs): logger = logging.get_logger(__name__) # pylint: disable=invalid-name +DEFAULT_NEGATIVE_PROMPT = ( + "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, " + "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, " + "underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, " + "jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, " + "fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. " + "Overall, the video is of poor quality." +) EXAMPLE_DOC_STRING = """ Examples: @@ -101,10 +109,10 @@ def __init__(self, *args, **kwargs): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -119,15 +127,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -160,7 +168,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -233,10 +241,10 @@ def __init__( # Copied from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -276,23 +284,23 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -331,7 +339,7 @@ def encode_prompt( prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) if do_classifier_free_guidance and negative_prompt_embeds is None: - negative_prompt = negative_prompt or "" + negative_prompt = negative_prompt if negative_prompt is not None else DEFAULT_NEGATIVE_PROMPT negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt if prompt is not None and type(prompt) is not type(negative_prompt): @@ -367,10 +375,10 @@ def prepare_latents( num_frames: int = 121, do_classifier_free_guidance: bool = True, input_frames_guidance: bool = False, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( @@ -507,9 +515,9 @@ def interrupt(self): def __call__( self, image: PipelineImageInput = None, - video: List[PipelineImageInput] = None, - prompt: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, + video: list[PipelineImageInput] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] | None = None, height: int = 704, width: int = 1280, num_frames: int = 121, @@ -518,24 +526,22 @@ def __call__( input_frames_guidance: bool = False, augment_sigma: float = 0.001, fps: int = 30, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. height (`int`, defaults to `720`): @@ -556,7 +562,7 @@ def __call__( The frames per second of the generated video. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -578,7 +584,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/cosmos/pipeline_output.py b/src/diffusers/pipelines/cosmos/pipeline_output.py index ec5f4826f62a..1ded292f8dfb 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_output.py +++ b/src/diffusers/pipelines/cosmos/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -17,8 +16,8 @@ class CosmosPipelineOutput(BaseOutput): Output class for Cosmos any-to-world/video pipelines. Args: - frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): - List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): + list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape `(batch_size, num_frames, channels, height, width)`. """ @@ -32,9 +31,9 @@ class CosmosImagePipelineOutput(BaseOutput): Output class for Cosmos any-to-image pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image] | np.ndarray diff --git a/src/diffusers/pipelines/dance_diffusion/pipeline_dance_diffusion.py b/src/diffusers/pipelines/dance_diffusion/pipeline_dance_diffusion.py index 5a70c4f5ff9a..eb8f8106061d 100644 --- a/src/diffusers/pipelines/dance_diffusion/pipeline_dance_diffusion.py +++ b/src/diffusers/pipelines/dance_diffusion/pipeline_dance_diffusion.py @@ -13,8 +13,6 @@ # limitations under the License. -from typing import List, Optional, Tuple, Union - import torch from ...models import UNet1DModel @@ -61,10 +59,10 @@ def __call__( self, batch_size: int = 1, num_inference_steps: int = 100, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - audio_length_in_s: Optional[float] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + audio_length_in_s: float | None = None, return_dict: bool = True, - ) -> Union[AudioPipelineOutput, Tuple]: + ) -> AudioPipelineOutput | tuple: r""" The call function to the pipeline for generation. diff --git a/src/diffusers/pipelines/ddim/pipeline_ddim.py b/src/diffusers/pipelines/ddim/pipeline_ddim.py index 39587ca5221d..dc92b34a7565 100644 --- a/src/diffusers/pipelines/ddim/pipeline_ddim.py +++ b/src/diffusers/pipelines/ddim/pipeline_ddim.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Optional, Tuple, Union - import torch from ...models import UNet2DModel @@ -60,13 +58,13 @@ def __init__(self, unet: UNet2DModel, scheduler: DDIMScheduler): def __call__( self, batch_size: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + generator: torch.Generator | list[torch.Generator] | None = None, eta: float = 0.0, num_inference_steps: int = 50, - use_clipped_model_output: Optional[bool] = None, - output_type: Optional[str] = "pil", + use_clipped_model_output: bool | None = None, + output_type: str | None = "pil", return_dict: bool = True, - ) -> Union[ImagePipelineOutput, Tuple]: + ) -> ImagePipelineOutput | tuple: r""" The call function to the pipeline for generation. diff --git a/src/diffusers/pipelines/ddpm/pipeline_ddpm.py b/src/diffusers/pipelines/ddpm/pipeline_ddpm.py index 0d7766a8cfd0..6d4796cbea1f 100644 --- a/src/diffusers/pipelines/ddpm/pipeline_ddpm.py +++ b/src/diffusers/pipelines/ddpm/pipeline_ddpm.py @@ -13,8 +13,6 @@ # limitations under the License. -from typing import List, Optional, Tuple, Union - import torch from ...models import UNet2DModel @@ -57,11 +55,11 @@ def __init__(self, unet: UNet2DModel, scheduler: DDPMScheduler): def __call__( self, batch_size: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + generator: torch.Generator | list[torch.Generator] | None = None, num_inference_steps: int = 1000, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, - ) -> Union[ImagePipelineOutput, Tuple]: + ) -> ImagePipelineOutput | tuple: r""" The call function to the pipeline for generation. diff --git a/src/diffusers/pipelines/deepfloyd_if/pipeline_if.py b/src/diffusers/pipelines/deepfloyd_if/pipeline_if.py index 8fa31f8504d3..b8c70fc6528c 100644 --- a/src/diffusers/pipelines/deepfloyd_if/pipeline_if.py +++ b/src/diffusers/pipelines/deepfloyd_if/pipeline_if.py @@ -2,7 +2,7 @@ import inspect import re import urllib.parse as ul -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch from transformers import CLIPImageProcessor, T5EncoderModel, T5Tokenizer @@ -100,10 +100,10 @@ class IFPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin): unet: UNet2DConditionModel scheduler: DDPMScheduler - feature_extractor: Optional[CLIPImageProcessor] - safety_checker: Optional[IFSafetyChecker] + feature_extractor: CLIPImageProcessor | None + safety_checker: IFSafetyChecker | None - watermarker: Optional[IFWatermarker] + watermarker: IFWatermarker | None bad_punct_regex = re.compile( r"[" @@ -131,9 +131,9 @@ def __init__( text_encoder: T5EncoderModel, unet: UNet2DConditionModel, scheduler: DDPMScheduler, - safety_checker: Optional[IFSafetyChecker], - feature_extractor: Optional[CLIPImageProcessor], - watermarker: Optional[IFWatermarker], + safety_checker: IFSafetyChecker | None, + feature_extractor: CLIPImageProcessor | None, + watermarker: IFWatermarker | None, requires_safety_checker: bool = True, ): super().__init__() @@ -168,20 +168,20 @@ def __init__( @torch.no_grad() def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], do_classifier_free_guidance: bool = True, num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + device: torch.device | None = None, + negative_prompt: str | list[str] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, clean_caption: bool = False, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): whether to use classifier free guidance or not @@ -189,7 +189,7 @@ def encode_prompt( number of images that should be generated per prompt device: (`torch.device`, *optional*): torch device to place the resulting embeddings on - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -269,7 +269,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif isinstance(negative_prompt, str): @@ -548,36 +548,36 @@ def _clean_caption(self, caption): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_inference_steps: int = 100, - timesteps: List[int] = None, + timesteps: list[int] = None, guidance_scale: float = 7.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - height: Optional[int] = None, - width: Optional[int] = None, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + height: int | None = None, + width: int | None = None, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, clean_caption: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, ): """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. num_inference_steps (`int`, *optional*, defaults to 100): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps` timesteps are used. Must be in descending order. guidance_scale (`float`, *optional*, defaults to 7.0): @@ -586,7 +586,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -599,7 +599,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. prompt_embeds (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py b/src/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py index 507927faf61b..3dadc63f4952 100644 --- a/src/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py +++ b/src/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py @@ -2,7 +2,7 @@ import inspect import re import urllib.parse as ul -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -124,10 +124,10 @@ class IFImg2ImgPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin): unet: UNet2DConditionModel scheduler: DDPMScheduler - feature_extractor: Optional[CLIPImageProcessor] - safety_checker: Optional[IFSafetyChecker] + feature_extractor: CLIPImageProcessor | None + safety_checker: IFSafetyChecker | None - watermarker: Optional[IFWatermarker] + watermarker: IFWatermarker | None bad_punct_regex = re.compile( r"[" @@ -155,9 +155,9 @@ def __init__( text_encoder: T5EncoderModel, unet: UNet2DConditionModel, scheduler: DDPMScheduler, - safety_checker: Optional[IFSafetyChecker], - feature_extractor: Optional[CLIPImageProcessor], - watermarker: Optional[IFWatermarker], + safety_checker: IFSafetyChecker | None, + feature_extractor: CLIPImageProcessor | None, + watermarker: IFWatermarker | None, requires_safety_checker: bool = True, ): super().__init__() @@ -192,20 +192,20 @@ def __init__( @torch.no_grad() def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], do_classifier_free_guidance: bool = True, num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + device: torch.device | None = None, + negative_prompt: str | list[str] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, clean_caption: bool = False, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): whether to use classifier free guidance or not @@ -213,7 +213,7 @@ def encode_prompt( number of images that should be generated per prompt device: (`torch.device`, *optional*): torch device to place the resulting embeddings on - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -293,7 +293,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif isinstance(negative_prompt, str): @@ -430,7 +430,7 @@ def check_inputs( and not isinstance(check_image_type, np.ndarray) ): raise ValueError( - "`image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or List[...] but is" + "`image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list[...] but is" f" {type(check_image_type)}" ) @@ -662,32 +662,35 @@ def prepare_intermediate_images( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - image: Union[ - PIL.Image.Image, torch.Tensor, np.ndarray, List[PIL.Image.Image], List[torch.Tensor], List[np.ndarray] - ] = None, + prompt: str | list[str] = None, + image: PIL.Image.Image + | torch.Tensor + | np.ndarray + | list[PIL.Image.Image] + | list[torch.Tensor] + | list[np.ndarray] = None, strength: float = 0.7, num_inference_steps: int = 80, - timesteps: List[int] = None, + timesteps: list[int] = None, guidance_scale: float = 10.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, clean_caption: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, ): """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. image (`torch.Tensor` or `PIL.Image.Image`): @@ -702,7 +705,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 80): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps` timesteps are used. Must be in descending order. guidance_scale (`float`, *optional*, defaults to 10.0): @@ -711,7 +714,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -720,7 +723,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. prompt_embeds (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py b/src/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py index 9bc15c3c6f62..4839a0860462 100644 --- a/src/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py +++ b/src/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py @@ -2,7 +2,7 @@ import inspect import re import urllib.parse as ul -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -129,10 +129,10 @@ class IFImg2ImgSuperResolutionPipeline(DiffusionPipeline, StableDiffusionLoraLoa scheduler: DDPMScheduler image_noising_scheduler: DDPMScheduler - feature_extractor: Optional[CLIPImageProcessor] - safety_checker: Optional[IFSafetyChecker] + feature_extractor: CLIPImageProcessor | None + safety_checker: IFSafetyChecker | None - watermarker: Optional[IFWatermarker] + watermarker: IFWatermarker | None bad_punct_regex = re.compile( r"[" @@ -161,9 +161,9 @@ def __init__( unet: UNet2DConditionModel, scheduler: DDPMScheduler, image_noising_scheduler: DDPMScheduler, - safety_checker: Optional[IFSafetyChecker], - feature_extractor: Optional[CLIPImageProcessor], - watermarker: Optional[IFWatermarker], + safety_checker: IFSafetyChecker | None, + feature_extractor: CLIPImageProcessor | None, + watermarker: IFWatermarker | None, requires_safety_checker: bool = True, ): super().__init__() @@ -345,20 +345,20 @@ def _clean_caption(self, caption): # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], do_classifier_free_guidance: bool = True, num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + device: torch.device | None = None, + negative_prompt: str | list[str] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, clean_caption: bool = False, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): whether to use classifier free guidance or not @@ -366,7 +366,7 @@ def encode_prompt( number of images that should be generated per prompt device: (`torch.device`, *optional*): torch device to place the resulting embeddings on - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -446,7 +446,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif isinstance(negative_prompt, str): @@ -586,7 +586,7 @@ def check_inputs( and not isinstance(check_image_type, np.ndarray) ): raise ValueError( - "`image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or List[...] but is" + "`image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list[...] but is" f" {type(check_image_type)}" ) @@ -617,7 +617,7 @@ def check_inputs( and not isinstance(check_image_type, np.ndarray) ): raise ValueError( - "`original_image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or List[...] but is" + "`original_image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list[...] but is" f" {type(check_image_type)}" ) @@ -745,26 +745,29 @@ def prepare_intermediate_images( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - image: Union[PIL.Image.Image, np.ndarray, torch.Tensor], - original_image: Union[ - PIL.Image.Image, torch.Tensor, np.ndarray, List[PIL.Image.Image], List[torch.Tensor], List[np.ndarray] - ] = None, + image: PIL.Image.Image | np.ndarray | torch.Tensor, + original_image: PIL.Image.Image + | torch.Tensor + | np.ndarray + | list[PIL.Image.Image] + | list[torch.Tensor] + | list[np.ndarray] = None, strength: float = 0.8, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_inference_steps: int = 50, - timesteps: List[int] = None, + timesteps: list[int] = None, guidance_scale: float = 4.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, noise_level: int = 250, clean_caption: bool = True, ): @@ -783,13 +786,13 @@ def __call__( denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will be maximum and the denoising process will run for the full number of iterations specified in `num_inference_steps`. A value of 1, therefore, essentially ignores `image`. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps` timesteps are used. Must be in descending order. guidance_scale (`float`, *optional*, defaults to 4.0): @@ -798,7 +801,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -807,7 +810,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. prompt_embeds (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py b/src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py index 9d6cf62020a9..03a9d6f7c5e8 100644 --- a/src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py +++ b/src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py @@ -2,7 +2,7 @@ import inspect import re import urllib.parse as ul -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -127,10 +127,10 @@ class IFInpaintingPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin): unet: UNet2DConditionModel scheduler: DDPMScheduler - feature_extractor: Optional[CLIPImageProcessor] - safety_checker: Optional[IFSafetyChecker] + feature_extractor: CLIPImageProcessor | None + safety_checker: IFSafetyChecker | None - watermarker: Optional[IFWatermarker] + watermarker: IFWatermarker | None bad_punct_regex = re.compile( r"[" @@ -158,9 +158,9 @@ def __init__( text_encoder: T5EncoderModel, unet: UNet2DConditionModel, scheduler: DDPMScheduler, - safety_checker: Optional[IFSafetyChecker], - feature_extractor: Optional[CLIPImageProcessor], - watermarker: Optional[IFWatermarker], + safety_checker: IFSafetyChecker | None, + feature_extractor: CLIPImageProcessor | None, + watermarker: IFWatermarker | None, requires_safety_checker: bool = True, ): super().__init__() @@ -196,20 +196,20 @@ def __init__( # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], do_classifier_free_guidance: bool = True, num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + device: torch.device | None = None, + negative_prompt: str | list[str] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, clean_caption: bool = False, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): whether to use classifier free guidance or not @@ -217,7 +217,7 @@ def encode_prompt( number of images that should be generated per prompt device: (`torch.device`, *optional*): torch device to place the resulting embeddings on - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -297,7 +297,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif isinstance(negative_prompt, str): @@ -437,7 +437,7 @@ def check_inputs( and not isinstance(check_image_type, np.ndarray) ): raise ValueError( - "`image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or List[...] but is" + "`image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list[...] but is" f" {type(check_image_type)}" ) @@ -468,7 +468,7 @@ def check_inputs( and not isinstance(check_image_type, np.ndarray) ): raise ValueError( - "`mask_image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or List[...] but is" + "`mask_image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list[...] but is" f" {type(check_image_type)}" ) @@ -754,35 +754,41 @@ def prepare_intermediate_images( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - image: Union[ - PIL.Image.Image, torch.Tensor, np.ndarray, List[PIL.Image.Image], List[torch.Tensor], List[np.ndarray] - ] = None, - mask_image: Union[ - PIL.Image.Image, torch.Tensor, np.ndarray, List[PIL.Image.Image], List[torch.Tensor], List[np.ndarray] - ] = None, + prompt: str | list[str] = None, + image: PIL.Image.Image + | torch.Tensor + | np.ndarray + | list[PIL.Image.Image] + | list[torch.Tensor] + | list[np.ndarray] = None, + mask_image: PIL.Image.Image + | torch.Tensor + | np.ndarray + | list[PIL.Image.Image] + | list[torch.Tensor] + | list[np.ndarray] = None, strength: float = 1.0, num_inference_steps: int = 50, - timesteps: List[int] = None, + timesteps: list[int] = None, guidance_scale: float = 7.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, clean_caption: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, ): """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. image (`torch.Tensor` or `PIL.Image.Image`): @@ -802,7 +808,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps` timesteps are used. Must be in descending order. guidance_scale (`float`, *optional*, defaults to 7.0): @@ -811,7 +817,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -820,7 +826,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. prompt_embeds (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py b/src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py index 0122c164d8b8..841382ad9c63 100644 --- a/src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py +++ b/src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py @@ -2,7 +2,7 @@ import inspect import re import urllib.parse as ul -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -131,10 +131,10 @@ class IFInpaintingSuperResolutionPipeline(DiffusionPipeline, StableDiffusionLora scheduler: DDPMScheduler image_noising_scheduler: DDPMScheduler - feature_extractor: Optional[CLIPImageProcessor] - safety_checker: Optional[IFSafetyChecker] + feature_extractor: CLIPImageProcessor | None + safety_checker: IFSafetyChecker | None - watermarker: Optional[IFWatermarker] + watermarker: IFWatermarker | None bad_punct_regex = re.compile( r"[" @@ -163,9 +163,9 @@ def __init__( unet: UNet2DConditionModel, scheduler: DDPMScheduler, image_noising_scheduler: DDPMScheduler, - safety_checker: Optional[IFSafetyChecker], - feature_extractor: Optional[CLIPImageProcessor], - watermarker: Optional[IFWatermarker], + safety_checker: IFSafetyChecker | None, + feature_extractor: CLIPImageProcessor | None, + watermarker: IFWatermarker | None, requires_safety_checker: bool = True, ): super().__init__() @@ -347,20 +347,20 @@ def _clean_caption(self, caption): # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], do_classifier_free_guidance: bool = True, num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + device: torch.device | None = None, + negative_prompt: str | list[str] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, clean_caption: bool = False, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): whether to use classifier free guidance or not @@ -368,7 +368,7 @@ def encode_prompt( number of images that should be generated per prompt device: (`torch.device`, *optional*): torch device to place the resulting embeddings on - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -448,7 +448,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif isinstance(negative_prompt, str): @@ -589,7 +589,7 @@ def check_inputs( and not isinstance(check_image_type, np.ndarray) ): raise ValueError( - "`image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or List[...] but is" + "`image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list[...] but is" f" {type(check_image_type)}" ) @@ -620,7 +620,7 @@ def check_inputs( and not isinstance(check_image_type, np.ndarray) ): raise ValueError( - "`original_image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or List[...] but is" + "`original_image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list[...] but is" f" {type(check_image_type)}" ) @@ -653,7 +653,7 @@ def check_inputs( and not isinstance(check_image_type, np.ndarray) ): raise ValueError( - "`mask_image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or List[...] but is" + "`mask_image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list[...] but is" f" {type(check_image_type)}" ) @@ -833,29 +833,35 @@ def prepare_intermediate_images( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - image: Union[PIL.Image.Image, np.ndarray, torch.Tensor], - original_image: Union[ - PIL.Image.Image, torch.Tensor, np.ndarray, List[PIL.Image.Image], List[torch.Tensor], List[np.ndarray] - ] = None, - mask_image: Union[ - PIL.Image.Image, torch.Tensor, np.ndarray, List[PIL.Image.Image], List[torch.Tensor], List[np.ndarray] - ] = None, + image: PIL.Image.Image | np.ndarray | torch.Tensor, + original_image: PIL.Image.Image + | torch.Tensor + | np.ndarray + | list[PIL.Image.Image] + | list[torch.Tensor] + | list[np.ndarray] = None, + mask_image: PIL.Image.Image + | torch.Tensor + | np.ndarray + | list[PIL.Image.Image] + | list[torch.Tensor] + | list[np.ndarray] = None, strength: float = 0.8, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_inference_steps: int = 100, - timesteps: List[int] = None, + timesteps: list[int] = None, guidance_scale: float = 4.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, noise_level: int = 0, clean_caption: bool = True, ): @@ -879,13 +885,13 @@ def __call__( denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will be maximum and the denoising process will run for the full number of iterations specified in `num_inference_steps`. A value of 1, therefore, essentially ignores `image`. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. num_inference_steps (`int`, *optional*, defaults to 100): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps` timesteps are used. Must be in descending order. guidance_scale (`float`, *optional*, defaults to 4.0): @@ -894,7 +900,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -903,7 +909,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. prompt_embeds (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py b/src/diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py index ffa60575fe33..52ebebb6f9b4 100644 --- a/src/diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py +++ b/src/diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py @@ -2,7 +2,7 @@ import inspect import re import urllib.parse as ul -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -87,10 +87,10 @@ class IFSuperResolutionPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixi scheduler: DDPMScheduler image_noising_scheduler: DDPMScheduler - feature_extractor: Optional[CLIPImageProcessor] - safety_checker: Optional[IFSafetyChecker] + feature_extractor: CLIPImageProcessor | None + safety_checker: IFSafetyChecker | None - watermarker: Optional[IFWatermarker] + watermarker: IFWatermarker | None bad_punct_regex = re.compile( r"[" @@ -119,9 +119,9 @@ def __init__( unet: UNet2DConditionModel, scheduler: DDPMScheduler, image_noising_scheduler: DDPMScheduler, - safety_checker: Optional[IFSafetyChecker], - feature_extractor: Optional[CLIPImageProcessor], - watermarker: Optional[IFWatermarker], + safety_checker: IFSafetyChecker | None, + feature_extractor: CLIPImageProcessor | None, + watermarker: IFWatermarker | None, requires_safety_checker: bool = True, ): super().__init__() @@ -303,20 +303,20 @@ def _clean_caption(self, caption): # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], do_classifier_free_guidance: bool = True, num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + device: torch.device | None = None, + negative_prompt: str | list[str] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, clean_caption: bool = False, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): whether to use classifier free guidance or not @@ -324,7 +324,7 @@ def encode_prompt( number of images that should be generated per prompt device: (`torch.device`, *optional*): torch device to place the resulting embeddings on - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -404,7 +404,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif isinstance(negative_prompt, str): @@ -547,7 +547,7 @@ def check_inputs( and not isinstance(check_image_type, np.ndarray) ): raise ValueError( - "`image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or List[...] but is" + "`image` has to be of type `torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list[...] but is" f" {type(check_image_type)}" ) @@ -615,24 +615,24 @@ def preprocess_image(self, image, num_images_per_prompt, device): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, height: int = None, width: int = None, - image: Union[PIL.Image.Image, np.ndarray, torch.Tensor] = None, + image: PIL.Image.Image | np.ndarray | torch.Tensor = None, num_inference_steps: int = 50, - timesteps: List[int] = None, + timesteps: list[int] = None, guidance_scale: float = 4.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, noise_level: int = 250, clean_caption: bool = True, ): @@ -640,7 +640,7 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. height (`int`, *optional*, defaults to None): @@ -652,7 +652,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*, defaults to None): + timesteps (`list[int]`, *optional*, defaults to None): Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps` timesteps are used. Must be in descending order. guidance_scale (`float`, *optional*, defaults to 4.0): @@ -661,7 +661,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -670,7 +670,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. prompt_embeds (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/deepfloyd_if/pipeline_output.py b/src/diffusers/pipelines/deepfloyd_if/pipeline_output.py index b8bae89cec03..7fe1cd013835 100644 --- a/src/diffusers/pipelines/deepfloyd_if/pipeline_output.py +++ b/src/diffusers/pipelines/deepfloyd_if/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Optional, Union import numpy as np import PIL.Image @@ -13,17 +12,17 @@ class IFPipelineOutput(BaseOutput): Output class for Stable Diffusion pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`): - List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`): + list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. - nsfw_detected (`List[bool]`): - List of flags denoting whether the corresponding generated image likely represents "not-safe-for-work" + nsfw_detected (`list[bool]`): + list of flags denoting whether the corresponding generated image likely represents "not-safe-for-work" (nsfw) content or a watermark. `None` if safety checking could not be performed. - watermark_detected (`List[bool]`): - List of flags denoting whether the corresponding generated image likely has a watermark. `None` if safety + watermark_detected (`list[bool]`): + list of flags denoting whether the corresponding generated image likely has a watermark. `None` if safety checking could not be performed. """ - images: Union[List[PIL.Image.Image], np.ndarray] - nsfw_detected: Optional[List[bool]] - watermark_detected: Optional[List[bool]] + images: list[PIL.Image.Image] | np.ndarray + nsfw_detected: list[bool] | None + watermark_detected: list[bool] | None diff --git a/src/diffusers/pipelines/deepfloyd_if/watermark.py b/src/diffusers/pipelines/deepfloyd_if/watermark.py index e03e3fab026a..d5fe99f681f7 100644 --- a/src/diffusers/pipelines/deepfloyd_if/watermark.py +++ b/src/diffusers/pipelines/deepfloyd_if/watermark.py @@ -1,5 +1,3 @@ -from typing import List - import PIL.Image import torch from PIL import Image @@ -16,7 +14,7 @@ def __init__(self): self.register_buffer("watermark_image", torch.zeros((62, 62, 4))) self.watermark_image_as_pil = None - def apply_watermark(self, images: List[PIL.Image.Image], sample_size=None): + def apply_watermark(self, images: list[PIL.Image.Image], sample_size=None): # Copied from https://github.com/deep-floyd/IF/blob/b77482e36ca2031cb94dbca1001fc1e6400bf4ab/deepfloyd_if/modules/base.py#L287 h = images[0].height diff --git a/src/diffusers/pipelines/deprecated/alt_diffusion/modeling_roberta_series.py b/src/diffusers/pipelines/deprecated/alt_diffusion/modeling_roberta_series.py index f69f905b56c5..ed72e505b9c3 100644 --- a/src/diffusers/pipelines/deprecated/alt_diffusion/modeling_roberta_series.py +++ b/src/diffusers/pipelines/deprecated/alt_diffusion/modeling_roberta_series.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import Optional, Tuple import torch from torch import nn @@ -18,22 +17,22 @@ class TransformationModelOutput(ModelOutput): last_hidden_state (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model. hidden_states (`tuple(torch.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): - Tuple of `torch.Tensor` (one for the output of the embeddings, if the model has an embedding layer, + one + tuple of `torch.Tensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. attentions (`tuple(torch.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): - Tuple of `torch.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + tuple of `torch.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. """ - projection_state: Optional[torch.Tensor] = None + projection_state: torch.Tensor | None = None last_hidden_state: torch.Tensor = None - hidden_states: Optional[Tuple[torch.Tensor]] = None - attentions: Optional[Tuple[torch.Tensor]] = None + hidden_states: tuple[torch.Tensor] | None = None + attentions: tuple[torch.Tensor] | None = None class RobertaSeriesConfig(XLMRobertaConfig): @@ -73,17 +72,17 @@ def __init__(self, config): def forward( self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - token_type_ids: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - head_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - return_dict: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + head_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + output_hidden_states: bool | None = None, ): r""" """ diff --git a/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py b/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py index 6f484aa3e298..9fab42916e9e 100644 --- a/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py +++ b/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch from packaging import version @@ -93,10 +93,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -111,15 +111,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -299,9 +299,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -331,16 +331,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -348,7 +348,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -447,7 +447,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -704,35 +704,35 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, guidance_rescale: float = 0.0, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -741,14 +741,14 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -756,7 +756,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -790,7 +790,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py b/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py index d6bf90120755..b6cd51c6d203 100644 --- a/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py +++ b/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -82,7 +82,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -121,10 +121,10 @@ def preprocess(image): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -139,15 +139,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -327,9 +327,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -359,16 +359,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -376,7 +376,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -475,7 +475,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -772,35 +772,35 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, image: PipelineImageInput = None, strength: float = 0.8, - num_inference_steps: Optional[int] = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - guidance_scale: Optional[float] = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - eta: Optional[float] = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - output_type: Optional[str] = "pil", + num_inference_steps: int | None = 50, + timesteps: list[int] = None, + sigmas: list[float] = None, + guidance_scale: float | None = 7.5, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + eta: float | None = 0.0, + generator: torch.Generator | list[torch.Generator] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, clip_skip: int = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a @@ -815,14 +815,14 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. This parameter is modulated by `strength`. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -830,7 +830,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. prompt_embeds (`torch.Tensor`, *optional*): @@ -856,7 +856,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_output.py b/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_output.py index dd174ae3c21f..259a86756965 100644 --- a/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_output.py +++ b/src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Optional, Union import numpy as np import PIL.Image @@ -16,13 +15,13 @@ class AltDiffusionPipelineOutput(BaseOutput): Output class for Alt Diffusion pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, num_channels)`. - nsfw_content_detected (`List[bool]`) - List indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or + nsfw_content_detected (`list[bool]`) + list indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or `None` if safety checking could not be performed. """ - images: Union[List[PIL.Image.Image], np.ndarray] - nsfw_content_detected: Optional[List[bool]] + images: list[PIL.Image.Image] | np.ndarray + nsfw_content_detected: list[bool] | None diff --git a/src/diffusers/pipelines/deprecated/audio_diffusion/pipeline_audio_diffusion.py b/src/diffusers/pipelines/deprecated/audio_diffusion/pipeline_audio_diffusion.py index 81fa999eb1fb..f63fc8aacbc8 100644 --- a/src/diffusers/pipelines/deprecated/audio_diffusion/pipeline_audio_diffusion.py +++ b/src/diffusers/pipelines/deprecated/audio_diffusion/pipeline_audio_diffusion.py @@ -14,7 +14,6 @@ from math import acos, sin -from typing import List, Tuple, Union import numpy as np import torch @@ -53,7 +52,7 @@ def __init__( vqvae: AutoencoderKL, unet: UNet2DConditionModel, mel: Mel, - scheduler: Union[DDIMScheduler, DDPMScheduler], + scheduler: DDIMScheduler | DDPMScheduler, ): super().__init__() self.register_modules(unet=unet, scheduler=scheduler, mel=mel, vqvae=vqvae) @@ -84,10 +83,7 @@ def __call__( noise: torch.Tensor = None, encoding: torch.Tensor = None, return_dict=True, - ) -> Union[ - Union[AudioPipelineOutput, ImagePipelineOutput], - Tuple[List[Image.Image], Tuple[int, List[np.ndarray]]], - ]: + ) -> AudioPipelineOutput | ImagePipelineOutput | tuple[list[Image.Image], tuple[int, list[np.ndarray]]]: """ The call function to the pipeline for generation. @@ -170,8 +166,8 @@ def __call__( ``` Returns: - `List[PIL Image]`: - A list of Mel spectrograms (`float`, `List[np.ndarray]`) with the sample rate and raw audio. + `list[PIL Image]`: + A list of Mel spectrograms (`float`, `list[np.ndarray]`) with the sample rate and raw audio. """ steps = steps or self.get_default_steps() @@ -268,13 +264,13 @@ def __call__( return BaseOutput(**AudioPipelineOutput(np.array(audios)[:, np.newaxis, :]), **ImagePipelineOutput(images)) @torch.no_grad() - def encode(self, images: List[Image.Image], steps: int = 50) -> np.ndarray: + def encode(self, images: list[Image.Image], steps: int = 50) -> np.ndarray: """ Reverse the denoising step process to recover a noisy image from the generated image. Args: - images (`List[PIL Image]`): - List of images to encode. + images (`list[PIL Image]`): + list of images to encode. steps (`int`): Number of encoding steps to perform (defaults to `50`). diff --git a/src/diffusers/pipelines/deprecated/latent_diffusion_uncond/pipeline_latent_diffusion_uncond.py b/src/diffusers/pipelines/deprecated/latent_diffusion_uncond/pipeline_latent_diffusion_uncond.py index 0bb24ed0b1ce..70a65e2ef5be 100644 --- a/src/diffusers/pipelines/deprecated/latent_diffusion_uncond/pipeline_latent_diffusion_uncond.py +++ b/src/diffusers/pipelines/deprecated/latent_diffusion_uncond/pipeline_latent_diffusion_uncond.py @@ -13,7 +13,6 @@ # limitations under the License. import inspect -from typing import List, Optional, Tuple, Union import torch @@ -47,13 +46,13 @@ def __init__(self, vqvae: VQModel, unet: UNet2DModel, scheduler: DDIMScheduler): def __call__( self, batch_size: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + generator: torch.Generator | list[torch.Generator] | None = None, eta: float = 0.0, num_inference_steps: int = 50, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, **kwargs, - ) -> Union[Tuple, ImagePipelineOutput]: + ) -> tuple | ImagePipelineOutput: r""" The call function to the pipeline for generation. diff --git a/src/diffusers/pipelines/deprecated/pndm/pipeline_pndm.py b/src/diffusers/pipelines/deprecated/pndm/pipeline_pndm.py index 71e3e156e0e4..fb116511f727 100644 --- a/src/diffusers/pipelines/deprecated/pndm/pipeline_pndm.py +++ b/src/diffusers/pipelines/deprecated/pndm/pipeline_pndm.py @@ -13,8 +13,6 @@ # limitations under the License. -from typing import List, Optional, Tuple, Union - import torch from ....models import UNet2DModel @@ -52,11 +50,11 @@ def __call__( self, batch_size: int = 1, num_inference_steps: int = 50, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + output_type: str | None = "pil", return_dict: bool = True, **kwargs, - ) -> Union[ImagePipelineOutput, Tuple]: + ) -> ImagePipelineOutput | tuple: r""" The call function to the pipeline for generation. diff --git a/src/diffusers/pipelines/deprecated/repaint/pipeline_repaint.py b/src/diffusers/pipelines/deprecated/repaint/pipeline_repaint.py index 56c6007ae886..3231d5e13049 100644 --- a/src/diffusers/pipelines/deprecated/repaint/pipeline_repaint.py +++ b/src/diffusers/pipelines/deprecated/repaint/pipeline_repaint.py @@ -13,8 +13,6 @@ # limitations under the License. -from typing import List, Optional, Tuple, Union - import numpy as np import PIL.Image import torch @@ -30,7 +28,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.preprocess -def _preprocess_image(image: Union[List, PIL.Image.Image, torch.Tensor]): +def _preprocess_image(image: list | PIL.Image.Image | torch.Tensor): deprecation_message = "The preprocess method is deprecated and will be removed in diffusers 1.0.0. Please use VaeImageProcessor.preprocess(...) instead" deprecate("preprocess", "1.0.0", deprecation_message, standard_warn=False) if isinstance(image, torch.Tensor): @@ -53,7 +51,7 @@ def _preprocess_image(image: Union[List, PIL.Image.Image, torch.Tensor]): return image -def _preprocess_mask(mask: Union[List, PIL.Image.Image, torch.Tensor]): +def _preprocess_mask(mask: list | PIL.Image.Image | torch.Tensor): if isinstance(mask, torch.Tensor): return mask elif isinstance(mask, PIL.Image.Image): @@ -98,16 +96,16 @@ def __init__(self, unet: UNet2DModel, scheduler: RePaintScheduler): @torch.no_grad() def __call__( self, - image: Union[torch.Tensor, PIL.Image.Image], - mask_image: Union[torch.Tensor, PIL.Image.Image], + image: torch.Tensor | PIL.Image.Image, + mask_image: torch.Tensor | PIL.Image.Image, num_inference_steps: int = 250, eta: float = 0.0, jump_length: int = 10, jump_n_sample: int = 10, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - ) -> Union[ImagePipelineOutput, Tuple]: + ) -> ImagePipelineOutput | tuple: r""" The call function to the pipeline for generation. diff --git a/src/diffusers/pipelines/deprecated/score_sde_ve/pipeline_score_sde_ve.py b/src/diffusers/pipelines/deprecated/score_sde_ve/pipeline_score_sde_ve.py index 3f04db7ad699..c6abdba42d3c 100644 --- a/src/diffusers/pipelines/deprecated/score_sde_ve/pipeline_score_sde_ve.py +++ b/src/diffusers/pipelines/deprecated/score_sde_ve/pipeline_score_sde_ve.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Optional, Tuple, Union - import torch from ....models import UNet2DModel @@ -48,11 +46,11 @@ def __call__( self, batch_size: int = 1, num_inference_steps: int = 2000, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + output_type: str | None = "pil", return_dict: bool = True, **kwargs, - ) -> Union[ImagePipelineOutput, Tuple]: + ) -> ImagePipelineOutput | tuple: r""" The call function to the pipeline for generation. diff --git a/src/diffusers/pipelines/deprecated/spectrogram_diffusion/midi_utils.py b/src/diffusers/pipelines/deprecated/spectrogram_diffusion/midi_utils.py index 8985a6f88800..76b8576468d2 100644 --- a/src/diffusers/pipelines/deprecated/spectrogram_diffusion/midi_utils.py +++ b/src/diffusers/pipelines/deprecated/spectrogram_diffusion/midi_utils.py @@ -16,7 +16,7 @@ import dataclasses import math import os -from typing import Any, Callable, List, Mapping, MutableMapping, Optional, Sequence, Tuple, Union +from typing import Any, Callable, Mapping, MutableMapping, Sequence import numpy as np import torch @@ -91,10 +91,10 @@ class NoteRepresentationConfig: @dataclasses.dataclass class NoteEventData: pitch: int - velocity: Optional[int] = None - program: Optional[int] = None - is_drum: Optional[bool] = None - instrument: Optional[int] = None + velocity: int | None = None + program: int | None = None + is_drum: bool | None = None + instrument: int | None = None @dataclasses.dataclass @@ -102,7 +102,7 @@ class NoteEncodingState: """Encoding state for note transcription, keeping track of active pitches.""" # velocity bin for active pitches and programs - active_pitches: MutableMapping[Tuple[int, int], int] = dataclasses.field(default_factory=dict) + active_pitches: MutableMapping[tuple[int, int], int] = dataclasses.field(default_factory=dict) @dataclasses.dataclass @@ -153,7 +153,7 @@ class Codec: and specified separately. """ - def __init__(self, max_shift_steps: int, steps_per_second: float, event_ranges: List[EventRange]): + def __init__(self, max_shift_steps: int, steps_per_second: float, event_ranges: list[EventRange]): """Define Codec. Args: @@ -197,7 +197,7 @@ def encode_event(self, event: Event) -> int: raise ValueError(f"Unknown event type: {event.type}") - def event_type_range(self, event_type: str) -> Tuple[int, int]: + def event_type_range(self, event_type: str) -> tuple[int, int]: """Return [min_id, max_id] for an event type.""" offset = 0 for er in self._event_ranges: @@ -280,7 +280,7 @@ def audio_to_frames( samples, hop_size: int, frame_rate: int, -) -> Tuple[Sequence[Sequence[int]], torch.Tensor]: +) -> tuple[Sequence[Sequence[int]], torch.Tensor]: """Convert audio samples to non-overlapping frames and frame times.""" frame_size = hop_size samples = np.pad(samples, [0, frame_size - len(samples) % frame_size], mode="constant") @@ -301,7 +301,7 @@ def audio_to_frames( def note_sequence_to_onsets_and_offsets_and_programs( ns: note_seq.NoteSequence, -) -> Tuple[Sequence[float], Sequence[NoteEventData]]: +) -> tuple[Sequence[float], Sequence[NoteEventData]]: """Extract onset & offset times and pitches & programs from a NoteSequence. The onset & offset times will not necessarily be in sorted order. @@ -348,7 +348,7 @@ def velocity_to_bin(velocity, num_velocity_bins): def note_event_data_to_events( - state: Optional[NoteEncodingState], + state: NoteEncodingState | None, value: NoteEventData, codec: Codec, ) -> Sequence[Event]: @@ -632,7 +632,7 @@ def __init__(self): self.tokenizer = Tokenizer(self.codec.num_classes) self.note_representation_config = NoteRepresentationConfig(onsets_only=False, include_ties=True) - def __call__(self, midi: Union[bytes, os.PathLike, str]): + def __call__(self, midi: bytes | os.PathLike | str): if not isinstance(midi, bytes): with open(midi, "rb") as f: midi = f.read() diff --git a/src/diffusers/pipelines/deprecated/spectrogram_diffusion/pipeline_spectrogram_diffusion.py b/src/diffusers/pipelines/deprecated/spectrogram_diffusion/pipeline_spectrogram_diffusion.py index be07b1b15ea8..269e7405d10d 100644 --- a/src/diffusers/pipelines/deprecated/spectrogram_diffusion/pipeline_spectrogram_diffusion.py +++ b/src/diffusers/pipelines/deprecated/spectrogram_diffusion/pipeline_spectrogram_diffusion.py @@ -14,7 +14,7 @@ # limitations under the License. import math -from typing import Any, Callable, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import torch @@ -129,14 +129,14 @@ def decode(self, encodings_and_masks, input_tokens, noise_time): @torch.no_grad() def __call__( self, - input_tokens: List[List[int]], - generator: Optional[torch.Generator] = None, + input_tokens: list[list[int]], + generator: torch.Generator | None = None, num_inference_steps: int = 100, return_dict: bool = True, output_type: str = "np", - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - ) -> Union[AudioPipelineOutput, Tuple]: + ) -> AudioPipelineOutput | tuple: if (callback_steps is None) or ( callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0) ): @@ -148,8 +148,8 @@ def __call__( The call function to the pipeline for generation. Args: - input_tokens (`List[List[int]]`): - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + input_tokens (`list[list[int]]`): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. num_inference_steps (`int`, *optional*, defaults to 100): diff --git a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py index 08f8c7e26fae..dae5e600d773 100644 --- a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py +++ b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -63,7 +63,7 @@ def preprocess(image): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -259,9 +259,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -292,16 +292,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -309,7 +309,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -408,7 +408,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -633,31 +633,31 @@ def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dt @torch.no_grad() def __call__( self, - prompt: Union[str, List[str]], - source_prompt: Union[str, List[str]], + prompt: str | list[str], + source_prompt: str | list[str], image: PipelineImageInput = None, strength: float = 0.8, - num_inference_steps: Optional[int] = 50, - guidance_scale: Optional[float] = 7.5, - source_guidance_scale: Optional[float] = 1, - num_images_per_prompt: Optional[int] = 1, - eta: Optional[float] = 0.1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + num_inference_steps: int | None = 50, + guidance_scale: float | None = 7.5, + source_guidance_scale: float | None = 1, + num_images_per_prompt: int | None = 1, + eta: float | None = 0.1, + generator: torch.Generator | list[torch.Generator] | None = None, + prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, + cross_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. - image (`torch.Tensor` `np.ndarray`, `PIL.Image.Image`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor` `np.ndarray`, `PIL.Image.Image`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image` or tensor representing an image batch to be used as the starting point. Can also accept image latents as `image`, but if passing latents directly it is not encoded again. strength (`float`, *optional*, defaults to 0.8): @@ -680,7 +680,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. prompt_embeds (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_onnx_stable_diffusion_inpaint_legacy.py b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_onnx_stable_diffusion_inpaint_legacy.py index fcd8bf317adf..f526dc419cea 100644 --- a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_onnx_stable_diffusion_inpaint_legacy.py +++ b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_onnx_stable_diffusion_inpaint_legacy.py @@ -1,5 +1,5 @@ import inspect -from typing import Callable, List, Optional, Union +from typing import Callable import numpy as np import PIL.Image @@ -76,7 +76,7 @@ class OnnxStableDiffusionInpaintPipelineLegacy(DiffusionPipeline): text_encoder: OnnxRuntimeModel tokenizer: CLIPTokenizer unet: OnnxRuntimeModel - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler] + scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler safety_checker: OnnxRuntimeModel feature_extractor: CLIPImageProcessor @@ -87,7 +87,7 @@ def __init__( text_encoder: OnnxRuntimeModel, tokenizer: CLIPTokenizer, unet: OnnxRuntimeModel, - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], + scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler, safety_checker: OnnxRuntimeModel, feature_extractor: CLIPImageProcessor, requires_safety_checker: bool = True, @@ -152,24 +152,24 @@ def __init__( # Copied from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion.OnnxStableDiffusionPipeline._encode_prompt def _encode_prompt( self, - prompt: Union[str, List[str]], - num_images_per_prompt: Optional[int], + prompt: str | list[str], + num_images_per_prompt: int | None, do_classifier_free_guidance: bool, - negative_prompt: Optional[str], - prompt_embeds: Optional[np.ndarray] = None, - negative_prompt_embeds: Optional[np.ndarray] = None, + negative_prompt: str | None, + prompt_embeds: np.ndarray | None = None, + negative_prompt_embeds: np.ndarray | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): prompt to be encoded num_images_per_prompt (`int`): number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`): + negative_prompt (`str` or `list[str]`): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). prompt_embeds (`np.ndarray`, *optional*): @@ -214,7 +214,7 @@ def _encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -297,28 +297,28 @@ def check_inputs( def __call__( self, - prompt: Union[str, List[str]], - image: Union[np.ndarray, PIL.Image.Image] = None, - mask_image: Union[np.ndarray, PIL.Image.Image] = None, + prompt: str | list[str], + image: np.ndarray | PIL.Image.Image = None, + mask_image: np.ndarray | PIL.Image.Image = None, strength: float = 0.8, - num_inference_steps: Optional[int] = 50, - guidance_scale: Optional[float] = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - eta: Optional[float] = 0.0, - generator: Optional[np.random.RandomState] = None, - prompt_embeds: Optional[np.ndarray] = None, - negative_prompt_embeds: Optional[np.ndarray] = None, - output_type: Optional[str] = "pil", + num_inference_steps: int | None = 50, + guidance_scale: float | None = 7.5, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + eta: float | None = 0.0, + generator: np.random.RandomState | None = None, + prompt_embeds: np.ndarray | None = None, + negative_prompt_embeds: np.ndarray | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, np.ndarray], None]] = None, + callback: Callable[[int, int, np.ndarray], None] | None = None, callback_steps: int = 1, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. image (`nd.ndarray` or `PIL.Image.Image`): `Image`, or tensor representing an image batch, that will be used as the starting point for the @@ -343,7 +343,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): diff --git a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py index ba0dd66c2938..650695b604c1 100644 --- a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py +++ b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -230,9 +230,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -263,16 +263,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -280,7 +280,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -379,7 +379,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -562,31 +562,31 @@ def prepare_latents(self, image, timestep, num_images_per_prompt, dtype, device, @torch.no_grad() def __call__( self, - prompt: Union[str, List[str]] = None, - image: Union[torch.Tensor, PIL.Image.Image] = None, - mask_image: Union[torch.Tensor, PIL.Image.Image] = None, + prompt: str | list[str] = None, + image: torch.Tensor | PIL.Image.Image = None, + mask_image: torch.Tensor | PIL.Image.Image = None, strength: float = 0.8, - num_inference_steps: Optional[int] = 50, - guidance_scale: Optional[float] = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - add_predicted_noise: Optional[bool] = False, - eta: Optional[float] = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + num_inference_steps: int | None = 50, + guidance_scale: float | None = 7.5, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + add_predicted_noise: bool | None = False, + eta: float | None = 0.0, + generator: torch.Generator | list[torch.Generator] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, + cross_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. image (`torch.Tensor` or `PIL.Image.Image`): @@ -611,7 +611,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`. instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). diff --git a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py index b7a0be57c12b..851820c00aed 100644 --- a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py +++ b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py @@ -13,7 +13,7 @@ import copy import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer @@ -163,9 +163,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -196,16 +196,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -213,7 +213,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -312,7 +312,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -614,30 +614,30 @@ def edit_model( @torch.no_grad() def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, + cross_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -649,7 +649,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -657,7 +657,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py index c236e73bf448..ea81be87a0f4 100644 --- a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py +++ b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer @@ -160,9 +160,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -193,16 +193,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -210,7 +210,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -309,7 +309,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -488,25 +488,25 @@ def _cumsum(self, input, dim, debug=False): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, parallel: int = 10, tolerance: float = 0.1, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, debug: bool = False, clip_skip: int = None, ): @@ -514,7 +514,7 @@ def __call__( The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -533,7 +533,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -541,7 +541,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py index 2a461ae20cc9..0955b6fe48a1 100644 --- a/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py +++ b/src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py @@ -14,7 +14,7 @@ import inspect from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -62,13 +62,13 @@ class Pix2PixInversionPipelineOutput(BaseOutput, TextualInversionLoaderMixin): Args: latents (`torch.Tensor`) inverted latents tensor - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ latents: torch.Tensor - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image] | np.ndarray EXAMPLE_DOC_STRING = """ @@ -328,7 +328,7 @@ def __init__( text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, unet: UNet2DConditionModel, - scheduler: Union[DDPMScheduler, DDIMScheduler, EulerAncestralDiscreteScheduler, LMSDiscreteScheduler], + scheduler: DDPMScheduler | DDIMScheduler | EulerAncestralDiscreteScheduler | LMSDiscreteScheduler, feature_extractor: CLIPImageProcessor, safety_checker: StableDiffusionSafetyChecker, inverse_scheduler: DDIMInverseScheduler, @@ -378,9 +378,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -411,16 +411,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -428,7 +428,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -527,7 +527,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -709,7 +709,7 @@ def construct_direction(self, embs_source: torch.Tensor, embs_target: torch.Tens return (embs_target.mean(0) - embs_source.mean(0)).unsqueeze(0) @torch.no_grad() - def get_embeds(self, prompt: List[str], batch_size: int = 16) -> torch.Tensor: + def get_embeds(self, prompt: list[str], batch_size: int = 16) -> torch.Tensor: num_prompts = len(prompt) embeds = [] for i in range(0, num_prompts, batch_size): @@ -818,33 +818,33 @@ def kl_divergence(self, hidden_states): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] | None = None, source_embeds: torch.Tensor = None, target_embeds: torch.Tensor = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, cross_attention_guidance_amount: float = 0.1, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: Optional[int] = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, + callback_steps: int | None = 1, + cross_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. source_embeds (`torch.Tensor`): @@ -866,7 +866,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -875,7 +875,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -1109,19 +1109,19 @@ def __call__( @replace_example_docstring(EXAMPLE_INVERT_DOC_STRING) def invert( self, - prompt: Optional[str] = None, + prompt: str | None = None, image: PipelineImageInput = None, num_inference_steps: int = 50, guidance_scale: float = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, cross_attention_guidance_amount: float = 0.1, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: Optional[int] = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, + callback_steps: int | None = 1, + cross_attention_kwargs: dict[str, Any] | None = None, lambda_auto_corr: float = 20.0, lambda_kl: float = 20.0, num_reg_steps: int = 5, @@ -1131,10 +1131,10 @@ def invert( Function used to generate inverted latents given a prompt and image. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - image (`torch.Tensor` `np.ndarray`, `PIL.Image.Image`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor` `np.ndarray`, `PIL.Image.Image`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, or tensor representing an image batch which will be used for conditioning. Can also accept image latents as `image`, if passing latents directly, it will not be encoded again. num_inference_steps (`int`, *optional*, defaults to 50): @@ -1146,7 +1146,7 @@ def invert( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/deprecated/stochastic_karras_ve/pipeline_stochastic_karras_ve.py b/src/diffusers/pipelines/deprecated/stochastic_karras_ve/pipeline_stochastic_karras_ve.py index 50b8b0bcbc1d..ce2c785c8d98 100644 --- a/src/diffusers/pipelines/deprecated/stochastic_karras_ve/pipeline_stochastic_karras_ve.py +++ b/src/diffusers/pipelines/deprecated/stochastic_karras_ve/pipeline_stochastic_karras_ve.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Optional, Tuple, Union - import torch from ....models import UNet2DModel @@ -46,11 +44,11 @@ def __call__( self, batch_size: int = 1, num_inference_steps: int = 50, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + output_type: str | None = "pil", return_dict: bool = True, **kwargs, - ) -> Union[Tuple, ImagePipelineOutput]: + ) -> tuple | ImagePipelineOutput: r""" The call function to the pipeline for generation. diff --git a/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py b/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py index 7c25713cd1d7..7be159d77af5 100644 --- a/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py +++ b/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import numpy as np import torch @@ -277,7 +277,7 @@ class UNetFlatConditionModel(ModelMixin, ConfigMixin): for all models (such as downloading or saving). Parameters: - sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`): + sample_size (`int` or `tuple[int, int]`, *optional*, defaults to `None`): Height and width of input/output sample. in_channels (`int`, *optional*, defaults to 4): Number of channels in the input sample. out_channels (`int`, *optional*, defaults to 4): Number of channels in the output. @@ -285,17 +285,17 @@ class UNetFlatConditionModel(ModelMixin, ConfigMixin): flip_sin_to_cos (`bool`, *optional*, defaults to `False`): Whether to flip the sin to cos in the time embedding. freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding. - down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlockFlat", "CrossAttnDownBlockFlat", "CrossAttnDownBlockFlat", "DownBlockFlat")`): + down_block_types (`tuple[str]`, *optional*, defaults to `("CrossAttnDownBlockFlat", "CrossAttnDownBlockFlat", "CrossAttnDownBlockFlat", "DownBlockFlat")`): The tuple of downsample blocks to use. mid_block_type (`str`, *optional*, defaults to `"UNetMidBlockFlatCrossAttn"`): Block type for middle of UNet, it can be one of `UNetMidBlockFlatCrossAttn`, `UNetMidBlockFlat`, or `UNetMidBlockFlatSimpleCrossAttn`. If `None`, the mid block layer is skipped. - up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlockFlat", "CrossAttnUpBlockFlat", "CrossAttnUpBlockFlat", "CrossAttnUpBlockFlat")`): + up_block_types (`tuple[str]`, *optional*, defaults to `("UpBlockFlat", "CrossAttnUpBlockFlat", "CrossAttnUpBlockFlat", "CrossAttnUpBlockFlat")`): The tuple of upsample blocks to use. - only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`): + only_cross_attention(`bool` or `tuple[bool]`, *optional*, default to `False`): Whether to include self-attention in the basic transformer blocks, see [`~models.attention.BasicTransformerBlock`]. - block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): + block_out_channels (`tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): The tuple of output channels for each block. layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block. downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution. @@ -305,15 +305,15 @@ class UNetFlatConditionModel(ModelMixin, ConfigMixin): norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization. If `None`, normalization and activation layers is skipped in post-processing. norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization. - cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280): + cross_attention_dim (`int` or `tuple[int]`, *optional*, defaults to 1280): The dimension of the cross attention features. - transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple]` , *optional*, defaults to 1): + transformer_layers_per_block (`int`, `tuple[int]`, or `tuple[tuple]` , *optional*, defaults to 1): The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for [`~models.unet_2d_blocks.CrossAttnDownBlockFlat`], [`~models.unet_2d_blocks.CrossAttnUpBlockFlat`], [`~models.unet_2d_blocks.UNetMidBlockFlatCrossAttn`]. - reverse_transformer_layers_per_block : (`Tuple[Tuple]`, *optional*, defaults to None): + reverse_transformer_layers_per_block : (`tuple[tuple]`, *optional*, defaults to None): The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`], in the upsampling - blocks of the U-Net. Only relevant if `transformer_layers_per_block` is of type `Tuple[Tuple]` and for + blocks of the U-Net. Only relevant if `transformer_layers_per_block` is of type `tuple[tuple]` and for [`~models.unet_2d_blocks.CrossAttnDownBlockFlat`], [`~models.unet_2d_blocks.CrossAttnUpBlockFlat`], [`~models.unet_2d_blocks.UNetMidBlockFlatCrossAttn`]. encoder_hid_dim (`int`, *optional*, defaults to None): @@ -368,63 +368,63 @@ class conditioning with `class_embed_type` equal to `None`. @register_to_config def __init__( self, - sample_size: Optional[int] = None, + sample_size: int | None = None, in_channels: int = 4, out_channels: int = 4, center_input_sample: bool = False, flip_sin_to_cos: bool = True, freq_shift: int = 0, - down_block_types: Tuple[str, ...] = ( + down_block_types: tuple[str] = ( "CrossAttnDownBlockFlat", "CrossAttnDownBlockFlat", "CrossAttnDownBlockFlat", "DownBlockFlat", ), - mid_block_type: Optional[str] = "UNetMidBlockFlatCrossAttn", - up_block_types: Tuple[str, ...] = ( + mid_block_type: str = "UNetMidBlockFlatCrossAttn", + up_block_types: tuple[str] = ( "UpBlockFlat", "CrossAttnUpBlockFlat", "CrossAttnUpBlockFlat", "CrossAttnUpBlockFlat", ), - only_cross_attention: Union[bool, Tuple[bool]] = False, - block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280), - layers_per_block: Union[int, Tuple[int]] = 2, + only_cross_attention: bool | tuple[bool] = False, + block_out_channels: tuple[int] = (320, 640, 1280, 1280), + layers_per_block: int | tuple[int] = 2, downsample_padding: int = 1, mid_block_scale_factor: float = 1, dropout: float = 0.0, act_fn: str = "silu", - norm_num_groups: Optional[int] = 32, + norm_num_groups: int | None = 32, norm_eps: float = 1e-5, - cross_attention_dim: Union[int, Tuple[int]] = 1280, - transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1, - reverse_transformer_layers_per_block: Optional[Tuple[Tuple[int]]] = None, - encoder_hid_dim: Optional[int] = None, - encoder_hid_dim_type: Optional[str] = None, - attention_head_dim: Union[int, Tuple[int]] = 8, - num_attention_heads: Optional[Union[int, Tuple[int]]] = None, + cross_attention_dim: int | tuple[int] = 1280, + transformer_layers_per_block: int | tuple[int] | tuple[tuple] = 1, + reverse_transformer_layers_per_block: tuple[tuple[int]] | None = None, + encoder_hid_dim: int | None = None, + encoder_hid_dim_type: str | None = None, + attention_head_dim: int | tuple[int] = 8, + num_attention_heads: int | tuple[int] | None = None, dual_cross_attention: bool = False, use_linear_projection: bool = False, - class_embed_type: Optional[str] = None, - addition_embed_type: Optional[str] = None, - addition_time_embed_dim: Optional[int] = None, - num_class_embeds: Optional[int] = None, + class_embed_type: str | None = None, + addition_embed_type: str | None = None, + addition_time_embed_dim: int | None = None, + num_class_embeds: int | None = None, upcast_attention: bool = False, resnet_time_scale_shift: str = "default", resnet_skip_time_act: bool = False, resnet_out_scale_factor: int = 1.0, time_embedding_type: str = "positional", - time_embedding_dim: Optional[int] = None, - time_embedding_act_fn: Optional[str] = None, - timestep_post_act: Optional[str] = None, - time_cond_proj_dim: Optional[int] = None, + time_embedding_dim: int | None = None, + time_embedding_act_fn: str | None = None, + timestep_post_act: str | None = None, + time_cond_proj_dim: int | None = None, conv_in_kernel: int = 3, conv_out_kernel: int = 3, - projection_class_embeddings_input_dim: Optional[int] = None, + projection_class_embeddings_input_dim: int | None = None, attention_type: str = "default", class_embeddings_concat: bool = False, - mid_block_only_cross_attention: Optional[bool] = None, - cross_attention_norm: Optional[str] = None, + mid_block_only_cross_attention: bool | None = None, + cross_attention_norm: str | None = None, addition_embed_type_num_heads=64, ): super().__init__() @@ -826,7 +826,7 @@ def __init__( ) @property - def attn_processors(self) -> Dict[str, AttentionProcessor]: + def attn_processors(self) -> dict[str, AttentionProcessor]: r""" Returns: `dict` of attention processors: A dictionary containing all attention processors used in the model with @@ -835,7 +835,7 @@ def attn_processors(self) -> Dict[str, AttentionProcessor]: # set recursively processors = {} - def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]): + def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: dict[str, AttentionProcessor]): if hasattr(module, "get_processor"): processors[f"{name}.processor"] = module.get_processor() @@ -849,7 +849,7 @@ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: return processors - def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): + def set_attn_processor(self, processor: AttentionProcessor | dict[str, AttentionProcessor]): r""" Sets the attention processor to use to compute attention. @@ -952,7 +952,7 @@ def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): # Recursively walk through all the children. # Any children which exposes the set_attention_slice method # gets the message - def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]): + def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: list[int]): if hasattr(module, "set_attention_slice"): module.set_attention_slice(slice_size.pop()) @@ -1037,19 +1037,19 @@ def unload_lora(self): def forward( self, sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], + timestep: torch.Tensor | float | int, encoder_hidden_states: torch.Tensor, - class_labels: Optional[torch.Tensor] = None, - timestep_cond: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, - down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None, - mid_block_additional_residual: Optional[torch.Tensor] = None, - down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, + class_labels: torch.Tensor | None = None, + timestep_cond: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + added_cond_kwargs: dict[str, torch.Tensor] | None = None, + down_block_additional_residuals: tuple[torch.Tensor] | None = None, + mid_block_additional_residual: torch.Tensor | None = None, + down_intrablock_additional_residuals: tuple[torch.Tensor] | None = None, + encoder_attention_mask: torch.Tensor | None = None, return_dict: bool = True, - ) -> Union[UNet2DConditionOutput, Tuple]: + ) -> UNet2DConditionOutput | tuple: r""" The [`UNetFlatConditionModel`] forward method. @@ -1579,8 +1579,8 @@ def __init__( self.gradient_checkpointing = False def forward( - self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None - ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]: + self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None + ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]: output_states = () for resnet in self.resnets: @@ -1608,7 +1608,7 @@ def __init__( temb_channels: int, dropout: float = 0.0, num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, + transformer_layers_per_block: int | tuple[int] = 1, resnet_eps: float = 1e-6, resnet_time_scale_shift: str = "default", resnet_act_fn: str = "swish", @@ -1695,13 +1695,13 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - additional_residuals: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]: + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + encoder_attention_mask: torch.Tensor | None = None, + additional_residuals: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]: output_states = () blocks = list(zip(self.resnets, self.attentions)) @@ -1751,7 +1751,7 @@ def __init__( prev_output_channel: int, out_channels: int, temb_channels: int, - resolution_idx: Optional[int] = None, + resolution_idx: int | None = None, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, @@ -1797,9 +1797,9 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - upsample_size: Optional[int] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, + upsample_size: int | None = None, *args, **kwargs, ) -> torch.Tensor: @@ -1853,10 +1853,10 @@ def __init__( out_channels: int, prev_output_channel: int, temb_channels: int, - resolution_idx: Optional[int] = None, + resolution_idx: int | None = None, dropout: float = 0.0, num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, + transformer_layers_per_block: int | tuple[int] = 1, resnet_eps: float = 1e-6, resnet_time_scale_shift: str = "default", resnet_act_fn: str = "swish", @@ -1940,13 +1940,13 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - res_hidden_states_tuple: Tuple[torch.Tensor, ...], - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - upsample_size: Optional[int] = None, - attention_mask: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + upsample_size: int | None = None, + attention_mask: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: if cross_attention_kwargs is not None: if cross_attention_kwargs.get("scale", None) is not None: @@ -2023,7 +2023,7 @@ class UNetMidBlockFlat(nn.Module): resnet_act_fn (`str`, *optional*, defaults to `swish`): The activation function for the resnet blocks. resnet_groups (`int`, *optional*, defaults to 32): The number of groups to use in the group normalization layers of the resnet blocks. - attn_groups (`Optional[int]`, *optional*, defaults to None): The number of groups for the attention blocks. + attn_groups (`int | None`, *optional*, defaults to None): The number of groups for the attention blocks. resnet_pre_norm (`bool`, *optional*, defaults to `True`): Whether to use pre-normalization for the resnet blocks. add_attention (`bool`, *optional*, defaults to `True`): Whether to add attention blocks. @@ -2048,7 +2048,7 @@ def __init__( resnet_time_scale_shift: str = "default", # default, spatial resnet_act_fn: str = "swish", resnet_groups: int = 32, - attn_groups: Optional[int] = None, + attn_groups: int | None = None, resnet_pre_norm: bool = True, add_attention: bool = True, attention_head_dim: int = 1, @@ -2154,7 +2154,7 @@ def __init__( self.gradient_checkpointing = False - def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor: hidden_states = self.resnets[0](hidden_states, temb) for attn, resnet in zip(self.attentions, self.resnets[1:]): if torch.is_grad_enabled() and self.gradient_checkpointing: @@ -2175,15 +2175,15 @@ def __init__( self, in_channels: int, temb_channels: int, - out_channels: Optional[int] = None, + out_channels: int | None = None, dropout: float = 0.0, num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, + transformer_layers_per_block: int | tuple[int] = 1, resnet_eps: float = 1e-6, resnet_time_scale_shift: str = "default", resnet_act_fn: str = "swish", resnet_groups: int = 32, - resnet_groups_out: Optional[int] = None, + resnet_groups_out: int | None = None, resnet_pre_norm: bool = True, num_attention_heads: int = 1, output_scale_factor: float = 1.0, @@ -2276,11 +2276,11 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + encoder_attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: if cross_attention_kwargs is not None: if cross_attention_kwargs.get("scale", None) is not None: @@ -2330,7 +2330,7 @@ def __init__( cross_attention_dim: int = 1280, skip_time_act: bool = False, only_cross_attention: bool = False, - cross_attention_norm: Optional[str] = None, + cross_attention_norm: str | None = None, ): super().__init__() @@ -2401,11 +2401,11 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - temb: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, + temb: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + encoder_attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} if cross_attention_kwargs.get("scale", None) is not None: diff --git a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py index 9ff8e9857791..101a1b72e7f9 100644 --- a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py +++ b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py @@ -1,5 +1,5 @@ import inspect -from typing import Callable, List, Optional, Union +from typing import Callable import PIL.Image import torch @@ -81,26 +81,26 @@ def __init__( @torch.no_grad() def image_variation( self, - image: Union[torch.Tensor, PIL.Image.Image], - height: Optional[int] = None, - width: Optional[int] = None, + image: torch.Tensor | PIL.Image.Image, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, ): r""" The call function to the pipeline for generation. Args: - image (`PIL.Image.Image`, `List[PIL.Image.Image]` or `torch.Tensor`): + image (`PIL.Image.Image`, `list[PIL.Image.Image]` or `torch.Tensor`): The image prompt or prompts to guide the image generation. height (`int`, *optional*, defaults to `self.image_unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -112,7 +112,7 @@ def image_variation( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -193,26 +193,26 @@ def image_variation( @torch.no_grad() def text_to_image( self, - prompt: Union[str, List[str]], - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str], + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide image generation. height (`int`, *optional*, defaults to `self.image_unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -224,7 +224,7 @@ def text_to_image( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -301,27 +301,27 @@ def text_to_image( @torch.no_grad() def dual_guided( self, - prompt: Union[PIL.Image.Image, List[PIL.Image.Image]], - image: Union[str, List[str]], + prompt: PIL.Image.Image | list[PIL.Image.Image], + image: str | list[str], text_to_image_strength: float = 0.5, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, - num_images_per_prompt: Optional[int] = 1, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide image generation. height (`int`, *optional*, defaults to `self.image_unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -333,7 +333,7 @@ def dual_guided( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -341,7 +341,7 @@ def dual_guided( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_dual_guided.py b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_dual_guided.py index 0252f4f6af7f..8f8fb712e023 100644 --- a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_dual_guided.py +++ b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_dual_guided.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Callable, List, Optional, Tuple, Union +from typing import Callable import numpy as np import PIL.Image @@ -158,7 +158,7 @@ def _encode_text_prompt(self, prompt, device, num_images_per_prompt, do_classifi Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): prompt to be encoded device: (`torch.device`): torch device @@ -249,7 +249,7 @@ def _encode_image_prompt(self, prompt, device, num_images_per_prompt, do_classif Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): prompt to be encoded device: (`torch.device`): torch device @@ -369,7 +369,7 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype latents = latents * self.scheduler.init_noise_sigma return latents - def set_transformer_params(self, mix_ratio: float = 0.5, condition_types: Tuple = ("text", "image")): + def set_transformer_params(self, mix_ratio: float = 0.5, condition_types: tuple = ("text", "image")): for name, module in self.image_unet.named_modules(): if isinstance(module, DualTransformer2DModel): module.mix_ratio = mix_ratio @@ -385,20 +385,20 @@ def set_transformer_params(self, mix_ratio: float = 0.5, condition_types: Tuple @torch.no_grad() def __call__( self, - prompt: Union[PIL.Image.Image, List[PIL.Image.Image]], - image: Union[str, List[str]], + prompt: PIL.Image.Image | list[PIL.Image.Image], + image: str | list[str], text_to_image_strength: float = 0.5, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, - num_images_per_prompt: Optional[int] = 1, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, **kwargs, ): @@ -406,7 +406,7 @@ def __call__( The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide image generation. height (`int`, *optional*, defaults to `self.image_unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -418,7 +418,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -426,7 +426,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py index 034a0226419b..348417ad11df 100644 --- a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py +++ b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Callable, List, Optional, Union +from typing import Callable import numpy as np import PIL.Image @@ -84,7 +84,7 @@ def _encode_prompt(self, prompt, device, num_images_per_prompt, do_classifier_fr Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): prompt to be encoded device: (`torch.device`): torch device @@ -92,7 +92,7 @@ def _encode_prompt(self, prompt, device, num_images_per_prompt, do_classifier_fr number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`): + negative_prompt (`str` or `list[str]`): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). """ @@ -122,7 +122,7 @@ def normalize_embeddings(encoder_output): # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance: - uncond_images: List[str] + uncond_images: list[str] if negative_prompt is None: uncond_images = [np.zeros((512, 512, 3)) + 0.5] * batch_size elif type(prompt) is not type(negative_prompt): @@ -196,7 +196,7 @@ def check_inputs(self, image, height, width, callback_steps): and not isinstance(image, list) ): raise ValueError( - "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is" + "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]` but is" f" {type(image)}" ) @@ -237,19 +237,19 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype @torch.no_grad() def __call__( self, - image: Union[PIL.Image.Image, List[PIL.Image.Image], torch.Tensor], - height: Optional[int] = None, - width: Optional[int] = None, + image: PIL.Image.Image | list[PIL.Image.Image] | torch.Tensor, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, **kwargs, ): @@ -257,7 +257,7 @@ def __call__( The call function to the pipeline for generation. Args: - image (`PIL.Image.Image`, `List[PIL.Image.Image]` or `torch.Tensor`): + image (`PIL.Image.Image`, `list[PIL.Image.Image]` or `torch.Tensor`): The image prompt or prompts to guide the image generation. height (`int`, *optional*, defaults to `self.image_unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -269,7 +269,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): diff --git a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py index 2f54f4fc98a4..8e2af1063421 100644 --- a/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py +++ b/src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Callable, List, Optional, Union +from typing import Callable import torch from transformers import CLIPImageProcessor, CLIPTextModelWithProjection, CLIPTokenizer @@ -108,7 +108,7 @@ def _encode_prompt(self, prompt, device, num_images_per_prompt, do_classifier_fr Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): prompt to be encoded device: (`torch.device`): torch device @@ -116,7 +116,7 @@ def _encode_prompt(self, prompt, device, num_images_per_prompt, do_classifier_fr number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`): + negative_prompt (`str` or `list[str]`): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). """ @@ -164,7 +164,7 @@ def normalize_embeddings(encoder_output): # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -323,19 +323,19 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype @torch.no_grad() def __call__( self, - prompt: Union[str, List[str]], - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str], + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, **kwargs, ): @@ -343,7 +343,7 @@ def __call__( The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide image generation. height (`int`, *optional*, defaults to `self.image_unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -355,7 +355,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): diff --git a/src/diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py b/src/diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py index e8617a54b691..8f3b9512e888 100644 --- a/src/diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py +++ b/src/diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Callable, List, Optional, Tuple, Union +from typing import Callable import torch from transformers import CLIPTextModel, CLIPTokenizer @@ -33,7 +33,7 @@ class LearnedClassifierFreeSamplingEmbeddings(ModelMixin, ConfigMixin): """ @register_to_config - def __init__(self, learnable: bool, hidden_size: Optional[int] = None, length: Optional[int] = None): + def __init__(self, learnable: bool, hidden_size: int | None = None, length: int | None = None): super().__init__() self.learnable = learnable @@ -163,23 +163,23 @@ def _encode_prompt(self, prompt, num_images_per_prompt, do_classifier_free_guida @torch.no_grad() def __call__( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_inference_steps: int = 100, guidance_scale: float = 5.0, truncation_rate: float = 1.0, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - ) -> Union[ImagePipelineOutput, Tuple]: + ) -> ImagePipelineOutput | tuple: """ The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide image generation. num_inference_steps (`int`, *optional*, defaults to 100): The number of denoising steps. More denoising steps usually lead to a higher quality image at the diff --git a/src/diffusers/pipelines/dit/pipeline_dit.py b/src/diffusers/pipelines/dit/pipeline_dit.py index 68ff6c9b559a..505360fc78a7 100644 --- a/src/diffusers/pipelines/dit/pipeline_dit.py +++ b/src/diffusers/pipelines/dit/pipeline_dit.py @@ -18,8 +18,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Dict, List, Optional, Tuple, Union - import torch from ...models import AutoencoderKL, DiTTransformer2DModel @@ -62,7 +60,7 @@ def __init__( transformer: DiTTransformer2DModel, vae: AutoencoderKL, scheduler: KarrasDiffusionSchedulers, - id2label: Optional[Dict[int, str]] = None, + id2label: dict[int, str] | None = None, ): super().__init__() self.register_modules(transformer=transformer, vae=vae, scheduler=scheduler) @@ -75,7 +73,7 @@ def __init__( self.labels[label.lstrip().rstrip()] = int(key) self.labels = dict(sorted(self.labels.items())) - def get_label_ids(self, label: Union[str, List[str]]) -> List[int]: + def get_label_ids(self, label: str | list[str]) -> list[int]: r""" Map label strings from ImageNet to corresponding class ids. @@ -103,19 +101,19 @@ def get_label_ids(self, label: Union[str, List[str]]) -> List[int]: @torch.no_grad() def __call__( self, - class_labels: List[int], + class_labels: list[int], guidance_scale: float = 4.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + generator: torch.Generator | list[torch.Generator] | None = None, num_inference_steps: int = 50, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, - ) -> Union[ImagePipelineOutput, Tuple]: + ) -> ImagePipelineOutput | tuple: r""" The call function to the pipeline for generation. Args: - class_labels (List[int]): - List of ImageNet class labels for the images to be generated. + class_labels (list[int]): + list of ImageNet class labels for the images to be generated. guidance_scale (`float`, *optional*, defaults to 4.0): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. diff --git a/src/diffusers/pipelines/easyanimate/pipeline_easyanimate.py b/src/diffusers/pipelines/easyanimate/pipeline_easyanimate.py index 92239c0d32f0..6ec8f44e6d1a 100755 --- a/src/diffusers/pipelines/easyanimate/pipeline_easyanimate.py +++ b/src/diffusers/pipelines/easyanimate/pipeline_easyanimate.py @@ -14,7 +14,7 @@ # limitations under the License. import inspect -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import torch from transformers import ( @@ -126,10 +126,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -144,15 +144,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -195,9 +195,9 @@ class EasyAnimatePipeline(DiffusionPipeline): Args: vae ([`AutoencoderKLMagvit`]): Variational Auto-Encoder (VAE) Model to encode and decode video to and from latent representations. - text_encoder (Optional[`~transformers.Qwen2VLForConditionalGeneration`, `~transformers.BertModel`]): + text_encoder (`~transformers.Qwen2VLForConditionalGeneration`, `~transformers.BertModel` | None): EasyAnimate uses [qwen2 vl](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct) in V5.1. - tokenizer (Optional[`~transformers.Qwen2Tokenizer`, `~transformers.BertTokenizer`]): + tokenizer (`~transformers.Qwen2Tokenizer`, `~transformers.BertTokenizer` | None): A `Qwen2Tokenizer` or `BertTokenizer` to tokenize text. transformer ([`EasyAnimateTransformer3DModel`]): The EasyAnimate model designed by EasyAnimate Team. @@ -211,8 +211,8 @@ class EasyAnimatePipeline(DiffusionPipeline): def __init__( self, vae: AutoencoderKLMagvit, - text_encoder: Union[Qwen2VLForConditionalGeneration, BertModel], - tokenizer: Union[Qwen2Tokenizer, BertTokenizer], + text_encoder: Qwen2VLForConditionalGeneration | BertModel, + tokenizer: Qwen2Tokenizer | BertTokenizer, transformer: EasyAnimateTransformer3DModel, scheduler: FlowMatchEulerDiscreteScheduler, ): @@ -240,23 +240,23 @@ def __init__( def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + negative_prompt: str | list[str] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, max_sequence_length: int = 256, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -266,7 +266,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -525,35 +525,33 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - num_frames: Optional[int] = 49, - height: Optional[int] = 512, - width: Optional[int] = 512, - num_inference_steps: Optional[int] = 50, - guidance_scale: Optional[float] = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - eta: Optional[float] = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - timesteps: Optional[List[int]] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + prompt: str | list[str] = None, + num_frames: int | None = 49, + height: int | None = 512, + width: int | None = 512, + num_inference_steps: int | None = 50, + guidance_scale: float | None = 5.0, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + eta: float | None = 0.0, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + timesteps: list[int] | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], guidance_rescale: float = 0.0, ): r""" Generates images or video using the EasyAnimate pipeline based on the provided prompts. Examples: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): Text prompts to guide the image or video generation. If not provided, use `prompt_embeds` instead. num_frames (`int`, *optional*): Length of the generated video (in frames). @@ -566,13 +564,13 @@ def __call__( down inference. guidance_scale (`float`, *optional*, defaults to 5.0): Encourages the model to align outputs with prompts. A higher value may decrease image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): Prompts indicating what to exclude in generation. If not specified, use `negative_prompt_embeds`. num_images_per_prompt (`int`, *optional*, defaults to 1): Number of images to generate for each prompt. eta (`float`, *optional*, defaults to 0.0): Applies to DDIM scheduling. Controlled by the eta parameter from the related literature. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A generator to ensure reproducibility in image generation. latents (`torch.Tensor`, *optional*): Predefined latent tensors to condition generation. @@ -590,15 +588,15 @@ def __call__( If `True`, returns a structured output. Otherwise returns a simple tuple. callback_on_step_end (`Callable`, *optional*): Functions called at the end of each denoising step. - callback_on_step_end_tensor_inputs (`List[str]`, *optional*): + callback_on_step_end_tensor_inputs (`list[str]`, *optional*): Tensor names to be included in callback function calls. guidance_rescale (`float`, *optional*, defaults to 0.0): Adjusts noise levels based on guidance scale. - original_size (`Tuple[int, int]`, *optional*, defaults to `(1024, 1024)`): + original_size (`tuple[int, int]`, *optional*, defaults to `(1024, 1024)`): Original dimensions of the output. - target_size (`Tuple[int, int]`, *optional*): + target_size (`tuple[int, int]`, *optional*): Desired output dimensions for calculations. - crops_coords_top_left (`Tuple[int, int]`, *optional*, defaults to `(0, 0)`): + crops_coords_top_left (`tuple[int, int]`, *optional*, defaults to `(0, 0)`): Coordinates for cropping. Returns: @@ -666,12 +664,18 @@ def __call__( ) # 4. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler): timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, mu=1 + self.scheduler, num_inference_steps, timestep_device, timesteps, mu=1 ) else: - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, timesteps + ) # 5. Prepare latent variables num_channels_latents = self.transformer.config.in_channels diff --git a/src/diffusers/pipelines/easyanimate/pipeline_easyanimate_control.py b/src/diffusers/pipelines/easyanimate/pipeline_easyanimate_control.py index f74a11f87d75..5e07996a661c 100755 --- a/src/diffusers/pipelines/easyanimate/pipeline_easyanimate_control.py +++ b/src/diffusers/pipelines/easyanimate/pipeline_easyanimate_control.py @@ -14,7 +14,7 @@ # limitations under the License. import inspect -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import numpy as np import torch @@ -228,10 +228,10 @@ def resize_mask(mask, latent, process_first_frame_only=True): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -246,15 +246,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -297,9 +297,9 @@ class EasyAnimateControlPipeline(DiffusionPipeline): Args: vae ([`AutoencoderKLMagvit`]): Variational Auto-Encoder (VAE) Model to encode and decode video to and from latent representations. - text_encoder (Optional[`~transformers.Qwen2VLForConditionalGeneration`, `~transformers.BertModel`]): + text_encoder (`~transformers.Qwen2VLForConditionalGeneration`, `~transformers.BertModel` | None): EasyAnimate uses [qwen2 vl](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct) in V5.1. - tokenizer (Optional[`~transformers.Qwen2Tokenizer`, `~transformers.BertTokenizer`]): + tokenizer (`~transformers.Qwen2Tokenizer`, `~transformers.BertTokenizer` | None): A `Qwen2Tokenizer` or `BertTokenizer` to tokenize text. transformer ([`EasyAnimateTransformer3DModel`]): The EasyAnimate model designed by EasyAnimate Team. @@ -313,8 +313,8 @@ class EasyAnimateControlPipeline(DiffusionPipeline): def __init__( self, vae: AutoencoderKLMagvit, - text_encoder: Union[Qwen2VLForConditionalGeneration, BertModel], - tokenizer: Union[Qwen2Tokenizer, BertTokenizer], + text_encoder: Qwen2VLForConditionalGeneration | BertModel, + tokenizer: Qwen2Tokenizer | BertTokenizer, transformer: EasyAnimateTransformer3DModel, scheduler: FlowMatchEulerDiscreteScheduler, ): @@ -351,23 +351,23 @@ def __init__( # Copied from diffusers.pipelines.easyanimate.pipeline_easyanimate.EasyAnimatePipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + negative_prompt: str | list[str] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, max_sequence_length: int = 256, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -377,7 +377,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -671,38 +671,36 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - num_frames: Optional[int] = 49, - height: Optional[int] = 512, - width: Optional[int] = 512, - control_video: Union[torch.FloatTensor] = None, - control_camera_video: Union[torch.FloatTensor] = None, - ref_image: Union[torch.FloatTensor] = None, - num_inference_steps: Optional[int] = 50, - guidance_scale: Optional[float] = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - eta: Optional[float] = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + prompt: str | list[str] = None, + num_frames: int | None = 49, + height: int | None = 512, + width: int | None = 512, + control_video: torch.FloatTensor = None, + control_camera_video: torch.FloatTensor = None, + ref_image: torch.FloatTensor = None, + num_inference_steps: int | None = 50, + guidance_scale: float | None = 5.0, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + eta: float | None = 0.0, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], guidance_rescale: float = 0.0, - timesteps: Optional[List[int]] = None, + timesteps: list[int] | None = None, ): r""" Generates images or video using the EasyAnimate pipeline based on the provided prompts. Examples: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): Text prompts to guide the image or video generation. If not provided, use `prompt_embeds` instead. num_frames (`int`, *optional*): Length of the generated video (in frames). @@ -715,13 +713,13 @@ def __call__( down inference. guidance_scale (`float`, *optional*, defaults to 5.0): Encourages the model to align outputs with prompts. A higher value may decrease image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): Prompts indicating what to exclude in generation. If not specified, use `negative_prompt_embeds`. num_images_per_prompt (`int`, *optional*, defaults to 1): Number of images to generate for each prompt. eta (`float`, *optional*, defaults to 0.0): Applies to DDIM scheduling. Controlled by the eta parameter from the related literature. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A generator to ensure reproducibility in image generation. latents (`torch.Tensor`, *optional*): Predefined latent tensors to condition generation. @@ -739,7 +737,7 @@ def __call__( If `True`, returns a structured output. Otherwise returns a simple tuple. callback_on_step_end (`Callable`, *optional*): Functions called at the end of each denoising step. - callback_on_step_end_tensor_inputs (`List[str]`, *optional*): + callback_on_step_end_tensor_inputs (`list[str]`, *optional*): Tensor names to be included in callback function calls. guidance_rescale (`float`, *optional*, defaults to 0.0): Adjusts noise levels based on guidance scale. @@ -810,12 +808,18 @@ def __call__( ) # 4. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler): timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, mu=1 + self.scheduler, num_inference_steps, timestep_device, timesteps, mu=1 ) else: - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, timesteps + ) timesteps = self.scheduler.timesteps # 5. Prepare latent variables diff --git a/src/diffusers/pipelines/easyanimate/pipeline_easyanimate_inpaint.py b/src/diffusers/pipelines/easyanimate/pipeline_easyanimate_inpaint.py index b16ef92d8e6b..872313898008 100755 --- a/src/diffusers/pipelines/easyanimate/pipeline_easyanimate_inpaint.py +++ b/src/diffusers/pipelines/easyanimate/pipeline_easyanimate_inpaint.py @@ -14,7 +14,7 @@ # limitations under the License. import inspect -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import numpy as np import torch @@ -270,10 +270,10 @@ def add_noise_to_reference_video(image, ratio=None, generator=None): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -288,15 +288,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -339,9 +339,9 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline): Args: vae ([`AutoencoderKLMagvit`]): Variational Auto-Encoder (VAE) Model to encode and decode video to and from latent representations. - text_encoder (Optional[`~transformers.Qwen2VLForConditionalGeneration`, `~transformers.BertModel`]): + text_encoder (`~transformers.Qwen2VLForConditionalGeneration`, `~transformers.BertModel` | None): EasyAnimate uses [qwen2 vl](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct) in V5.1. - tokenizer (Optional[`~transformers.Qwen2Tokenizer`, `~transformers.BertTokenizer`]): + tokenizer (`~transformers.Qwen2Tokenizer`, `~transformers.BertTokenizer` | None): A `Qwen2Tokenizer` or `BertTokenizer` to tokenize text. transformer ([`EasyAnimateTransformer3DModel`]): The EasyAnimate model designed by EasyAnimate Team. @@ -355,8 +355,8 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline): def __init__( self, vae: AutoencoderKLMagvit, - text_encoder: Union[Qwen2VLForConditionalGeneration, BertModel], - tokenizer: Union[Qwen2Tokenizer, BertTokenizer], + text_encoder: Qwen2VLForConditionalGeneration | BertModel, + tokenizer: Qwen2Tokenizer | BertTokenizer, transformer: EasyAnimateTransformer3DModel, scheduler: FlowMatchEulerDiscreteScheduler, ): @@ -393,23 +393,23 @@ def __init__( # Copied from diffusers.pipelines.easyanimate.pipeline_easyanimate.EasyAnimatePipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + negative_prompt: str | list[str] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, max_sequence_length: int = 256, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -419,7 +419,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -789,40 +789,38 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - num_frames: Optional[int] = 49, - video: Union[torch.FloatTensor] = None, - mask_video: Union[torch.FloatTensor] = None, - masked_video_latents: Union[torch.FloatTensor] = None, - height: Optional[int] = 512, - width: Optional[int] = 512, - num_inference_steps: Optional[int] = 50, - guidance_scale: Optional[float] = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - eta: Optional[float] = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + prompt: str | list[str] = None, + num_frames: int | None = 49, + video: torch.FloatTensor = None, + mask_video: torch.FloatTensor = None, + masked_video_latents: torch.FloatTensor = None, + height: int | None = 512, + width: int | None = 512, + num_inference_steps: int | None = 50, + guidance_scale: float | None = 5.0, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + eta: float | None = 0.0, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], guidance_rescale: float = 0.0, strength: float = 1.0, noise_aug_strength: float = 0.0563, - timesteps: Optional[List[int]] = None, + timesteps: list[int] | None = None, ): r""" The call function to the pipeline for generation with HunyuanDiT. Examples: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. num_frames (`int`, *optional*): Length of the video to be generated in seconds. This parameter influences the number of frames and @@ -843,7 +841,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 5.0): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is effective when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to exclude in image generation. If not defined, you need to provide `negative_prompt_embeds`. This parameter is ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -852,7 +850,7 @@ def __call__( A parameter defined in the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`] and is ignored in other schedulers. It adjusts noise level during the inference process. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) for setting random seeds which helps in making generation deterministic. latents (`torch.Tensor`, *optional*): @@ -874,11 +872,11 @@ def __call__( return_dict (`bool`, *optional*, defaults to `True`): If set to `True`, a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] will be returned; otherwise, a tuple containing the generated images and safety flags will be returned. - callback_on_step_end (`Callable[[int, int, Dict], None]`, `PipelineCallback`, `MultiPipelineCallbacks`, + callback_on_step_end (`Callable[[int, int], None]`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): A callback function (or a list of them) that will be executed at the end of each denoising step, allowing for custom processing during generation. - callback_on_step_end_tensor_inputs (`List[str]`, *optional*): + callback_on_step_end_tensor_inputs (`list[str]`, *optional*): Specifies which tensor inputs should be included in the callback function. If not defined, all tensor inputs will be passed, facilitating enhanced logging or monitoring of the generation process. guidance_rescale (`float`, *optional*, defaults to 0.0): @@ -956,12 +954,18 @@ def __call__( ) # 4. set timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler): timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, mu=1 + self.scheduler, num_inference_steps, timestep_device, timesteps, mu=1 ) else: - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, timesteps + ) timesteps, num_inference_steps = self.get_timesteps( num_inference_steps=num_inference_steps, strength=strength, device=device ) diff --git a/src/diffusers/pipelines/easyanimate/pipeline_output.py b/src/diffusers/pipelines/easyanimate/pipeline_output.py index c761a3b1079f..4fd904ae7dfe 100644 --- a/src/diffusers/pipelines/easyanimate/pipeline_output.py +++ b/src/diffusers/pipelines/easyanimate/pipeline_output.py @@ -11,8 +11,8 @@ class EasyAnimatePipelineOutput(BaseOutput): Output class for EasyAnimate pipelines. Args: - frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): - List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): + list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape `(batch_size, num_frames, channels, height, width)`. """ diff --git a/src/diffusers/pipelines/flux/modeling_flux.py b/src/diffusers/pipelines/flux/modeling_flux.py index d7f2f45812b3..916e3fbe5953 100644 --- a/src/diffusers/pipelines/flux/modeling_flux.py +++ b/src/diffusers/pipelines/flux/modeling_flux.py @@ -14,7 +14,6 @@ from dataclasses import dataclass -from typing import Optional import torch import torch.nn as nn @@ -26,7 +25,7 @@ @dataclass class ReduxImageEncoderOutput(BaseOutput): - image_embeds: Optional[torch.Tensor] = None + image_embeds: torch.Tensor | None = None class ReduxImageEncoder(ModelMixin, ConfigMixin): diff --git a/src/diffusers/pipelines/flux/pipeline_flux.py b/src/diffusers/pipelines/flux/pipeline_flux.py index 5041e352f73d..be2bbe2acc6a 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux.py +++ b/src/diffusers/pipelines/flux/pipeline_flux.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -87,10 +87,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -105,15 +105,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -217,11 +217,11 @@ def __init__( def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -266,9 +266,9 @@ def _get_t5_prompt_embeds( def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): device = device or self._execution_device @@ -310,21 +310,21 @@ def _get_clip_prompt_embeds( def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Optional[Union[str, List[str]]] = None, - device: Optional[torch.device] = None, + prompt: str | list[str], + prompt_2: str | list[str] | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 512, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in all text-encoders device: (`torch.device`): @@ -653,49 +653,49 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt: Union[str, List[str]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + negative_prompt: str | list[str] = None, + negative_prompt_2: str | list[str] | None = None, true_cfg_scale: float = 1.0, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 28, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 3.5, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - negative_ip_adapter_image: Optional[PipelineImageInput] = None, - negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + negative_ip_adapter_image: PipelineImageInput | None = None, + negative_ip_adapter_image_embeds: list[torch.Tensor] | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is not greater than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders. true_cfg_scale (`float`, *optional*, defaults to 1.0): @@ -708,7 +708,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -720,7 +720,7 @@ def __call__( the [paper](https://huggingface.co/papers/2210.03142) to learn more. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -734,13 +734,13 @@ def __call__( Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, pooled text embeddings will be generated from `prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not provided, embeddings are computed from the `ip_adapter_image` input argument. negative_ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + negative_ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not provided, embeddings are computed from the `ip_adapter_image` input argument. @@ -766,7 +766,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -876,10 +876,15 @@ def __call__( self.scheduler.config.get("base_shift", 0.5), self.scheduler.config.get("max_shift", 1.15), ) + + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, - device, + timestep_device, sigmas=sigmas, mu=mu, ) diff --git a/src/diffusers/pipelines/flux/pipeline_flux_control.py b/src/diffusers/pipelines/flux/pipeline_flux_control.py index 848d7bd39254..84e21b59ce58 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_control.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_control.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -100,10 +100,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -118,15 +118,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -229,11 +229,11 @@ def __init__( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -279,9 +279,9 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): device = device or self._execution_device @@ -324,21 +324,21 @@ def _get_clip_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Optional[Union[str, List[str]]] = None, - device: Optional[torch.device] = None, + prompt: str | list[str], + prompt_2: str | list[str] | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 512, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in all text-encoders device: (`torch.device`): @@ -636,38 +636,38 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 28, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 3.5, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead - control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: - `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,: + `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`): The ControlNet input condition to provide guidance to the `unet` for generation. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or @@ -681,7 +681,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -693,7 +693,7 @@ def __call__( the [paper](https://huggingface.co/papers/2210.03142) to learn more. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -720,7 +720,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -829,10 +829,14 @@ def __call__( self.scheduler.config.get("base_shift", 0.5), self.scheduler.config.get("max_shift", 1.15), ) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, - device, + timestep_device, sigmas=sigmas, mu=mu, ) diff --git a/src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py b/src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py index 262345c75afc..b455c611e0ae 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -103,7 +103,7 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -118,10 +118,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -136,15 +136,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -239,11 +239,11 @@ def __init__( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -289,9 +289,9 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): device = device or self._execution_device @@ -334,21 +334,21 @@ def _get_clip_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Optional[Union[str, List[str]]] = None, - device: Optional[torch.device] = None, + prompt: str | list[str], + prompt_2: str | list[str] | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 512, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in all text-encoders device: (`torch.device`): @@ -635,46 +635,46 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, image: PipelineImageInput = None, control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, strength: float = 0.6, num_inference_steps: int = 28, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 7.0, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image latents as `image`, but if passing latents directly it is not encoded again. - control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: - `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,: + `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`): The ControlNet input condition to provide guidance to the `unet` for generation. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or @@ -694,7 +694,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -706,7 +706,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -733,7 +733,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -810,10 +810,14 @@ def __call__( self.scheduler.config.get("base_shift", 0.5), self.scheduler.config.get("max_shift", 1.15), ) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, - device, + timestep_device, sigmas=sigmas, mu=mu, ) diff --git a/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py b/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py index 6915a83a7ca7..2d1e05493a11 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -130,7 +130,7 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -145,10 +145,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -163,15 +163,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -279,11 +279,11 @@ def __init__( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -329,9 +329,9 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): device = device or self._execution_device @@ -374,21 +374,21 @@ def _get_clip_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Optional[Union[str, List[str]]] = None, - device: Optional[torch.device] = None, + prompt: str | list[str], + prompt_2: str | list[str] | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 512, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in all text-encoders device: (`torch.device`): @@ -806,62 +806,62 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, image: PipelineImageInput = None, control_image: PipelineImageInput = None, mask_image: PipelineImageInput = None, masked_image_latents: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, strength: float = 0.6, num_inference_steps: int = 28, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 7.0, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image latents as `image`, but if passing latents directly it is not encoded again. - control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: - `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,: + `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`): The ControlNet input condition to provide guidance to the `unet` for generation. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or width are passed, `image` is resized accordingly. If multiple ControlNets are specified in `init`, images must be passed as a list such that each element of the list can be correctly batched for input to a single ControlNet. - mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one color channel (L) instead of 3, so the expected shape for pytorch tensor would be `(B, 1, H, W)`, `(B, H, W)`, `(1, H, W)`, `(H, W)`. And for numpy array would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W, 1)`, or `(H, W)`. - mask_image_latent (`torch.Tensor`, `List[torch.Tensor]`): + mask_image_latent (`torch.Tensor`, `list[torch.Tensor]`): `Tensor` representing an image batch to mask `image` generated by VAE. If not provided, the mask latents tensor will be generated by `mask_image`. height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): @@ -877,7 +877,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -889,7 +889,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -916,7 +916,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1013,10 +1013,14 @@ def __call__( self.scheduler.config.get("base_shift", 0.5), self.scheduler.config.get("max_shift", 1.15), ) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, - device, + timestep_device, sigmas=sigmas, mu=mu, ) diff --git a/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py b/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py index 507ec687347c..d8dcdfcd4640 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_controlnet.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import torch @@ -102,7 +102,7 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -117,10 +117,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -135,15 +135,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -214,9 +214,10 @@ def __init__( text_encoder_2: T5EncoderModel, tokenizer_2: T5TokenizerFast, transformer: FluxTransformer2DModel, - controlnet: Union[ - FluxControlNetModel, List[FluxControlNetModel], Tuple[FluxControlNetModel], FluxMultiControlNetModel - ], + controlnet: FluxControlNetModel + | list[FluxControlNetModel] + | tuple[FluxControlNetModel] + | FluxMultiControlNetModel, image_encoder: CLIPVisionModelWithProjection = None, feature_extractor: CLIPImageProcessor = None, ): @@ -247,11 +248,11 @@ def __init__( def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -296,9 +297,9 @@ def _get_t5_prompt_embeds( def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): device = device or self._execution_device @@ -340,21 +341,21 @@ def _get_clip_prompt_embeds( def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Optional[Union[str, List[str]]] = None, - device: Optional[torch.device] = None, + prompt: str | list[str], + prompt_2: str | list[str] | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 512, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in all text-encoders device: (`torch.device`): @@ -678,47 +679,47 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt: Union[str, List[str]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + negative_prompt: str | list[str] = None, + negative_prompt_2: str | list[str] | None = None, true_cfg_scale: float = 1.0, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 28, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 7.0, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, control_image: PipelineImageInput = None, - control_mode: Optional[Union[int, List[int]]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - negative_ip_adapter_image: Optional[PipelineImageInput] = None, - negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + control_mode: int | list[int] | None = None, + controlnet_conditioning_scale: float | list[float] = 1.0, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + negative_ip_adapter_image: PipelineImageInput | None = None, + negative_ip_adapter_image_embeds: list[torch.Tensor] | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): @@ -728,7 +729,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -738,27 +739,27 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the ControlNet starts applying. - control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the ControlNet stops applying. - control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: - `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,: + `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`): The ControlNet input condition to provide guidance to the `unet` for generation. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or width are passed, `image` is resized accordingly. If multiple ControlNets are specified in `init`, images must be passed as a list such that each element of the list can be correctly batched for input to a single ControlNet. - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set the corresponding scale as a list. - control_mode (`int` or `List[int]`,, *optional*, defaults to None): + control_mode (`int` or `list[int]`,, *optional*, defaults to None): The control mode when applying ControlNet-Union. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -772,13 +773,13 @@ def __call__( Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, pooled text embeddings will be generated from `prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not provided, embeddings are computed from the `ip_adapter_image` input argument. negative_ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + negative_ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not provided, embeddings are computed from the `ip_adapter_image` input argument. @@ -796,7 +797,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1002,10 +1003,14 @@ def __call__( self.scheduler.config.get("base_shift", 0.5), self.scheduler.config.get("max_shift", 1.15), ) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, - device, + timestep_device, sigmas=sigmas, mu=mu, ) diff --git a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py index 582c7bbad84e..fdaff9b0af8a 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py @@ -1,5 +1,5 @@ import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import torch @@ -97,7 +97,7 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -112,10 +112,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -130,15 +130,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -209,9 +209,10 @@ def __init__( text_encoder_2: T5EncoderModel, tokenizer_2: T5TokenizerFast, transformer: FluxTransformer2DModel, - controlnet: Union[ - FluxControlNetModel, List[FluxControlNetModel], Tuple[FluxControlNetModel], FluxMultiControlNetModel - ], + controlnet: FluxControlNetModel + | list[FluxControlNetModel] + | tuple[FluxControlNetModel] + | FluxMultiControlNetModel, ): super().__init__() if isinstance(controlnet, (list, tuple)): @@ -239,11 +240,11 @@ def __init__( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -289,9 +290,9 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): device = device or self._execution_device @@ -334,21 +335,21 @@ def _get_clip_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Optional[Union[str, List[str]]] = None, - device: Optional[torch.device] = None, + prompt: str | list[str], + prompt_2: str | list[str] | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 512, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in all text-encoders device: (`torch.device`): @@ -635,43 +636,43 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, image: PipelineImageInput = None, control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, strength: float = 0.6, num_inference_steps: int = 28, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 7.0, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, - control_mode: Optional[Union[int, List[int]]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, + control_mode: int | list[int] | None = None, + controlnet_conditioning_scale: float | list[float] = 1.0, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. - image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.FloatTensor`): + image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.FloatTensor`): The image(s) to modify with the pipeline. - control_image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.FloatTensor`): + control_image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.FloatTensor`): The ControlNet input condition. Image to control the generation. height (`int`, *optional*, defaults to self.default_sample_size * self.vae_scale_factor): The height in pixels of the generated image. @@ -682,21 +683,21 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 28): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. guidance_scale (`float`, *optional*, defaults to 7.0): Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://huggingface.co/papers/2207.12598). - control_mode (`int` or `List[int]`, *optional*): + control_mode (`int` or `list[int]`, *optional*): The mode for the ControlNet. If multiple ControlNets are used, this should be a list. - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original transformer. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or more [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -714,7 +715,7 @@ def __call__( Additional keyword arguments to be passed to the joint attention mechanism. callback_on_step_end (`Callable`, *optional*): A function that calls at the end of each denoising step during the inference. - callback_on_step_end_tensor_inputs (`List[str]`, *optional*): + callback_on_step_end_tensor_inputs (`list[str]`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. max_sequence_length (`int`, *optional*, defaults to 512): The maximum length of the sequence to be generated. @@ -873,10 +874,14 @@ def __call__( self.scheduler.config.get("base_shift", 0.5), self.scheduler.config.get("max_shift", 1.15), ) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, - device, + timestep_device, sigmas=sigmas, mu=mu, ) diff --git a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py index f7f34ef231af..eed671152bc9 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py @@ -1,5 +1,5 @@ import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL @@ -99,7 +99,7 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -114,10 +114,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -132,15 +132,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -211,9 +211,10 @@ def __init__( text_encoder_2: T5EncoderModel, tokenizer_2: T5TokenizerFast, transformer: FluxTransformer2DModel, - controlnet: Union[ - FluxControlNetModel, List[FluxControlNetModel], Tuple[FluxControlNetModel], FluxMultiControlNetModel - ], + controlnet: FluxControlNetModel + | list[FluxControlNetModel] + | tuple[FluxControlNetModel] + | FluxMultiControlNetModel, ): super().__init__() if isinstance(controlnet, (list, tuple)): @@ -250,11 +251,11 @@ def __init__( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -300,9 +301,9 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): device = device or self._execution_device @@ -345,21 +346,21 @@ def _get_clip_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Optional[Union[str, List[str]]] = None, - device: Optional[torch.device] = None, + prompt: str | list[str], + prompt_2: str | list[str] | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 512, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in all text-encoders device: (`torch.device`): @@ -739,51 +740,51 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, image: PipelineImageInput = None, mask_image: PipelineImageInput = None, masked_image_latents: PipelineImageInput = None, control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, strength: float = 0.6, - padding_mask_crop: Optional[int] = None, - sigmas: Optional[List[float]] = None, + padding_mask_crop: int | None = None, + sigmas: list[float] | None = None, num_inference_steps: int = 28, guidance_scale: float = 7.0, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, - control_mode: Optional[Union[int, List[int]]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, + control_mode: int | list[int] | None = None, + controlnet_conditioning_scale: float | list[float] = 1.0, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. - image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.FloatTensor`): + image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.FloatTensor`): The image(s) to inpaint. - mask_image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.FloatTensor`): + mask_image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.FloatTensor`): The mask image(s) to use for inpainting. White pixels in the mask will be repainted, while black pixels will be preserved. masked_image_latents (`torch.FloatTensor`, *optional*): Pre-generated masked image latents. - control_image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.FloatTensor`): + control_image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.FloatTensor`): The ControlNet input condition. Image to control the generation. height (`int`, *optional*, defaults to self.default_sample_size * self.vae_scale_factor): The height in pixels of the generated image. @@ -796,25 +797,25 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 28): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. guidance_scale (`float`, *optional*, defaults to 7.0): Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://huggingface.co/papers/2207.12598). - control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the ControlNet starts applying. - control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the ControlNet stops applying. - control_mode (`int` or `List[int]`, *optional*): + control_mode (`int` or `list[int]`, *optional*): The mode for the ControlNet. If multiple ControlNets are used, this should be a list. - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original transformer. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or more [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -832,7 +833,7 @@ def __call__( Additional keyword arguments to be passed to the joint attention mechanism. callback_on_step_end (`Callable`, *optional*): A function that calls at the end of each denoising step during the inference. - callback_on_step_end_tensor_inputs (`List[str]`, *optional*): + callback_on_step_end_tensor_inputs (`list[str]`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. max_sequence_length (`int`, *optional*, defaults to 512): The maximum length of the sequence to be generated. @@ -1020,10 +1021,14 @@ def __call__( self.scheduler.config.get("base_shift", 0.5), self.scheduler.config.get("max_shift", 1.15), ) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, - device, + timestep_device, sigmas=sigmas, mu=mu, ) diff --git a/src/diffusers/pipelines/flux/pipeline_flux_fill.py b/src/diffusers/pipelines/flux/pipeline_flux_fill.py index 5cb9c82204b2..cf929f53fc6d 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_fill.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_fill.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -94,10 +94,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -112,15 +112,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -153,7 +153,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -244,11 +244,11 @@ def __init__( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -294,9 +294,9 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): device = device or self._execution_device @@ -419,21 +419,21 @@ def prepare_mask_latents( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Optional[Union[str, List[str]]] = None, - device: Optional[torch.device] = None, + prompt: str | list[str], + prompt_2: str | list[str] | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 512, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in all text-encoders device: (`torch.device`): @@ -753,52 +753,52 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - image: Optional[torch.FloatTensor] = None, - mask_image: Optional[torch.FloatTensor] = None, - masked_image_latents: Optional[torch.FloatTensor] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + image: torch.FloatTensor | None = None, + mask_image: torch.FloatTensor | None = None, + masked_image_latents: torch.FloatTensor | None = None, + height: int | None = None, + width: int | None = None, strength: float = 1.0, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 30.0, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)`. - mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one color channel (L) instead of 3, so the expected shape for pytorch tensor would be `(B, 1, H, W)`, `(B, H, W)`, `(1, H, W)`, `(H, W)`. And for numpy array would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W, 1)`, or `(H, W)`. - mask_image_latent (`torch.Tensor`, `List[torch.Tensor]`): + mask_image_latent (`torch.Tensor`, `list[torch.Tensor]`): `Tensor` representing an image batch to mask `image` generated by VAE. If not provided, the mask latents tensor will be generated by `mask_image`. height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): @@ -814,7 +814,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -826,7 +826,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -853,7 +853,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -932,10 +932,14 @@ def __call__( self.scheduler.config.get("base_shift", 0.5), self.scheduler.config.get("max_shift", 1.15), ) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, - device, + timestep_device, sigmas=sigmas, mu=mu, ) diff --git a/src/diffusers/pipelines/flux/pipeline_flux_img2img.py b/src/diffusers/pipelines/flux/pipeline_flux_img2img.py index ab9140dae921..cadff7736ff4 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_img2img.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_img2img.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -95,7 +95,7 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -110,10 +110,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -128,15 +128,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -238,11 +238,11 @@ def __init__( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -288,9 +288,9 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): device = device or self._execution_device @@ -333,21 +333,21 @@ def _get_clip_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Optional[Union[str, List[str]]] = None, - device: Optional[torch.device] = None, + prompt: str | list[str], + prompt_2: str | list[str] | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 512, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in all text-encoders device: (`torch.device`): @@ -735,47 +735,47 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt: Union[str, List[str]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + negative_prompt: str | list[str] = None, + negative_prompt_2: str | list[str] | None = None, true_cfg_scale: float = 1.0, image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, strength: float = 0.6, num_inference_steps: int = 28, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 7.0, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - negative_ip_adapter_image: Optional[PipelineImageInput] = None, - negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + negative_ip_adapter_image: PipelineImageInput | None = None, + negative_ip_adapter_image_embeds: list[torch.Tensor] | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a @@ -794,7 +794,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -806,7 +806,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -820,13 +820,13 @@ def __call__( Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, pooled text embeddings will be generated from `prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not provided, embeddings are computed from the `ip_adapter_image` input argument. negative_ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + negative_ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not provided, embeddings are computed from the `ip_adapter_image` input argument. @@ -844,7 +844,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -940,10 +940,14 @@ def __call__( self.scheduler.config.get("base_shift", 0.5), self.scheduler.config.get("max_shift", 1.15), ) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, - device, + timestep_device, sigmas=sigmas, mu=mu, ) diff --git a/src/diffusers/pipelines/flux/pipeline_flux_inpaint.py b/src/diffusers/pipelines/flux/pipeline_flux_inpaint.py index 3bfe82cf4382..b8ce25a4f5a9 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_inpaint.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_inpaint.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -91,7 +91,7 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -106,10 +106,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -124,15 +124,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -241,11 +241,11 @@ def __init__( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -291,9 +291,9 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): device = device or self._execution_device @@ -336,21 +336,21 @@ def _get_clip_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Optional[Union[str, List[str]]] = None, - device: Optional[torch.device] = None, + prompt: str | list[str], + prompt_2: str | list[str] | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 512, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in all text-encoders device: (`torch.device`): @@ -776,63 +776,63 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt: Union[str, List[str]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + negative_prompt: str | list[str] = None, + negative_prompt_2: str | list[str] | None = None, true_cfg_scale: float = 1.0, image: PipelineImageInput = None, mask_image: PipelineImageInput = None, masked_image_latents: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, - padding_mask_crop: Optional[int] = None, + height: int | None = None, + width: int | None = None, + padding_mask_crop: int | None = None, strength: float = 0.6, num_inference_steps: int = 28, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 7.0, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - negative_ip_adapter_image: Optional[PipelineImageInput] = None, - negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + negative_ip_adapter_image: PipelineImageInput | None = None, + negative_ip_adapter_image_embeds: list[torch.Tensor] | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image latents as `image`, but if passing latents directly it is not encoded again. - mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one color channel (L) instead of 3, so the expected shape for pytorch tensor would be `(B, 1, H, W)`, `(B, H, W)`, `(1, H, W)`, `(H, W)`. And for numpy array would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W, 1)`, or `(H, W)`. - mask_image_latent (`torch.Tensor`, `List[torch.Tensor]`): + mask_image_latent (`torch.Tensor`, `list[torch.Tensor]`): `Tensor` representing an image batch to mask `image` generated by VAE. If not provided, the mask latents tensor will be generated by `mask_image`. height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): @@ -855,7 +855,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -867,7 +867,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -881,13 +881,13 @@ def __call__( Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, pooled text embeddings will be generated from `prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not provided, embeddings are computed from the `ip_adapter_image` input argument. negative_ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + negative_ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not provided, embeddings are computed from the `ip_adapter_image` input argument. @@ -905,7 +905,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1015,10 +1015,14 @@ def __call__( self.scheduler.config.get("base_shift", 0.5), self.scheduler.config.get("max_shift", 1.15), ) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, - device, + timestep_device, sigmas=sigmas, mu=mu, ) diff --git a/src/diffusers/pipelines/flux/pipeline_flux_kontext.py b/src/diffusers/pipelines/flux/pipeline_flux_kontext.py index 94ae460afcd0..f4bbe42ef850 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_kontext.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_kontext.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -117,10 +117,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -135,15 +135,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -176,7 +176,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -263,11 +263,11 @@ def __init__( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -313,9 +313,9 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): device = device or self._execution_device @@ -358,21 +358,21 @@ def _get_clip_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Optional[Union[str, List[str]]] = None, - device: Optional[torch.device] = None, + prompt: str | list[str], + prompt_2: str | list[str] | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 512, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in all text-encoders device: (`torch.device`): @@ -668,15 +668,15 @@ def disable_vae_tiling(self): def prepare_latents( self, - image: Optional[torch.Tensor], + image: torch.Tensor | None, batch_size: int, num_channels_latents: int, height: int, width: int, dtype: torch.dtype, device: torch.device, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ): if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( @@ -752,33 +752,33 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - image: Optional[PipelineImageInput] = None, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt: Union[str, List[str]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, + image: PipelineImageInput | None = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + negative_prompt: str | list[str] = None, + negative_prompt_2: str | list[str] | None = None, true_cfg_scale: float = 1.0, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 28, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 3.5, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - negative_ip_adapter_image: Optional[PipelineImageInput] = None, - negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + negative_ip_adapter_image: PipelineImageInput | None = None, + negative_ip_adapter_image_embeds: list[torch.Tensor] | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, max_area: int = 1024**2, _auto_resize: bool = True, @@ -787,23 +787,23 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image latents as `image`, but if passing latents directly it is not encoded again. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is not greater than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders. true_cfg_scale (`float`, *optional*, defaults to 1.0): @@ -815,7 +815,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -827,7 +827,7 @@ def __call__( the [paper](https://huggingface.co/papers/2210.03142) to learn more. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -842,13 +842,13 @@ def __call__( If not provided, pooled text embeddings will be generated from `prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not provided, embeddings are computed from the `ip_adapter_image` input argument. negative_ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + negative_ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not provided, embeddings are computed from the `ip_adapter_image` input argument. @@ -874,7 +874,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py b/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py index b6f957981e14..313682dc7e33 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py @@ -2,7 +2,7 @@ # author: @vuongminh1907 import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -141,10 +141,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -159,15 +159,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -200,7 +200,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -296,11 +296,11 @@ def __init__( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -346,9 +346,9 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): device = device or self._execution_device @@ -391,21 +391,21 @@ def _get_clip_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Optional[Union[str, List[str]]] = None, - device: Optional[torch.device] = None, + prompt: str | list[str], + prompt_2: str | list[str] | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 512, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in all text-encoders device: (`torch.device`): @@ -742,7 +742,7 @@ def disable_vae_tiling(self): def prepare_latents( self, - image: Optional[torch.Tensor], + image: torch.Tensor | None, timestep: int, batch_size: int, num_channels_latents: int, @@ -750,9 +750,9 @@ def prepare_latents( width: int, dtype: torch.dtype, device: torch.device, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - image_reference: Optional[torch.Tensor] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + image_reference: torch.Tensor | None = None, ): if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( @@ -941,37 +941,37 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - image: Optional[PipelineImageInput] = None, - image_reference: Optional[PipelineImageInput] = None, + image: PipelineImageInput | None = None, + image_reference: PipelineImageInput | None = None, mask_image: PipelineImageInput = None, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt: Union[str, List[str]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + negative_prompt: str | list[str] = None, + negative_prompt_2: str | list[str] | None = None, true_cfg_scale: float = 1.0, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, strength: float = 1.0, - padding_mask_crop: Optional[int] = None, + padding_mask_crop: int | None = None, num_inference_steps: int = 28, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 3.5, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - negative_ip_adapter_image: Optional[PipelineImageInput] = None, - negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + negative_ip_adapter_image: PipelineImageInput | None = None, + negative_ip_adapter_image_embeds: list[torch.Tensor] | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, max_area: int = 1024**2, _auto_resize: bool = True, @@ -980,37 +980,37 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be be inpainted (which parts of the image to be masked out with `mask_image` and repainted according to `prompt` and `image_reference`). For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image latents as `image`, but if passing latents directly it is not encoded again. - image_reference (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image_reference (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point for the masked area. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)` If it is a numpy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image latents as `image`, but if passing latents directly it is not encoded again. - mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one color channel (L) instead of 3, so the expected shape for pytorch tensor would be `(B, 1, H, W)`, `(B, H, W)`, `(1, H, W)`, `(H, W)`. And for numpy array would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W, 1)`, or `(H, W)`. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is not greater than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders. true_cfg_scale (`float`, *optional*, defaults to 1.0): @@ -1036,7 +1036,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -1048,7 +1048,7 @@ def __call__( the [paper](https://huggingface.co/papers/2210.03142) to learn more. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -1063,13 +1063,13 @@ def __call__( If not provided, pooled text embeddings will be generated from `prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not provided, embeddings are computed from the `ip_adapter_image` input argument. negative_ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + negative_ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not provided, embeddings are computed from the `ip_adapter_image` input argument. @@ -1095,7 +1095,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/flux/pipeline_flux_prior_redux.py b/src/diffusers/pipelines/flux/pipeline_flux_prior_redux.py index e79db337b2e3..330e2623b287 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_prior_redux.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_prior_redux.py @@ -13,8 +13,6 @@ # limitations under the License. -from typing import List, Optional, Union - import torch from PIL import Image from transformers import ( @@ -196,11 +194,11 @@ def encode_image(self, image, device, num_images_per_prompt): # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -246,9 +244,9 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): device = device or self._execution_device @@ -291,21 +289,21 @@ def _get_clip_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Optional[Union[str, List[str]]] = None, - device: Optional[torch.device] = None, + prompt: str | list[str], + prompt_2: str | list[str] | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 512, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in all text-encoders device: (`torch.device`): @@ -373,28 +371,28 @@ def encode_prompt( def __call__( self, image: PipelineImageInput, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - prompt_embeds_scale: Optional[Union[float, List[float]]] = 1.0, - pooled_prompt_embeds_scale: Optional[Union[float, List[float]]] = 1.0, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + prompt_embeds_scale: float | list[float] | None = 1.0, + pooled_prompt_embeds_scale: float | list[float] | None = 1.0, return_dict: bool = True, ): r""" Function invoked when calling the pipeline for generation. Args: - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. **experimental feature**: to use this feature, make sure to explicitly load text encoders to the pipeline. Prompts will be ignored if text encoders are not loaded. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. prompt_embeds (`torch.FloatTensor`, *optional*): Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. diff --git a/src/diffusers/pipelines/flux/pipeline_output.py b/src/diffusers/pipelines/flux/pipeline_output.py index 69e742d3e026..7f24bdbf5d60 100644 --- a/src/diffusers/pipelines/flux/pipeline_output.py +++ b/src/diffusers/pipelines/flux/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -14,14 +13,14 @@ class FluxPipelineOutput(BaseOutput): Output class for Flux image generation pipelines. Args: - images (`List[PIL.Image.Image]` or `torch.Tensor` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or numpy array or torch tensor of shape `(batch_size, + images (`list[PIL.Image.Image]` or `torch.Tensor` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or numpy array or torch tensor of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. Torch tensors can represent either the denoised images or the intermediate latents ready to be passed to the decoder. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image] | np.ndarray @dataclass @@ -30,8 +29,8 @@ class FluxPriorReduxPipelineOutput(BaseOutput): Output class for Flux Prior Redux pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ diff --git a/src/diffusers/pipelines/flux2/__init__.py b/src/diffusers/pipelines/flux2/__init__.py index d986c9a63011..f6e1d5206630 100644 --- a/src/diffusers/pipelines/flux2/__init__.py +++ b/src/diffusers/pipelines/flux2/__init__.py @@ -23,6 +23,7 @@ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) else: _import_structure["pipeline_flux2"] = ["Flux2Pipeline"] + _import_structure["pipeline_flux2_klein"] = ["Flux2KleinPipeline"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: if not (is_transformers_available() and is_torch_available()): @@ -31,6 +32,7 @@ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 else: from .pipeline_flux2 import Flux2Pipeline + from .pipeline_flux2_klein import Flux2KleinPipeline else: import sys diff --git a/src/diffusers/pipelines/flux2/image_processor.py b/src/diffusers/pipelines/flux2/image_processor.py index f1a8742491f7..e0a1b80ce533 100644 --- a/src/diffusers/pipelines/flux2/image_processor.py +++ b/src/diffusers/pipelines/flux2/image_processor.py @@ -13,7 +13,6 @@ # limitations under the License. import math -from typing import List import PIL.Image @@ -148,7 +147,7 @@ def _resize_and_crop( # Taken from # https://github.com/black-forest-labs/flux2/blob/5a5d316b1b42f6b59a8c9194b77c8256be848432/src/flux2/sampling.py#L310C1-L339C19 @staticmethod - def concatenate_images(images: List[PIL.Image.Image]) -> PIL.Image.Image: + def concatenate_images(images: list[PIL.Image.Image]) -> PIL.Image.Image: """ Concatenate a list of PIL images horizontally with center alignment and white background. """ diff --git a/src/diffusers/pipelines/flux2/pipeline_flux2.py b/src/diffusers/pipelines/flux2/pipeline_flux2.py index b54a43dd89a5..6cd0563fcc19 100644 --- a/src/diffusers/pipelines/flux2/pipeline_flux2.py +++ b/src/diffusers/pipelines/flux2/pipeline_flux2.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL @@ -63,9 +63,9 @@ # Adapted from # https://github.com/black-forest-labs/flux2/blob/5a5d316b1b42f6b59a8c9194b77c8256be848432/src/flux2/text_encoder.py#L68 def format_input( - prompts: List[str], + prompts: list[str], system_message: str = SYSTEM_MESSAGE, - images: Optional[Union[List[PIL.Image.Image], List[List[PIL.Image.Image]]]] = None, + images: list[PIL.Image.Image, list[list[PIL.Image.Image]]] | None = None, ): """ Format a batch of text prompts into the conversation format expected by apply_chat_template. Optionally, add images @@ -130,10 +130,10 @@ def format_input( # Adapted from # https://github.com/black-forest-labs/flux2/blob/5a5d316b1b42f6b59a8c9194b77c8256be848432/src/flux2/text_encoder.py#L49C5-L66C19 def _validate_and_process_images( - images: List[List[PIL.Image.Image]] | List[PIL.Image.Image], + images: list[list[PIL.Image.Image]] | list[PIL.Image.Image], image_processor: Flux2ImageProcessor, upsampling_max_image_size: int, -) -> List[List[PIL.Image.Image]]: +) -> list[list[PIL.Image.Image]]: # Simple validation: ensure it's a list of PIL images or list of lists of PIL images if not images: return [] @@ -177,10 +177,10 @@ def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float: # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -195,15 +195,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -236,7 +236,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -304,12 +304,12 @@ def __init__( def _get_mistral_3_small_prompt_embeds( text_encoder: Mistral3ForConditionalGeneration, tokenizer: AutoProcessor, - prompt: Union[str, List[str]], - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, + prompt: str | list[str], + dtype: torch.dtype | None = None, + device: torch.device | None = None, max_sequence_length: int = 512, system_message: str = SYSTEM_MESSAGE, - hidden_states_layers: List[int] = (10, 20, 30), + hidden_states_layers: list[int] = (10, 20, 30), ): dtype = text_encoder.dtype if dtype is None else dtype device = text_encoder.device if device is None else device @@ -355,7 +355,7 @@ def _get_mistral_3_small_prompt_embeds( @staticmethod def _prepare_text_ids( x: torch.Tensor, # (B, L, D) or (L, D) - t_coord: Optional[torch.Tensor] = None, + t_coord: torch.Tensor | None = None, ): B, L, _ = x.shape out_ids = [] @@ -405,7 +405,7 @@ def _prepare_latent_ids( @staticmethod def _prepare_image_ids( - image_latents: List[torch.Tensor], # [(1, C, H, W), (1, C, H, W), ...] + image_latents: list[torch.Tensor], # [(1, C, H, W), (1, C, H, W), ...] scale: int = 10, ): r""" @@ -415,7 +415,7 @@ def _prepare_image_ids( dimensions. Args: - image_latents (List[torch.Tensor]): + image_latents (list[torch.Tensor]): A list of image latent feature tensors, typically of shape (C, H, W). scale (int, optional): A factor used to define the time separation (T-coordinate) between latents. T-coordinate for the i-th @@ -508,11 +508,11 @@ def _unpack_latents_with_ids(x: torch.Tensor, x_ids: torch.Tensor) -> list[torch def upsample_prompt( self, - prompt: Union[str, List[str]], - images: Union[List[PIL.Image.Image], List[List[PIL.Image.Image]]] = None, + prompt: str | list[str], + images: list[PIL.Image.Image, list[list[PIL.Image.Image]]] = None, temperature: float = 0.15, device: torch.device = None, - ) -> List[str]: + ) -> list[str]: prompt = [prompt] if isinstance(prompt, str) else prompt device = self.text_encoder.device if device is None else device @@ -570,12 +570,12 @@ def upsample_prompt( def encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 512, - text_encoder_out_layers: Tuple[int] = (10, 20, 30), + text_encoder_out_layers: tuple[int] = (10, 20, 30), ): device = device or self._execution_device @@ -625,7 +625,7 @@ def prepare_latents( dtype, device, generator: torch.Generator, - latents: Optional[torch.Tensor] = None, + latents: torch.Tensor | None = None, ): # VAE applies 8x compression on images but we must also account for packing which requires # latent height and width to be divisible by 2. @@ -651,7 +651,7 @@ def prepare_latents( def prepare_image_latents( self, - images: List[torch.Tensor], + images: list[torch.Tensor], batch_size, generator: torch.Generator, device, @@ -725,8 +725,8 @@ def guidance_scale(self): return self._guidance_scale @property - def joint_attention_kwargs(self): - return self._joint_attention_kwargs + def attention_kwargs(self): + return self._attention_kwargs @property def num_timesteps(self): @@ -744,37 +744,37 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - image: Optional[Union[List[PIL.Image.Image], PIL.Image.Image]] = None, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + image: list[PIL.Image.Image, PIL.Image.Image] | None = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, - guidance_scale: Optional[float] = 4.0, + sigmas: list[float] | None = None, + guidance_scale: float | None = 4.0, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, - text_encoder_out_layers: Tuple[int] = (10, 20, 30), + text_encoder_out_layers: tuple[int] = (10, 20, 30), caption_upsample_temperature: float = None, ): r""" Function invoked when calling the pipeline for generation. Args: - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image latents as `image`, but if passing latents directly it is not encoded again. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. guidance_scale (`float`, *optional*, defaults to 1.0): @@ -790,13 +790,13 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -825,7 +825,7 @@ def __call__( will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`. - text_encoder_out_layers (`Tuple[int]`): + text_encoder_out_layers (`tuple[int]`): Layer indices to use in the `text_encoder` to derive the final prompt embeddings. caption_upsample_temperature (`float`): When specified, we will try to perform caption upsampling for potentially improved outputs. We @@ -975,7 +975,7 @@ def __call__( encoder_hidden_states=prompt_embeds, txt_ids=text_ids, # B, text_seq_len, 4 img_ids=latent_image_ids, # B, image_seq_len, 4 - joint_attention_kwargs=self._attention_kwargs, + joint_attention_kwargs=self.attention_kwargs, return_dict=False, )[0] diff --git a/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py b/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py new file mode 100644 index 000000000000..936d2c3804ab --- /dev/null +++ b/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py @@ -0,0 +1,918 @@ +# Copyright 2025 Black Forest Labs and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +from typing import Any, Callable + +import numpy as np +import PIL +import torch +from transformers import Qwen2TokenizerFast, Qwen3ForCausalLM + +from ...loaders import Flux2LoraLoaderMixin +from ...models import AutoencoderKLFlux2, Flux2Transformer2DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import is_torch_xla_available, logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ..pipeline_utils import DiffusionPipeline +from .image_processor import Flux2ImageProcessor +from .pipeline_output import Flux2PipelineOutput + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import Flux2KleinPipeline + + >>> pipe = Flux2KleinPipeline.from_pretrained( + ... "black-forest-labs/FLUX.2-klein-base-9B", torch_dtype=torch.bfloat16 + ... ) + >>> pipe.to("cuda") + >>> prompt = "A cat holding a sign that says hello world" + >>> # Depending on the variant being used, the pipeline call will slightly vary. + >>> # Refer to the pipeline documentation for more details. + >>> image = pipe(prompt, num_inference_steps=50, guidance_scale=4.0).images[0] + >>> image.save("flux2_output.png") + ``` +""" + + +# Copied from diffusers.pipelines.flux2.pipeline_flux2.compute_empirical_mu +def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float: + a1, b1 = 8.73809524e-05, 1.89833333 + a2, b2 = 0.00016927, 0.45666666 + + if image_seq_len > 4300: + mu = a2 * image_seq_len + b2 + return float(mu) + + m_200 = a2 * image_seq_len + b2 + m_10 = a1 * image_seq_len + b1 + + a = (m_200 - m_10) / 190.0 + b = m_200 - 200.0 * a + mu = a * num_steps + b + + return float(mu) + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`list[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`list[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents +def retrieve_latents( + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" +): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +class Flux2KleinPipeline(DiffusionPipeline, Flux2LoraLoaderMixin): + r""" + The Flux2 Klein pipeline for text-to-image generation. + + Reference: + [https://bfl.ai/blog/flux2-klein-towards-interactive-visual-intelligence](https://bfl.ai/blog/flux2-klein-towards-interactive-visual-intelligence) + + Args: + transformer ([`Flux2Transformer2DModel`]): + Conditional Transformer (MMDiT) architecture to denoise the encoded image latents. + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + vae ([`AutoencoderKLFlux2`]): + Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. + text_encoder ([`Qwen3ForCausalLM`]): + [Qwen3ForCausalLM](https://huggingface.co/docs/transformers/en/model_doc/qwen3#transformers.Qwen3ForCausalLM) + tokenizer (`Qwen2TokenizerFast`): + Tokenizer of class + [Qwen2TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/qwen2#transformers.Qwen2TokenizerFast). + """ + + model_cpu_offload_seq = "text_encoder->transformer->vae" + _callback_tensor_inputs = ["latents", "prompt_embeds"] + + def __init__( + self, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKLFlux2, + text_encoder: Qwen3ForCausalLM, + tokenizer: Qwen2TokenizerFast, + transformer: Flux2Transformer2DModel, + is_distilled: bool = False, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + scheduler=scheduler, + transformer=transformer, + ) + + self.register_to_config(is_distilled=is_distilled) + + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 + # Flux latents are turned into 2x2 patches and packed. This means the latent width and height has to be divisible + # by the patch size. So the vae scale factor is multiplied by the patch size to account for this + self.image_processor = Flux2ImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) + self.tokenizer_max_length = 512 + self.default_sample_size = 128 + + @staticmethod + def _get_qwen3_prompt_embeds( + text_encoder: Qwen3ForCausalLM, + tokenizer: Qwen2TokenizerFast, + prompt: str | list[str], + dtype: torch.dtype | None = None, + device: torch.device | None = None, + max_sequence_length: int = 512, + hidden_states_layers: list[int] = (9, 18, 27), + ): + dtype = text_encoder.dtype if dtype is None else dtype + device = text_encoder.device if device is None else device + + prompt = [prompt] if isinstance(prompt, str) else prompt + + all_input_ids = [] + all_attention_masks = [] + + for single_prompt in prompt: + messages = [{"role": "user", "content": single_prompt}] + text = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + inputs = tokenizer( + text, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=max_sequence_length, + ) + + all_input_ids.append(inputs["input_ids"]) + all_attention_masks.append(inputs["attention_mask"]) + + input_ids = torch.cat(all_input_ids, dim=0).to(device) + attention_mask = torch.cat(all_attention_masks, dim=0).to(device) + + # Forward pass through the model + output = text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + use_cache=False, + ) + + # Only use outputs from intermediate layers and stack them + out = torch.stack([output.hidden_states[k] for k in hidden_states_layers], dim=1) + out = out.to(dtype=dtype, device=device) + + batch_size, num_channels, seq_len, hidden_dim = out.shape + prompt_embeds = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_channels * hidden_dim) + + return prompt_embeds + + @staticmethod + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._prepare_text_ids + def _prepare_text_ids( + x: torch.Tensor, # (B, L, D) or (L, D) + t_coord: torch.Tensor | None = None, + ): + B, L, _ = x.shape + out_ids = [] + + for i in range(B): + t = torch.arange(1) if t_coord is None else t_coord[i] + h = torch.arange(1) + w = torch.arange(1) + l = torch.arange(L) + + coords = torch.cartesian_prod(t, h, w, l) + out_ids.append(coords) + + return torch.stack(out_ids) + + @staticmethod + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._prepare_latent_ids + def _prepare_latent_ids( + latents: torch.Tensor, # (B, C, H, W) + ): + r""" + Generates 4D position coordinates (T, H, W, L) for latent tensors. + + Args: + latents (torch.Tensor): + Latent tensor of shape (B, C, H, W) + + Returns: + torch.Tensor: + Position IDs tensor of shape (B, H*W, 4) All batches share the same coordinate structure: T=0, + H=[0..H-1], W=[0..W-1], L=0 + """ + + batch_size, _, height, width = latents.shape + + t = torch.arange(1) # [0] - time dimension + h = torch.arange(height) + w = torch.arange(width) + l = torch.arange(1) # [0] - layer dimension + + # Create position IDs: (H*W, 4) + latent_ids = torch.cartesian_prod(t, h, w, l) + + # Expand to batch: (B, H*W, 4) + latent_ids = latent_ids.unsqueeze(0).expand(batch_size, -1, -1) + + return latent_ids + + @staticmethod + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._prepare_image_ids + def _prepare_image_ids( + image_latents: list[torch.Tensor], # [(1, C, H, W), (1, C, H, W), ...] + scale: int = 10, + ): + r""" + Generates 4D time-space coordinates (T, H, W, L) for a sequence of image latents. + + This function creates a unique coordinate for every pixel/patch across all input latent with different + dimensions. + + Args: + image_latents (list[torch.Tensor]): + A list of image latent feature tensors, typically of shape (C, H, W). + scale (int, optional): + A factor used to define the time separation (T-coordinate) between latents. T-coordinate for the i-th + latent is: 'scale + scale * i'. Defaults to 10. + + Returns: + torch.Tensor: + The combined coordinate tensor. Shape: (1, N_total, 4) Where N_total is the sum of (H * W) for all + input latents. + + Coordinate Components (Dimension 4): + - T (Time): The unique index indicating which latent image the coordinate belongs to. + - H (Height): The row index within that latent image. + - W (Width): The column index within that latent image. + - L (Seq. Length): A sequence length dimension, which is always fixed at 0 (size 1) + """ + + if not isinstance(image_latents, list): + raise ValueError(f"Expected `image_latents` to be a list, got {type(image_latents)}.") + + # create time offset for each reference image + t_coords = [scale + scale * t for t in torch.arange(0, len(image_latents))] + t_coords = [t.view(-1) for t in t_coords] + + image_latent_ids = [] + for x, t in zip(image_latents, t_coords): + x = x.squeeze(0) + _, height, width = x.shape + + x_ids = torch.cartesian_prod(t, torch.arange(height), torch.arange(width), torch.arange(1)) + image_latent_ids.append(x_ids) + + image_latent_ids = torch.cat(image_latent_ids, dim=0) + image_latent_ids = image_latent_ids.unsqueeze(0) + + return image_latent_ids + + @staticmethod + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._patchify_latents + def _patchify_latents(latents): + batch_size, num_channels_latents, height, width = latents.shape + latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 1, 3, 5, 2, 4) + latents = latents.reshape(batch_size, num_channels_latents * 4, height // 2, width // 2) + return latents + + @staticmethod + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._unpatchify_latents + def _unpatchify_latents(latents): + batch_size, num_channels_latents, height, width = latents.shape + latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), 2, 2, height, width) + latents = latents.permute(0, 1, 4, 2, 5, 3) + latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), height * 2, width * 2) + return latents + + @staticmethod + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._pack_latents + def _pack_latents(latents): + """ + pack latents: (batch_size, num_channels, height, width) -> (batch_size, height * width, num_channels) + """ + + batch_size, num_channels, height, width = latents.shape + latents = latents.reshape(batch_size, num_channels, height * width).permute(0, 2, 1) + + return latents + + @staticmethod + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._unpack_latents_with_ids + def _unpack_latents_with_ids(x: torch.Tensor, x_ids: torch.Tensor) -> list[torch.Tensor]: + """ + using position ids to scatter tokens into place + """ + x_list = [] + for data, pos in zip(x, x_ids): + _, ch = data.shape # noqa: F841 + h_ids = pos[:, 1].to(torch.int64) + w_ids = pos[:, 2].to(torch.int64) + + h = torch.max(h_ids) + 1 + w = torch.max(w_ids) + 1 + + flat_ids = h_ids * w + w_ids + + out = torch.zeros((h * w, ch), device=data.device, dtype=data.dtype) + out.scatter_(0, flat_ids.unsqueeze(1).expand(-1, ch), data) + + # reshape from (H * W, C) to (H, W, C) and permute to (C, H, W) + + out = out.view(h, w, ch).permute(2, 0, 1) + x_list.append(out) + + return torch.stack(x_list, dim=0) + + def encode_prompt( + self, + prompt: str | list[str], + device: torch.device | None = None, + num_images_per_prompt: int = 1, + prompt_embeds: torch.Tensor | None = None, + max_sequence_length: int = 512, + text_encoder_out_layers: tuple[int] = (9, 18, 27), + ): + device = device or self._execution_device + + if prompt is None: + prompt = "" + + prompt = [prompt] if isinstance(prompt, str) else prompt + + if prompt_embeds is None: + prompt_embeds = self._get_qwen3_prompt_embeds( + text_encoder=self.text_encoder, + tokenizer=self.tokenizer, + prompt=prompt, + device=device, + max_sequence_length=max_sequence_length, + hidden_states_layers=text_encoder_out_layers, + ) + + batch_size, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + text_ids = self._prepare_text_ids(prompt_embeds) + text_ids = text_ids.to(device) + return prompt_embeds, text_ids + + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._encode_vae_image + def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): + if image.ndim != 4: + raise ValueError(f"Expected image dims 4, got {image.ndim}.") + + image_latents = retrieve_latents(self.vae.encode(image), generator=generator, sample_mode="argmax") + image_latents = self._patchify_latents(image_latents) + + latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(image_latents.device, image_latents.dtype) + latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps) + image_latents = (image_latents - latents_bn_mean) / latents_bn_std + + return image_latents + + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline.prepare_latents + def prepare_latents( + self, + batch_size, + num_latents_channels, + height, + width, + dtype, + device, + generator: torch.Generator, + latents: torch.Tensor | None = None, + ): + # VAE applies 8x compression on images but we must also account for packing which requires + # latent height and width to be divisible by 2. + height = 2 * (int(height) // (self.vae_scale_factor * 2)) + width = 2 * (int(width) // (self.vae_scale_factor * 2)) + + shape = (batch_size, num_latents_channels * 4, height // 2, width // 2) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + if latents is None: + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + else: + latents = latents.to(device=device, dtype=dtype) + + latent_ids = self._prepare_latent_ids(latents) + latent_ids = latent_ids.to(device) + + latents = self._pack_latents(latents) # [B, C, H, W] -> [B, H*W, C] + return latents, latent_ids + + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline.prepare_image_latents + def prepare_image_latents( + self, + images: list[torch.Tensor], + batch_size, + generator: torch.Generator, + device, + dtype, + ): + image_latents = [] + for image in images: + image = image.to(device=device, dtype=dtype) + imagge_latent = self._encode_vae_image(image=image, generator=generator) + image_latents.append(imagge_latent) # (1, 128, 32, 32) + + image_latent_ids = self._prepare_image_ids(image_latents) + + # Pack each latent and concatenate + packed_latents = [] + for latent in image_latents: + # latent: (1, 128, 32, 32) + packed = self._pack_latents(latent) # (1, 1024, 128) + packed = packed.squeeze(0) # (1024, 128) - remove batch dim + packed_latents.append(packed) + + # Concatenate all reference tokens along sequence dimension + image_latents = torch.cat(packed_latents, dim=0) # (N*1024, 128) + image_latents = image_latents.unsqueeze(0) # (1, N*1024, 128) + + image_latents = image_latents.repeat(batch_size, 1, 1) + image_latent_ids = image_latent_ids.repeat(batch_size, 1, 1) + image_latent_ids = image_latent_ids.to(device) + + return image_latents, image_latent_ids + + def check_inputs( + self, + prompt, + height, + width, + prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + guidance_scale=None, + ): + if ( + height is not None + and height % (self.vae_scale_factor * 2) != 0 + or width is not None + and width % (self.vae_scale_factor * 2) != 0 + ): + logger.warning( + f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and {width}. Dimensions will be resized accordingly" + ) + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if guidance_scale > 1.0 and self.config.is_distilled: + logger.warning(f"Guidance scale {guidance_scale} is ignored for step-wise distilled models.") + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1 and not self.config.is_distilled + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + image: list[PIL.Image.Image] | PIL.Image.Image | None = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, + num_inference_steps: int = 50, + sigmas: list[float] | None = None, + guidance_scale: float = 4.0, + num_images_per_prompt: int = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: str | list[str] | None = None, + output_type: str = "pil", + return_dict: bool = True, + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int, dict], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + max_sequence_length: int = 512, + text_encoder_out_layers: tuple[int] = (9, 18, 27), + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both + numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list + or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a + list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image + latents as `image`, but if passing latents directly it is not encoded again. + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + guidance_scale (`float`, *optional*, defaults to 4.0): + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. + of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting + `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to + the text `prompt`, usually at the expense of lower image quality. For step-wise distilled models, + `guidance_scale` is ignored. + height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The height in pixels of the generated image. This is set to 1024 by default for the best results. + width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The width in pixels of the generated image. This is set to 1024 by default for the best results. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. Note that "" is used as the negative prompt in this pipeline. + If not provided, will be generated from "". + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.qwenimage.QwenImagePipelineOutput`] instead of a plain tuple. + attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`. + text_encoder_out_layers (`tuple[int]`): + Layer indices to use in the `text_encoder` to derive the final prompt embeddings. + + Examples: + + Returns: + [`~pipelines.flux2.Flux2PipelineOutput`] or `tuple`: [`~pipelines.flux2.Flux2PipelineOutput`] if + `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the + generated images. + """ + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt=prompt, + height=height, + width=width, + prompt_embeds=prompt_embeds, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + guidance_scale=guidance_scale, + ) + + self._guidance_scale = guidance_scale + self._attention_kwargs = attention_kwargs + self._current_timestep = None + self._interrupt = False + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + + # 3. prepare text embeddings + prompt_embeds, text_ids = self.encode_prompt( + prompt=prompt, + prompt_embeds=prompt_embeds, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + text_encoder_out_layers=text_encoder_out_layers, + ) + + if self.do_classifier_free_guidance: + negative_prompt = "" + if prompt is not None and isinstance(prompt, list): + negative_prompt = [negative_prompt] * len(prompt) + negative_prompt_embeds, negative_text_ids = self.encode_prompt( + prompt=negative_prompt, + prompt_embeds=negative_prompt_embeds, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + text_encoder_out_layers=text_encoder_out_layers, + ) + + # 4. process images + if image is not None and not isinstance(image, list): + image = [image] + + condition_images = None + if image is not None: + for img in image: + self.image_processor.check_image_input(img) + + condition_images = [] + for img in image: + image_width, image_height = img.size + if image_width * image_height > 1024 * 1024: + img = self.image_processor._resize_to_target_area(img, 1024 * 1024) + image_width, image_height = img.size + + multiple_of = self.vae_scale_factor * 2 + image_width = (image_width // multiple_of) * multiple_of + image_height = (image_height // multiple_of) * multiple_of + img = self.image_processor.preprocess(img, height=image_height, width=image_width, resize_mode="crop") + condition_images.append(img) + height = height or image_height + width = width or image_width + + height = height or self.default_sample_size * self.vae_scale_factor + width = width or self.default_sample_size * self.vae_scale_factor + + # 5. prepare latent variables + num_channels_latents = self.transformer.config.in_channels // 4 + latents, latent_ids = self.prepare_latents( + batch_size=batch_size * num_images_per_prompt, + num_latents_channels=num_channels_latents, + height=height, + width=width, + dtype=prompt_embeds.dtype, + device=device, + generator=generator, + latents=latents, + ) + + image_latents = None + image_latent_ids = None + if condition_images is not None: + image_latents, image_latent_ids = self.prepare_image_latents( + images=condition_images, + batch_size=batch_size * num_images_per_prompt, + generator=generator, + device=device, + dtype=self.vae.dtype, + ) + + # 6. Prepare timesteps + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas + if hasattr(self.scheduler.config, "use_flow_sigmas") and self.scheduler.config.use_flow_sigmas: + sigmas = None + image_seq_len = latents.shape[1] + mu = compute_empirical_mu(image_seq_len=image_seq_len, num_steps=num_inference_steps) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + sigmas=sigmas, + mu=mu, + ) + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + # 7. Denoising loop + # We set the index here to remove DtoH sync, helpful especially during compilation. + # Check out more details here: https://github.com/huggingface/diffusers/pull/11696 + self.scheduler.set_begin_index(0) + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + self._current_timestep = t + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latents.shape[0]).to(latents.dtype) + + latent_model_input = latents.to(self.transformer.dtype) + latent_image_ids = latent_ids + + if image_latents is not None: + latent_model_input = torch.cat([latents, image_latents], dim=1).to(self.transformer.dtype) + latent_image_ids = torch.cat([latent_ids, image_latent_ids], dim=1) + + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( + hidden_states=latent_model_input, # (B, image_seq_len, C) + timestep=timestep / 1000, + guidance=None, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, # B, text_seq_len, 4 + img_ids=latent_image_ids, # B, image_seq_len, 4 + joint_attention_kwargs=self.attention_kwargs, + return_dict=False, + )[0] + + noise_pred = noise_pred[:, : latents.size(1) :] + + if self.do_classifier_free_guidance: + with self.transformer.cache_context("uncond"): + neg_noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + guidance=None, + encoder_hidden_states=negative_prompt_embeds, + txt_ids=negative_text_ids, + img_ids=latent_image_ids, + joint_attention_kwargs=self._attention_kwargs, + return_dict=False, + )[0] + neg_noise_pred = neg_noise_pred[:, : latents.size(1) :] + noise_pred = neg_noise_pred + guidance_scale * (noise_pred - neg_noise_pred) + + # compute the previous noisy sample x_t -> x_t-1 + latents_dtype = latents.dtype + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + if latents.dtype != latents_dtype: + if torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 + latents = latents.to(latents_dtype) + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + self._current_timestep = None + + latents = self._unpack_latents_with_ids(latents, latent_ids) + + latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) + latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to( + latents.device, latents.dtype + ) + latents = latents * latents_bn_std + latents_bn_mean + latents = self._unpatchify_latents(latents) + if output_type == "latent": + image = latents + else: + image = self.vae.decode(latents, return_dict=False)[0] + image = self.image_processor.postprocess(image, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return Flux2PipelineOutput(images=image) diff --git a/src/diffusers/pipelines/flux2/pipeline_output.py b/src/diffusers/pipelines/flux2/pipeline_output.py index 58e8ad49c210..34ae9b574453 100644 --- a/src/diffusers/pipelines/flux2/pipeline_output.py +++ b/src/diffusers/pipelines/flux2/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -13,11 +12,11 @@ class Flux2PipelineOutput(BaseOutput): Output class for Flux2 image generation pipelines. Args: - images (`List[PIL.Image.Image]` or `torch.Tensor` or `np.ndarray`) + images (`list[PIL.Image.Image]` or `torch.Tensor` or `np.ndarray`) List of denoised PIL images of length `batch_size` or numpy array or torch tensor of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. Torch tensors can represent either the denoised images or the intermediate latents ready to be passed to the decoder. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image, np.ndarray] diff --git a/src/diffusers/pipelines/free_init_utils.py b/src/diffusers/pipelines/free_init_utils.py index 4495c5ea2683..04a385edd819 100644 --- a/src/diffusers/pipelines/free_init_utils.py +++ b/src/diffusers/pipelines/free_init_utils.py @@ -13,7 +13,6 @@ # limitations under the License. import math -from typing import Tuple, Union import torch import torch.fft as fft @@ -73,8 +72,8 @@ def free_init_enabled(self): def _get_free_init_freq_filter( self, - shape: Tuple[int, ...], - device: Union[str, torch.dtype], + shape: tuple[int, ...], + device: str | torch.dtype, filter_type: str, order: float, spatial_stop_frequency: float, diff --git a/src/diffusers/pipelines/free_noise_utils.py b/src/diffusers/pipelines/free_noise_utils.py index 2910afaf237b..5990e680ba07 100644 --- a/src/diffusers/pipelines/free_noise_utils.py +++ b/src/diffusers/pipelines/free_noise_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Callable, Dict, List, Optional, Tuple, Union +from typing import Callable import torch import torch.nn as nn @@ -48,7 +48,7 @@ class SplitInferenceModule(nn.Module): The size of each chunk after splitting the input tensor. split_dim (`int`, defaults to `0`): The dimension along which the input tensors are split. - input_kwargs_to_split (`List[str]`, defaults to `["hidden_states"]`): + input_kwargs_to_split (`list[str]`, defaults to `["hidden_states"]`): A list of keyword arguments (strings) that represent the input tensors to be split. Workflow: @@ -80,7 +80,7 @@ def __init__( module: nn.Module, split_size: int = 1, split_dim: int = 0, - input_kwargs_to_split: List[str] = ["hidden_states"], + input_kwargs_to_split: list[str] = ["hidden_states"], ) -> None: super().__init__() @@ -89,7 +89,7 @@ def __init__( self.split_dim = split_dim self.input_kwargs_to_split = set(input_kwargs_to_split) - def forward(self, *args, **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor]]: + def forward(self, *args, **kwargs) -> torch.Tensor | tuple[torch.Tensor]: r"""Forward method for the `SplitInferenceModule`. This method processes the input by splitting specified keyword arguments along a given dimension, running the @@ -99,13 +99,13 @@ def forward(self, *args, **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor]]: Args: *args (`Any`): Positional arguments that are passed directly to the `module` without modification. - **kwargs (`Dict[str, torch.Tensor]`): + **kwargs (`dict[str, torch.Tensor]`): Keyword arguments passed to the underlying `module`. Only keyword arguments whose names match the entries in `input_kwargs_to_split` and are of type `torch.Tensor` will be split. The remaining keyword arguments are passed unchanged. Returns: - `Union[torch.Tensor, Tuple[torch.Tensor]]`: + `torch.Tensor | tuple[torch.Tensor]`: The outputs obtained from `SplitInferenceModule` are the same as if the underlying module was inferred without it. - If the underlying module returns a single tensor, the result will be a single concatenated tensor @@ -145,7 +145,7 @@ def forward(self, *args, **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor]]: class AnimateDiffFreeNoiseMixin: r"""Mixin class for [FreeNoise](https://huggingface.co/papers/2310.15169).""" - def _enable_free_noise_in_block(self, block: Union[CrossAttnDownBlockMotion, DownBlockMotion, UpBlockMotion]): + def _enable_free_noise_in_block(self, block: CrossAttnDownBlockMotion | DownBlockMotion | UpBlockMotion): r"""Helper function to enable FreeNoise in transformer blocks.""" for motion_module in block.motion_modules: @@ -186,7 +186,7 @@ def _enable_free_noise_in_block(self, block: Union[CrossAttnDownBlockMotion, Dow basic_transfomer_block._chunk_size, basic_transfomer_block._chunk_dim ) - def _disable_free_noise_in_block(self, block: Union[CrossAttnDownBlockMotion, DownBlockMotion, UpBlockMotion]): + def _disable_free_noise_in_block(self, block: CrossAttnDownBlockMotion | DownBlockMotion | UpBlockMotion): r"""Helper function to disable FreeNoise in transformer blocks.""" for motion_module in block.motion_modules: @@ -255,16 +255,16 @@ def _check_inputs_free_noise( def _encode_prompt_free_noise( self, - prompt: Union[str, Dict[int, str]], + prompt: str | dict[int, str], num_frames: int, device: torch.device, num_videos_per_prompt: int, do_classifier_free_guidance: bool, - negative_prompt: Optional[Union[str, Dict[int, str]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | dict[int, str] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ) -> torch.Tensor: if negative_prompt is None: negative_prompt = "" @@ -362,8 +362,8 @@ def _prepare_latents_free_noise( width: int, dtype: torch.dtype, device: torch.device, - generator: Optional[torch.Generator] = None, - latents: Optional[torch.Tensor] = None, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, ): if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( @@ -443,13 +443,14 @@ def _lerp( def enable_free_noise( self, - context_length: Optional[int] = 16, + context_length: int | None = 16, context_stride: int = 4, weighting_scheme: str = "pyramid", noise_type: str = "shuffle_context", - prompt_interpolation_callback: Optional[ - Callable[[DiffusionPipeline, int, int, torch.Tensor, torch.Tensor], torch.Tensor] - ] = None, + prompt_interpolation_callback: Callable[ + [DiffusionPipeline, int, int, torch.Tensor, torch.Tensor], torch.Tensor + ] + | None = None, ) -> None: r""" Enable long video generation using FreeNoise. @@ -529,7 +530,7 @@ def disable_free_noise(self) -> None: self._disable_free_noise_in_block(block) def _enable_split_inference_motion_modules_( - self, motion_modules: List[AnimateDiffTransformer3D], spatial_split_size: int + self, motion_modules: list[AnimateDiffTransformer3D], spatial_split_size: int ) -> None: for motion_module in motion_modules: motion_module.proj_in = SplitInferenceModule(motion_module.proj_in, spatial_split_size, 0, ["input"]) @@ -545,19 +546,19 @@ def _enable_split_inference_motion_modules_( motion_module.proj_out = SplitInferenceModule(motion_module.proj_out, spatial_split_size, 0, ["input"]) def _enable_split_inference_attentions_( - self, attentions: List[Transformer2DModel], temporal_split_size: int + self, attentions: list[Transformer2DModel], temporal_split_size: int ) -> None: for i in range(len(attentions)): attentions[i] = SplitInferenceModule( attentions[i], temporal_split_size, 0, ["hidden_states", "encoder_hidden_states"] ) - def _enable_split_inference_resnets_(self, resnets: List[ResnetBlock2D], temporal_split_size: int) -> None: + def _enable_split_inference_resnets_(self, resnets: list[ResnetBlock2D], temporal_split_size: int) -> None: for i in range(len(resnets)): resnets[i] = SplitInferenceModule(resnets[i], temporal_split_size, 0, ["input_tensor", "temb"]) def _enable_split_inference_samplers_( - self, samplers: Union[List[Downsample2D], List[Upsample2D]], temporal_split_size: int + self, samplers: list[Downsample2D] | list[Upsample2D], temporal_split_size: int ) -> None: for i in range(len(samplers)): samplers[i] = SplitInferenceModule(samplers[i], temporal_split_size, 0, ["hidden_states"]) diff --git a/src/diffusers/pipelines/glm_image/__init__.py b/src/diffusers/pipelines/glm_image/__init__.py new file mode 100644 index 000000000000..140b9cc760cc --- /dev/null +++ b/src/diffusers/pipelines/glm_image/__init__.py @@ -0,0 +1,59 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, + is_transformers_version, +) + + +_dummy_objects = {} +_additional_imports = {} +_import_structure = {"pipeline_output": ["GlmImagePipelineOutput"]} + +# Import transformers components so they can be resolved during pipeline loading + +if is_transformers_available() and is_transformers_version(">=", "4.57.4"): + try: + from transformers import GlmImageForConditionalGeneration, GlmImageProcessor + + _additional_imports["GlmImageForConditionalGeneration"] = GlmImageForConditionalGeneration + _additional_imports["GlmImageProcessor"] = GlmImageProcessor + except ImportError: + pass + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["pipeline_glm_image"] = ["GlmImagePipeline"] +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 + else: + from .pipeline_glm_image import GlmImagePipeline +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) + for name, value in _additional_imports.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/glm_image/pipeline_glm_image.py b/src/diffusers/pipelines/glm_image/pipeline_glm_image.py new file mode 100644 index 000000000000..859b371b2514 --- /dev/null +++ b/src/diffusers/pipelines/glm_image/pipeline_glm_image.py @@ -0,0 +1,1049 @@ +# Copyright 2025 The CogVideoX team, Tsinghua University & ZhipuAI and The HuggingFace Team. +# All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import re +from typing import Any, Callable + +import numpy as np +import PIL +import torch +from transformers import ByT5Tokenizer, PreTrainedModel, ProcessorMixin, T5EncoderModel + +from ...callbacks import MultiPipelineCallbacks, PipelineCallback +from ...image_processor import VaeImageProcessor +from ...models import AutoencoderKL, GlmImageTransformer2DModel +from ...models.transformers.transformer_glm_image import GlmImageKVCache +from ...pipelines.pipeline_utils import DiffusionPipeline +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import is_torch_xla_available, is_transformers_version, logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from .pipeline_output import GlmImagePipelineOutput + + +# Because it's not released in stable as of 13/01/2026. So this is just a proxy. +GlmImageProcessor = ProcessorMixin +GlmImageForConditionalGeneration = PreTrainedModel +if is_transformers_version(">=", "5.0.0.dev0"): + from transformers import GlmImageForConditionalGeneration, GlmImageProcessor + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```python + >>> import torch + >>> from diffusers import GlmImagePipeline + + >>> pipe = GlmImagePipeline.from_pretrained("zai-org/GLM-Image", torch_dtype=torch.bfloat16) + >>> pipe.to("cuda") + + >>> prompt = "A photo of an astronaut riding a horse on mars" + >>> image = pipe(prompt).images[0] + >>> image.save("output.png") + ``` +""" + + +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + base_shift: float = 0.25, + max_shift: float = 0.75, +) -> float: + m = (image_seq_len / base_seq_len) ** 0.5 + mu = m * max_shift + base_shift + return mu + + +# Copied from diffusers.pipelines.cogview4.pipeline_cogview4.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`list[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`list[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + accepts_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + + if timesteps is not None and sigmas is not None: + if not accepts_timesteps and not accepts_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep or sigma schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif timesteps is not None and sigmas is None: + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif timesteps is None and sigmas is not None: + if not accepts_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents +def retrieve_latents( + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" +): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +class GlmImagePipeline(DiffusionPipeline): + r""" + Pipeline for text-to-image generation using GLM-Image. + + This pipeline integrates both the AR (autoregressive) model for token generation and the DiT (diffusion + transformer) model for image decoding. + + Args: + tokenizer (`PreTrainedTokenizer`): + Tokenizer for the text encoder. + processor (`AutoProcessor`): + Processor for the AR model to handle chat templates and tokenization. + text_encoder ([`T5EncoderModel`]): + Frozen text-encoder for glyph embeddings. + vision_language_encoder ([`GlmImageForConditionalGeneration`]): + The AR model that generates image tokens from text prompts. + vae ([`AutoencoderKL`]): + Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. + transformer ([`GlmImageTransformer2DModel`]): + A text conditioned transformer to denoise the encoded image latents (DiT). + scheduler ([`SchedulerMixin`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + """ + + _optional_components = [] + model_cpu_offload_seq = "vision_language_encoder->text_encoder->transformer->vae" + _callback_tensor_inputs = ["latents", "prompt_embeds"] + + def __init__( + self, + tokenizer: ByT5Tokenizer, + processor: GlmImageProcessor, + text_encoder: T5EncoderModel, + vision_language_encoder: GlmImageForConditionalGeneration, + vae: AutoencoderKL, + transformer: GlmImageTransformer2DModel, + scheduler: FlowMatchEulerDiscreteScheduler, + ): + super().__init__() + + self.register_modules( + tokenizer=tokenizer, + processor=processor, + text_encoder=text_encoder, + vision_language_encoder=vision_language_encoder, + vae=vae, + transformer=transformer, + scheduler=scheduler, + ) + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + + self.default_sample_size = ( + self.transformer.config.sample_size + if hasattr(self, "transformer") + and self.transformer is not None + and hasattr(self.transformer.config, "sample_size") + else 128 + ) + + @staticmethod + def _compute_generation_params( + image_grid_thw, + is_text_to_image: bool, + ): + grid_sizes = [] + grid_hw = [] + + for i in range(image_grid_thw.shape[0]): + t, h, w = image_grid_thw[i].tolist() + grid_sizes.append(int(h * w)) + grid_hw.append((int(h), int(w))) + + if not is_text_to_image: + max_new_tokens = grid_sizes[-1] + 1 + large_image_start_offset = 0 + target_grid_h, target_grid_w = grid_hw[-1] + else: + total_tokens = sum(grid_sizes) + max_new_tokens = total_tokens + 1 + large_image_start_offset = sum(grid_sizes[1:]) + target_grid_h, target_grid_w = grid_hw[0] + return max_new_tokens, large_image_start_offset, target_grid_h, target_grid_w + + @staticmethod + def _extract_large_image_tokens( + outputs: torch.Tensor, input_length: int, large_image_start_offset: int, large_image_tokens: int + ) -> torch.Tensor: + generated_tokens = outputs[0][input_length:] + large_image_start = large_image_start_offset + large_image_end = large_image_start + large_image_tokens + return generated_tokens[large_image_start:large_image_end] + + @staticmethod + def _upsample_token_ids(token_ids: torch.Tensor, token_h: int, token_w: int) -> torch.Tensor: + token_ids = token_ids.view(1, 1, token_h, token_w) + token_ids = torch.nn.functional.interpolate(token_ids.float(), scale_factor=2, mode="nearest").to( + dtype=torch.long + ) + token_ids = token_ids.view(1, -1) + return token_ids + + @staticmethod + def _validate_and_normalize_images( + image: list[PIL.Image.Image] | list[list[PIL.Image.Image]], + batch_size: int, + ) -> list[list[PIL.Image.Image]]: + """ + Validate and normalize image inputs to List[List[PIL.Image]]. + + Rules: + - batch_size > 1: Only accepts List[List[PIL.Image]], each sublist must have equal length + - batch_size == 1: Accepts List[PIL.Image] for legacy compatibility (converted to [[img1, img2, ...]]) + - Other formats raise ValueError + + Args: + image: Input images in various formats + batch_size: Number of prompts in the batch + + Returns: + Normalized images as List[List[PIL.Image]], or None if no images provided + """ + if image is None or len(image) == 0: + return None + + first_element = image[0] + + if batch_size == 1: + # Legacy format: List[PIL.Image] -> [[img1, img2, ...]] + if not isinstance(first_element, (list, tuple)): + return [list(image)] + # Already in List[List[PIL.Image]] format + if len(image) != 1: + raise ValueError( + f"For batch_size=1 with List[List[PIL.Image]] format, expected 1 image list, got {len(image)}." + ) + return [list(image[0])] + + # batch_size > 1: must be List[List[PIL.Image]] + if not isinstance(first_element, (list, tuple)): + raise ValueError( + f"For batch_size > 1, images must be List[List[PIL.Image]] format. " + f"Got List[{type(first_element).__name__}] instead. " + f"Each prompt requires its own list of condition images." + ) + + if len(image) != batch_size: + raise ValueError(f"Number of image lists ({len(image)}) must match batch size ({batch_size}).") + + # Validate homogeneous: all sublists must have same length + num_input_images_per_prompt = len(image[0]) + for idx, imgs in enumerate(image): + if len(imgs) != num_input_images_per_prompt: + raise ValueError( + f"All prompts must have the same number of condition images. " + f"Prompt 0 has {num_input_images_per_prompt} images, but prompt {idx} has {len(imgs)} images." + ) + + return [list(imgs) for imgs in image] + + def generate_prior_tokens( + self, + prompt: str | list[str], + height: int, + width: int, + image: list[list[PIL.Image.Image]] | None = None, + device: torch.device | None = None, + generator: torch.Generator | None = None, + ): + """ + Generate prior tokens for the DiT model using the AR model. + + Args: + prompt: Single prompt or list of prompts + height: Target image height + width: Target image width + image: Normalized image input as List[List[PIL.Image]]. Should be pre-validated + using _validate_and_normalize_images() before calling this method. + device: Target device + generator: Random generator for reproducibility + + Returns: + Tuple of: + - prior_token_ids: Tensor of shape (batch_size, num_tokens) with upsampled prior tokens + - prior_token_image_ids_per_sample: List of tensors, one per sample. Each tensor contains + the upsampled prior token ids for all condition images in that sample. None for t2i. + - source_image_grid_thw_per_sample: List of tensors, one per sample. Each tensor has shape + (num_condition_images, 3) with upsampled grid info. None for t2i. + """ + device = device or self._execution_device + + # Normalize prompt to list format + prompt_list = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt_list) + + # Image is already normalized by _validate_and_normalize_images(): None or List[List[PIL.Image]] + is_text_to_image = image is None + # Build messages for each sample in the batch + all_messages = [] + for idx, p in enumerate(prompt_list): + content = [] + if not is_text_to_image: + for img in image[idx]: + content.append({"type": "image", "image": img}) + content.append({"type": "text", "text": p}) + all_messages.append([{"role": "user", "content": content}]) + # Process with the processor (supports batch with left padding) + inputs = self.processor.apply_chat_template( + all_messages, + tokenize=True, + padding=True if batch_size > 1 else False, + target_h=height, + target_w=width, + return_dict=True, + return_tensors="pt", + ).to(device) + + image_grid_thw = inputs.get("image_grid_thw") + images_per_sample = inputs.get("images_per_sample") + + # Determine number of condition images and grids per sample + num_condition_images = 0 if is_text_to_image else len(image[0]) + if images_per_sample is not None: + num_grids_per_sample = images_per_sample[0].item() + else: + # Fallback for batch_size=1: total grids is for single sample + num_grids_per_sample = image_grid_thw.shape[0] + + # Compute generation params (same for all samples in homogeneous batch) + first_sample_grids = image_grid_thw[:num_grids_per_sample] + max_new_tokens, large_image_offset, token_h, token_w = self._compute_generation_params( + image_grid_thw=first_sample_grids, is_text_to_image=is_text_to_image + ) + + # Generate source image tokens (prior_token_image_ids) for i2i mode + prior_token_image_ids = None + source_image_grid_thw = None + if not is_text_to_image: + # Extract source grids by selecting condition image indices (skip target grids) + # Grid order from processor: [s0_cond1, s0_cond2, ..., s0_target, s1_cond1, s1_cond2, ..., s1_target, ...] + # We need indices: [0, 1, ..., num_condition_images-1, num_grids_per_sample, num_grids_per_sample+1, ...] + source_indices = [] + for sample_idx in range(batch_size): + base = sample_idx * num_grids_per_sample + source_indices.extend(range(base, base + num_condition_images)) + source_grids = image_grid_thw[source_indices] + + if len(source_grids) > 0: + prior_token_image_embed = self.vision_language_encoder.get_image_features( + inputs["pixel_values"], source_grids + ).pooler_output + prior_token_image_embed = torch.cat(prior_token_image_embed, dim=0) + prior_token_image_ids_d32 = self.vision_language_encoder.get_image_tokens( + prior_token_image_embed, source_grids + ) + # Upsample each source image's prior tokens to match VAE/DiT resolution + split_sizes = source_grids.prod(dim=-1).tolist() + prior_ids_per_source = torch.split(prior_token_image_ids_d32, split_sizes) + upsampled_prior_ids = [] + for i, prior_ids in enumerate(prior_ids_per_source): + t, h, w = source_grids[i].tolist() + upsampled = self._upsample_token_ids(prior_ids, int(h), int(w)) + upsampled_prior_ids.append(upsampled.squeeze(0)) + prior_token_image_ids = torch.cat(upsampled_prior_ids, dim=0) + # Upsample grid dimensions for later splitting + upsampled_grids = source_grids.clone() + upsampled_grids[:, 1] = upsampled_grids[:, 1] * 2 + upsampled_grids[:, 2] = upsampled_grids[:, 2] * 2 + source_image_grid_thw = upsampled_grids + + # Generate with AR model + # Set torch random seed from generator for reproducibility + # (transformers generate() doesn't accept generator parameter) + if generator is not None: + seed = generator.initial_seed() + torch.manual_seed(seed) + if device is not None and device.type == "cuda": + torch.cuda.manual_seed(seed) + outputs = self.vision_language_encoder.generate( + **inputs, + max_new_tokens=max_new_tokens, + do_sample=True, + ) + + # Extract and upsample prior tokens for each sample + # For left-padded inputs, generated tokens start after the padded input sequence + all_prior_token_ids = [] + max_input_length = inputs["input_ids"].shape[-1] + for idx in range(batch_size): + # For left-padded sequences, generated tokens start at max_input_length + # (padding is on the left, so all sequences end at the same position) + prior_token_ids_d32 = self._extract_large_image_tokens( + outputs[idx : idx + 1], max_input_length, large_image_offset, token_h * token_w + ) + prior_token_ids = self._upsample_token_ids(prior_token_ids_d32, token_h, token_w) + all_prior_token_ids.append(prior_token_ids) + prior_token_ids = torch.cat(all_prior_token_ids, dim=0) + + # Split prior_token_image_ids and source_image_grid_thw into per-sample lists for easier consumption + prior_token_image_ids_per_sample = None + source_image_grid_thw_per_sample = None + if prior_token_image_ids is not None and source_image_grid_thw is not None: + # Split grids: each sample has num_condition_images grids + source_image_grid_thw_per_sample = list(torch.split(source_image_grid_thw, num_condition_images)) + # Split prior_token_image_ids: tokens per sample may vary due to different image sizes + tokens_per_image = source_image_grid_thw.prod(dim=-1).tolist() + tokens_per_sample = [] + for i in range(batch_size): + start_idx = i * num_condition_images + end_idx = start_idx + num_condition_images + tokens_per_sample.append(sum(tokens_per_image[start_idx:end_idx])) + prior_token_image_ids_per_sample = list(torch.split(prior_token_image_ids, tokens_per_sample)) + + return prior_token_ids, prior_token_image_ids_per_sample, source_image_grid_thw_per_sample + + def get_glyph_texts(self, prompt): + """Extract glyph texts from prompt(s). Returns a list of lists for batch processing.""" + if isinstance(prompt, str): + prompt = [prompt] + all_ocr_texts = [] + for p in prompt: + ocr_texts = ( + re.findall(r"'([^']*)'", p) + + re.findall(r"\u201c([^\u201c\u201d]*)\u201d", p) + + re.findall(r'"([^"]*)"', p) + + re.findall(r"「([^「」]*)」", p) + ) + all_ocr_texts.append(ocr_texts) + return all_ocr_texts + + def _get_glyph_embeds( + self, + prompt: str | list[str] = None, + max_sequence_length: int = 2048, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + """Get glyph embeddings for each prompt in the batch.""" + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + # get_glyph_texts now returns a list of lists (one per prompt) + all_glyph_texts = self.get_glyph_texts(prompt) + + all_glyph_embeds = [] + for glyph_texts in all_glyph_texts: + if len(glyph_texts) == 0: + glyph_texts = [""] + input_ids = self.tokenizer( + glyph_texts, + max_length=max_sequence_length, + truncation=True, + ).input_ids + input_ids = [ + [self.tokenizer.pad_token_id] * ((len(input_ids) + 1) % 2) + input_ids_ for input_ids_ in input_ids + ] + max_length = max(len(input_ids_) for input_ids_ in input_ids) + attention_mask = torch.tensor( + [[1] * len(input_ids_) + [0] * (max_length - len(input_ids_)) for input_ids_ in input_ids], + device=device, + ) + input_ids = torch.tensor( + [ + input_ids_ + [self.tokenizer.pad_token_id] * (max_length - len(input_ids_)) + for input_ids_ in input_ids + ], + device=device, + ) + outputs = self.text_encoder(input_ids, attention_mask=attention_mask) + glyph_embeds = outputs.last_hidden_state[attention_mask.bool()].unsqueeze(0) + all_glyph_embeds.append(glyph_embeds) + + # Pad to same sequence length and stack (use left padding to match transformers) + max_seq_len = max(emb.size(1) for emb in all_glyph_embeds) + padded_embeds = [] + for emb in all_glyph_embeds: + if emb.size(1) < max_seq_len: + pad = torch.zeros(emb.size(0), max_seq_len - emb.size(1), emb.size(2), device=device, dtype=emb.dtype) + emb = torch.cat([pad, emb], dim=1) # left padding + padded_embeds.append(emb) + + glyph_embeds = torch.cat(padded_embeds, dim=0) + return glyph_embeds.to(device=device, dtype=dtype) + + def encode_prompt( + self, + prompt: str | list[str], + do_classifier_free_guidance: bool = True, + num_images_per_prompt: int = 1, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + max_sequence_length: int = 2048, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `list[str]`, *optional*): + prompt to be encoded + do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): + Whether to use classifier free guidance or not. + num_images_per_prompt (`int`, *optional*, defaults to 1): + Number of images that should be generated per prompt. torch device to place the resulting embeddings on + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + device: (`torch.device`, *optional*): + torch device + dtype: (`torch.dtype`, *optional*): + torch dtype + max_sequence_length (`int`, defaults to `2048`): + Maximum sequence length in encoded prompt. Can be set to other values but may lead to poorer results. + """ + device = device or self._execution_device + + prompt = [prompt] if isinstance(prompt, str) else prompt + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds = self._get_glyph_embeds(prompt, max_sequence_length, device, dtype) + + # Repeat embeddings for num_images_per_prompt + if num_images_per_prompt > 1: + prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) + + # For GLM-Image, negative_prompt must be "" instead of None + if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + negative_prompt_embeds = self._get_glyph_embeds(negative_prompt, max_sequence_length, device, dtype) + + if num_images_per_prompt > 1: + negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) + + return prompt_embeds, negative_prompt_embeds + + def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None): + if latents is not None: + return latents.to(device) + + shape = ( + batch_size, + num_channels_latents, + int(height) // self.vae_scale_factor, + int(width) // self.vae_scale_factor, + ) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + return latents + + def check_inputs( + self, + prompt, + height, + width, + callback_on_step_end_tensor_inputs, + prompt_embeds=None, + negative_prompt_embeds=None, + prior_token_ids=None, + prior_token_image_ids=None, + source_image_grid_thw=None, + image=None, + ): + if ( + height is not None + and height % (self.vae_scale_factor * self.transformer.config.patch_size * 2) != 0 + or width is not None + and width % (self.transformer.config.patch_size * 2) != 0 + ): + # GLM-Image uses 32× downsampling, so the image dimensions must be multiples of 32. + raise ValueError( + f"`height` and `width` have to be divisible by {self.vae_scale_factor * 4} but are {height} and {width}." + ) + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + if prompt is None and prior_token_ids is None: + raise ValueError( + "Provide either `prompt` or `prior_token_ids`. Cannot leave both `prompt` and `prior_token_ids` undefined." + ) + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}." + ) + # Validate prior token inputs: for i2i mode, all three must be provided together + # For t2i mode, only prior_token_ids is needed (prior_token_image_ids and source_image_grid_thw should be None) + prior_image_inputs = [prior_token_image_ids, source_image_grid_thw] + num_prior_image_inputs = sum(x is not None for x in prior_image_inputs) + if num_prior_image_inputs > 0 and num_prior_image_inputs < len(prior_image_inputs): + raise ValueError( + "`prior_token_image_ids` and `source_image_grid_thw` must be provided together for i2i mode. " + f"Got prior_token_image_ids={prior_token_image_ids is not None}, " + f"source_image_grid_thw={source_image_grid_thw is not None}." + ) + if num_prior_image_inputs > 0 and prior_token_ids is None: + raise ValueError( + "`prior_token_ids` must be provided when `prior_token_image_ids` and `source_image_grid_thw` are provided." + ) + if num_prior_image_inputs > 0 and image is None: + raise ValueError( + "`image` must be provided when `prior_token_image_ids` and `source_image_grid_thw` are provided " + "for i2i mode, as the images are needed for VAE encoding to build the KV cache." + ) + + if prior_token_ids is not None and prompt_embeds is None and prompt is None: + raise ValueError("`prompt_embeds` or `prompt` must also be provided with `prior_token_ids`.") + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: str | list[str] | None = None, + image: torch.Tensor + | PIL.Image.Image + | np.ndarray + | list[torch.Tensor] + | list[PIL.Image.Image] + | list[np.ndarray] + | None = None, + height: int | None = None, + width: int | None = None, + num_inference_steps: int = 50, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + guidance_scale: float = 1.5, + num_images_per_prompt: int = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prior_token_ids: torch.Tensor | None = None, + prior_token_image_ids: list[torch.Tensor] | None = None, + source_image_grid_thw: list[torch.Tensor] | None = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + output_type: str = "pil", + return_dict: bool = True, + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int, dict], None] + | PipelineCallback + | MultiPipelineCallbacks + | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + max_sequence_length: int = 2048, + ) -> GlmImagePipelineOutput | tuple: + """ + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `list[str]`, *optional*): + The prompt or prompts to guide the image generation. Must contain shape info in the format 'H + W' where H and W are token dimensions (d32). Example: "A beautiful sunset36 24" + generates a 1152x768 image. + image: Optional condition images for image-to-image generation. + height (`int`, *optional*): + The height in pixels. If not provided, derived from prompt shape info. + width (`int`, *optional*): + The width in pixels. If not provided, derived from prompt shape info. + num_inference_steps (`int`, *optional*, defaults to `50`): + The number of denoising steps for DiT. + guidance_scale (`float`, *optional*, defaults to `1.5`): + Guidance scale for classifier-free guidance. + num_images_per_prompt (`int`, *optional*, defaults to `1`): + The number of images to generate per prompt. + generator (`torch.Generator`, *optional*): + Random generator for reproducibility. + output_type (`str`, *optional*, defaults to `"pil"`): + Output format: "pil", "np", or "latent". + + Examples: + + Returns: + [`GlmImagePipelineOutput`] or `tuple`: Generated images. + """ + + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + # 1. Check inputs + self.check_inputs( + prompt, + height, + width, + callback_on_step_end_tensor_inputs, + prompt_embeds, + negative_prompt_embeds, + prior_token_ids, + prior_token_image_ids, + source_image_grid_thw, + image, + ) + self._guidance_scale = guidance_scale + self._attention_kwargs = attention_kwargs + self._current_timestep = None + self._interrupt = False + + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + + # 2. Validate and normalize image format + normalized_image = self._validate_and_normalize_images(image, batch_size) + + # 3. Generate prior tokens (batch mode) + # Get a single generator for AR model (use first if list provided) + ar_generator = generator[0] if isinstance(generator, list) else generator + if prior_token_ids is None: + prior_token_ids, prior_token_image_ids_per_sample, source_image_grid_thw_per_sample = ( + self.generate_prior_tokens( + prompt=prompt, + image=normalized_image, + height=height, + width=width, + device=device, + generator=ar_generator, + ) + ) + else: + # User provided prior_token_ids directly (from generate_prior_tokens) + prior_token_image_ids_per_sample = prior_token_image_ids + source_image_grid_thw_per_sample = source_image_grid_thw + + # 4. Preprocess images for VAE encoding + preprocessed_images = None + if normalized_image is not None: + preprocessed_images = [] + for prompt_images in normalized_image: + prompt_preprocessed = [] + for img in prompt_images: + image_height, image_width = img.size[::-1] if isinstance(img, PIL.Image.Image) else img.shape[:2] + multiple_of = self.vae_scale_factor * self.transformer.config.patch_size + image_height = (image_height // multiple_of) * multiple_of + image_width = (image_width // multiple_of) * multiple_of + img = self.image_processor.preprocess(img, height=image_height, width=image_width) + prompt_preprocessed.append(img) + height = height or image_height + width = width or image_width + preprocessed_images.append(prompt_preprocessed) + + # 5. Encode input prompt + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt, + self.do_classifier_free_guidance, + num_images_per_prompt=num_images_per_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + max_sequence_length=max_sequence_length, + device=device, + dtype=self.dtype, + ) + + # 6. Prepare latents and (optional) image kv cache + latent_channels = self.transformer.config.in_channels + latents = self.prepare_latents( + batch_size=batch_size * num_images_per_prompt, + num_channels_latents=latent_channels, + height=height, + width=width, + dtype=prompt_embeds.dtype, + device=device, + generator=generator, + latents=latents, + ) + kv_caches = GlmImageKVCache(num_layers=self.transformer.config.num_layers) + + if normalized_image is not None: + kv_caches.set_mode("write") + latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, self.vae.config.latent_channels, 1, 1) + latents_std = torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.latent_channels, 1, 1) + + latents_mean = latents_mean.to(device=device, dtype=prompt_embeds.dtype) + latents_std = latents_std.to(device=device, dtype=prompt_embeds.dtype) + + # Process each sample's condition images + for prompt_idx in range(batch_size): + prompt_images = preprocessed_images[prompt_idx] + prompt_prior_ids = prior_token_image_ids_per_sample[prompt_idx] + prompt_grid_thw = source_image_grid_thw_per_sample[prompt_idx] + + # Split this sample's prior_token_image_ids by each image's token count + split_sizes = prompt_grid_thw.prod(dim=-1).tolist() + prior_ids_per_image = torch.split(prompt_prior_ids, split_sizes) + # Process each condition image for this sample + for condition_image, condition_image_prior_token_id in zip(prompt_images, prior_ids_per_image): + condition_image = condition_image.to(device=device, dtype=prompt_embeds.dtype) + condition_latent = retrieve_latents( + self.vae.encode(condition_image), generator=generator, sample_mode="argmax" + ) + condition_latent = (condition_latent - latents_mean) / latents_std + + _ = self.transformer( + hidden_states=condition_latent, + encoder_hidden_states=torch.zeros_like(prompt_embeds)[:1, :0, ...], + prior_token_id=condition_image_prior_token_id, + prior_token_drop=torch.full_like(condition_image_prior_token_id, False, dtype=torch.bool), + timestep=torch.zeros((1,), device=device), + target_size=torch.tensor([condition_image.shape[-2:]], device=device), + crop_coords=torch.zeros((1, 2), device=device), + attention_kwargs=attention_kwargs, + kv_caches=kv_caches, + ) + # Move to next sample's cache slot + kv_caches.next_sample() + + # 7. Prepare additional timestep conditions + target_size = (height, width) + target_size = torch.tensor([target_size], dtype=prompt_embeds.dtype, device=device) + crops_coords_top_left = torch.tensor([crops_coords_top_left], dtype=prompt_embeds.dtype, device=device) + + target_size = target_size.repeat(batch_size * num_images_per_prompt, 1) + crops_coords_top_left = crops_coords_top_left.repeat(batch_size * num_images_per_prompt, 1) + + # Prepare timesteps + image_seq_len = ((height // self.vae_scale_factor) * (width // self.vae_scale_factor)) // ( + self.transformer.config.patch_size**2 + ) + timesteps = ( + np.linspace(self.scheduler.config.num_train_timesteps, 1.0, num_inference_steps + 1)[:-1] + if timesteps is None + else np.array(timesteps) + ) + timesteps = timesteps.astype(np.int64).astype(np.float32) + sigmas = timesteps / self.scheduler.config.num_train_timesteps if sigmas is None else sigmas + mu = calculate_shift( + image_seq_len, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("base_shift", 0.25), + self.scheduler.config.get("max_shift", 0.75), + ) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, device, timesteps, sigmas, mu=mu + ) + self._num_timesteps = len(timesteps) + + # 8. Denoising loop + transformer_dtype = self.transformer.dtype + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + + # Repeat prior_token_ids for num_images_per_prompt + if num_images_per_prompt > 1: + prior_token_ids = prior_token_ids.repeat_interleave(num_images_per_prompt, dim=0) + prior_token_drop_cond = torch.full_like(prior_token_ids, False, dtype=torch.bool) + prior_token_drop_uncond = torch.full_like(prior_token_ids, True, dtype=torch.bool) + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + self._current_timestep = t + latent_model_input = latents.to(transformer_dtype) + + timestep = t.expand(latents.shape[0]) - 1 + + if prior_token_image_ids_per_sample is not None: + kv_caches.set_mode("read") + + noise_pred_cond = self.transformer( + hidden_states=latent_model_input, + encoder_hidden_states=prompt_embeds, + prior_token_id=prior_token_ids, + prior_token_drop=prior_token_drop_cond, + timestep=timestep, + target_size=target_size, + crop_coords=crops_coords_top_left, + attention_kwargs=attention_kwargs, + return_dict=False, + kv_caches=kv_caches, + )[0].float() + + # perform guidance + if self.do_classifier_free_guidance: + if prior_token_image_ids_per_sample is not None: + kv_caches.set_mode("skip") + noise_pred_uncond = self.transformer( + hidden_states=latent_model_input, + encoder_hidden_states=negative_prompt_embeds, + prior_token_id=prior_token_ids, + prior_token_drop=prior_token_drop_uncond, + timestep=timestep, + target_size=target_size, + crop_coords=crops_coords_top_left, + attention_kwargs=attention_kwargs, + return_dict=False, + kv_caches=kv_caches, + )[0].float() + + noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_cond - noise_pred_uncond) + else: + noise_pred = noise_pred_cond + + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, self.scheduler.sigmas[i], callback_kwargs) + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + self._current_timestep = None + kv_caches.clear() + + if not output_type == "latent": + latents = latents.to(self.vae.dtype) + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.latent_channels, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = ( + torch.tensor(self.vae.config.latents_std) + .view(1, self.vae.config.latent_channels, 1, 1) + .to(latents.device, latents.dtype) + ) + latents = latents * latents_std + latents_mean + image = self.vae.decode(latents, return_dict=False, generator=generator)[0] + image = self.image_processor.postprocess(image, output_type=output_type) + else: + image = latents + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return GlmImagePipelineOutput(images=image) diff --git a/src/diffusers/pipelines/glm_image/pipeline_output.py b/src/diffusers/pipelines/glm_image/pipeline_output.py new file mode 100644 index 000000000000..d4fd061335d4 --- /dev/null +++ b/src/diffusers/pipelines/glm_image/pipeline_output.py @@ -0,0 +1,20 @@ +from dataclasses import dataclass + +import numpy as np +import PIL.Image + +from ...utils import BaseOutput + + +@dataclass +class GlmImagePipelineOutput(BaseOutput): + """ + Output class for CogView3 pipelines. + + Args: + images (`List[PIL.Image.Image]` or `np.ndarray`) + List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. + """ + + images: list[PIL.Image.Image] | np.ndarray diff --git a/src/diffusers/pipelines/helios/__init__.py b/src/diffusers/pipelines/helios/__init__.py new file mode 100644 index 000000000000..ae08f5997279 --- /dev/null +++ b/src/diffusers/pipelines/helios/__init__.py @@ -0,0 +1,48 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["pipeline_helios"] = ["HeliosPipeline"] + _import_structure["pipeline_helios_pyramid"] = ["HeliosPyramidPipeline"] +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * + else: + from .pipeline_helios import HeliosPipeline + from .pipeline_helios_pyramid import HeliosPyramidPipeline +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/helios/pipeline_helios.py b/src/diffusers/pipelines/helios/pipeline_helios.py new file mode 100644 index 000000000000..87a8600badab --- /dev/null +++ b/src/diffusers/pipelines/helios/pipeline_helios.py @@ -0,0 +1,916 @@ +# Copyright 2025 The Helios Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import html +from typing import Any, Callable + +import numpy as np +import regex as re +import torch +from transformers import AutoTokenizer, UMT5EncoderModel + +from ...callbacks import MultiPipelineCallbacks, PipelineCallback +from ...image_processor import PipelineImageInput +from ...loaders import HeliosLoraLoaderMixin +from ...models import AutoencoderKLWan, HeliosTransformer3DModel +from ...schedulers import HeliosScheduler +from ...utils import is_ftfy_available, is_torch_xla_available, logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor +from ..pipeline_utils import DiffusionPipeline +from .pipeline_output import HeliosPipelineOutput + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +if is_ftfy_available(): + import ftfy + + +EXAMPLE_DOC_STRING = """ + Examples: + ```python + >>> import torch + >>> from diffusers.utils import export_to_video + >>> from diffusers import AutoencoderKLWan, HeliosPipeline + + >>> # Available models: BestWishYsh/Helios-Base, BestWishYsh/Helios-Mid, BestWishYsh/Helios-Distilled + >>> model_id = "BestWishYsh/Helios-Base" + >>> vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32) + >>> pipe = HeliosPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16) + >>> pipe.to("cuda") + + >>> prompt = "A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window." + >>> negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" + + >>> output = pipe( + ... prompt=prompt, + ... negative_prompt=negative_prompt, + ... height=384, + ... width=640, + ... num_frames=132, + ... guidance_scale=5.0, + ... ).frames[0] + >>> export_to_video(output, "output.mp4", fps=24) + ``` +""" + + +def basic_clean(text): + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + return text.strip() + + +def whitespace_clean(text): + text = re.sub(r"\s+", " ", text) + text = text.strip() + return text + + +def prompt_clean(text): + text = whitespace_clean(basic_clean(text)) + return text + + +# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +class HeliosPipeline(DiffusionPipeline, HeliosLoraLoaderMixin): + r""" + Pipeline for text-to-video / image-to-video / video-to-video generation using Helios. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods + implemented for all pipelines (downloading, saving, running on a particular device, etc.). + + Args: + tokenizer ([`T5Tokenizer`]): + Tokenizer from [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5Tokenizer), + specifically the [google/umt5-xxl](https://huggingface.co/google/umt5-xxl) variant. + text_encoder ([`T5EncoderModel`]): + [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically + the [google/umt5-xxl](https://huggingface.co/google/umt5-xxl) variant. + transformer ([`HeliosTransformer3DModel`]): + Conditional Transformer to denoise the input latents. + scheduler ([`HeliosScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + vae ([`AutoencoderKLWan`]): + Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. + """ + + model_cpu_offload_seq = "text_encoder->transformer->vae" + _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] + _optional_components = ["transformer"] + + def __init__( + self, + tokenizer: AutoTokenizer, + text_encoder: UMT5EncoderModel, + vae: AutoencoderKLWan, + scheduler: HeliosScheduler, + transformer: HeliosTransformer3DModel, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + scheduler=scheduler, + ) + self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if getattr(self, "vae", None) else 4 + self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if getattr(self, "vae", None) else 8 + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + + def _get_t5_prompt_embeds( + self, + prompt: str | list[str] = None, + num_videos_per_prompt: int = 1, + max_sequence_length: int = 226, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + prompt = [prompt_clean(u) for u in prompt] + batch_size = len(prompt) + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask + seq_lens = mask.gt(0).sum(dim=1).long() + + prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + prompt_embeds = torch.stack( + [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0 + ) + + # duplicate text embeddings for each generation per prompt, using mps friendly method + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) + + return prompt_embeds, text_inputs.attention_mask.bool() + + def encode_prompt( + self, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, + do_classifier_free_guidance: bool = True, + num_videos_per_prompt: int = 1, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + max_sequence_length: int = 226, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `list[str]`, *optional*): + prompt to be encoded + negative_prompt (`str` or `list[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): + Whether to use classifier free guidance or not. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos that should be generated per prompt. torch device to place the resulting embeddings on + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + device: (`torch.device`, *optional*): + torch device + dtype: (`torch.dtype`, *optional*): + torch dtype + """ + device = device or self._execution_device + + prompt = [prompt] if isinstance(prompt, str) else prompt + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds, _ = self._get_t5_prompt_embeds( + prompt=prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + + if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + + if prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + + negative_prompt_embeds, _ = self._get_t5_prompt_embeds( + prompt=negative_prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + + return prompt_embeds, negative_prompt_embeds + + def check_inputs( + self, + prompt, + negative_prompt, + height, + width, + prompt_embeds=None, + negative_prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + image=None, + video=None, + ): + if height % 16 != 0 or width % 16 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.") + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`: {negative_prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + elif negative_prompt is not None and ( + not isinstance(negative_prompt, str) and not isinstance(negative_prompt, list) + ): + raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}") + + if image is not None and video is not None: + raise ValueError("image and video cannot be provided simultaneously") + + def prepare_latents( + self, + batch_size: int, + num_channels_latents: int = 16, + height: int = 384, + width: int = 640, + num_frames: int = 33, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + + num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + shape = ( + batch_size, + num_channels_latents, + num_latent_frames, + int(height) // self.vae_scale_factor_spatial, + int(width) // self.vae_scale_factor_spatial, + ) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + return latents + + def prepare_image_latents( + self, + image: torch.Tensor, + latents_mean: torch.Tensor, + latents_std: torch.Tensor, + num_latent_frames_per_chunk: int, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + fake_latents: torch.Tensor | None = None, + ) -> torch.Tensor: + device = device or self._execution_device + if latents is None: + image = image.unsqueeze(2).to(device=device, dtype=self.vae.dtype) + latents = self.vae.encode(image).latent_dist.sample(generator=generator) + latents = (latents - latents_mean) * latents_std + if fake_latents is None: + min_frames = (num_latent_frames_per_chunk - 1) * self.vae_scale_factor_temporal + 1 + fake_video = image.repeat(1, 1, min_frames, 1, 1).to(device=device, dtype=self.vae.dtype) + fake_latents_full = self.vae.encode(fake_video).latent_dist.sample(generator=generator) + fake_latents_full = (fake_latents_full - latents_mean) * latents_std + fake_latents = fake_latents_full[:, :, -1:, :, :] + return latents.to(device=device, dtype=dtype), fake_latents.to(device=device, dtype=dtype) + + def prepare_video_latents( + self, + video: torch.Tensor, + latents_mean: torch.Tensor, + latents_std: torch.Tensor, + num_latent_frames_per_chunk: int, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + device = device or self._execution_device + video = video.to(device=device, dtype=self.vae.dtype) + if latents is None: + num_frames = video.shape[2] + min_frames = (num_latent_frames_per_chunk - 1) * self.vae_scale_factor_temporal + 1 + num_chunks = num_frames // min_frames + if num_chunks == 0: + raise ValueError( + f"Video must have at least {min_frames} frames " + f"(got {num_frames} frames). " + f"Required: (num_latent_frames_per_chunk - 1) * {self.vae_scale_factor_temporal} + 1 = ({num_latent_frames_per_chunk} - 1) * {self.vae_scale_factor_temporal} + 1 = {min_frames}" + ) + total_valid_frames = num_chunks * min_frames + start_frame = num_frames - total_valid_frames + + first_frame = video[:, :, 0:1, :, :] + first_frame_latent = self.vae.encode(first_frame).latent_dist.sample(generator=generator) + first_frame_latent = (first_frame_latent - latents_mean) * latents_std + + latents_chunks = [] + for i in range(num_chunks): + chunk_start = start_frame + i * min_frames + chunk_end = chunk_start + min_frames + video_chunk = video[:, :, chunk_start:chunk_end, :, :] + chunk_latents = self.vae.encode(video_chunk).latent_dist.sample(generator=generator) + chunk_latents = (chunk_latents - latents_mean) * latents_std + latents_chunks.append(chunk_latents) + latents = torch.cat(latents_chunks, dim=2) + return first_frame_latent.to(device=device, dtype=dtype), latents.to(device=device, dtype=dtype) + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, + height: int = 384, + width: int = 640, + num_frames: int = 132, + num_inference_steps: int = 50, + sigmas: list[float] = None, + guidance_scale: float = 5.0, + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "np", + return_dict: bool = True, + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + max_sequence_length: int = 512, + # ------------ I2V ------------ + image: PipelineImageInput | None = None, + image_latents: torch.Tensor | None = None, + fake_image_latents: torch.Tensor | None = None, + add_noise_to_image_latents: bool = True, + image_noise_sigma_min: float = 0.111, + image_noise_sigma_max: float = 0.135, + # ------------ V2V ------------ + video: PipelineImageInput | None = None, + video_latents: torch.Tensor | None = None, + add_noise_to_video_latents: bool = True, + video_noise_sigma_min: float = 0.111, + video_noise_sigma_max: float = 0.135, + # ------------ Stage 1 ------------ + history_sizes: list = [16, 2, 1], + num_latent_frames_per_chunk: int = 9, + keep_first_frame: bool = True, + is_skip_first_chunk: bool = False, + ): + r""" + The call function to the pipeline for generation. + + Args: + prompt (`str` or `list[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, pass `prompt_embeds` instead. + negative_prompt (`str` or `list[str]`, *optional*): + The prompt or prompts to avoid during image generation. If not defined, pass `negative_prompt_embeds` + instead. Ignored when not using guidance (`guidance_scale` < `1`). + height (`int`, defaults to `384`): + The height in pixels of the generated image. + width (`int`, defaults to `640`): + The width in pixels of the generated image. + num_frames (`int`, defaults to `132`): + The number of frames in the generated video. + num_inference_steps (`int`, defaults to `50`): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + guidance_scale (`float`, defaults to `5.0`): + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. + of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting + `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to + the text `prompt`, usually at the expense of lower image quality. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make + generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor is generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not + provided, text embeddings are generated from the `prompt` input argument. + output_type (`str`, *optional*, defaults to `"np"`): + The output format of the generated image. Choose between `PIL.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`HeliosPipelineOutput`] instead of a plain tuple. + attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): + A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of + each denoising step during the inference. with the following arguments: `callback_on_step_end(self: + DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a + list of all tensors as specified by `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`list`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int`, defaults to `512`): + The maximum sequence length of the text encoder. If the prompt is longer than this, it will be + truncated. If the prompt is shorter, it will be padded to this length. + + Examples: + + Returns: + [`~HeliosPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`HeliosPipelineOutput`] is returned, otherwise a `tuple` is returned where + the first element is a list with the generated images and the second element is a list of `bool`s + indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content. + """ + + history_sizes = sorted(history_sizes, reverse=True) # From big to small + + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + negative_prompt, + height, + width, + prompt_embeds, + negative_prompt_embeds, + callback_on_step_end_tensor_inputs, + image, + video, + ) + + num_frames = max(num_frames, 1) + + self._guidance_scale = guidance_scale + self._attention_kwargs = attention_kwargs + self._current_timestep = None + self._interrupt = False + + device = self._execution_device + vae_dtype = self.vae.dtype + + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(device, self.vae.dtype) + ) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( + device, self.vae.dtype + ) + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + # 3. Encode input prompt + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + num_videos_per_prompt=num_videos_per_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + max_sequence_length=max_sequence_length, + device=device, + ) + + transformer_dtype = self.transformer.dtype + prompt_embeds = prompt_embeds.to(transformer_dtype) + if negative_prompt_embeds is not None: + negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype) + + # 4. Prepare image or video + if image is not None: + image = self.video_processor.preprocess(image, height=height, width=width) + image_latents, fake_image_latents = self.prepare_image_latents( + image, + latents_mean=latents_mean, + latents_std=latents_std, + num_latent_frames_per_chunk=num_latent_frames_per_chunk, + dtype=torch.float32, + device=device, + generator=generator, + latents=image_latents, + fake_latents=fake_image_latents, + ) + + if image_latents is not None and add_noise_to_image_latents: + image_noise_sigma = ( + torch.rand(1, device=device, generator=generator) * (image_noise_sigma_max - image_noise_sigma_min) + + image_noise_sigma_min + ) + image_latents = ( + image_noise_sigma * randn_tensor(image_latents.shape, generator=generator, device=device) + + (1 - image_noise_sigma) * image_latents + ) + fake_image_noise_sigma = ( + torch.rand(1, device=device, generator=generator) * (video_noise_sigma_max - video_noise_sigma_min) + + video_noise_sigma_min + ) + fake_image_latents = ( + fake_image_noise_sigma * randn_tensor(fake_image_latents.shape, generator=generator, device=device) + + (1 - fake_image_noise_sigma) * fake_image_latents + ) + + if video is not None: + video = self.video_processor.preprocess_video(video, height=height, width=width) + image_latents, video_latents = self.prepare_video_latents( + video, + latents_mean=latents_mean, + latents_std=latents_std, + num_latent_frames_per_chunk=num_latent_frames_per_chunk, + dtype=torch.float32, + device=device, + generator=generator, + latents=video_latents, + ) + + if video_latents is not None and add_noise_to_video_latents: + image_noise_sigma = ( + torch.rand(1, device=device, generator=generator) * (image_noise_sigma_max - image_noise_sigma_min) + + image_noise_sigma_min + ) + image_latents = ( + image_noise_sigma * randn_tensor(image_latents.shape, generator=generator, device=device) + + (1 - image_noise_sigma) * image_latents + ) + + noisy_latents_chunks = [] + num_latent_chunks = video_latents.shape[2] // num_latent_frames_per_chunk + for i in range(num_latent_chunks): + chunk_start = i * num_latent_frames_per_chunk + chunk_end = chunk_start + num_latent_frames_per_chunk + latent_chunk = video_latents[:, :, chunk_start:chunk_end, :, :] + + chunk_frames = latent_chunk.shape[2] + frame_sigmas = ( + torch.rand(chunk_frames, device=device, generator=generator) + * (video_noise_sigma_max - video_noise_sigma_min) + + video_noise_sigma_min + ) + frame_sigmas = frame_sigmas.view(1, 1, chunk_frames, 1, 1) + + noisy_chunk = ( + frame_sigmas * randn_tensor(latent_chunk.shape, generator=generator, device=device) + + (1 - frame_sigmas) * latent_chunk + ) + noisy_latents_chunks.append(noisy_chunk) + video_latents = torch.cat(noisy_latents_chunks, dim=2) + + # 5. Prepare latent variables + num_channels_latents = self.transformer.config.in_channels + window_num_frames = (num_latent_frames_per_chunk - 1) * self.vae_scale_factor_temporal + 1 + num_latent_chunk = max(1, (num_frames + window_num_frames - 1) // window_num_frames) + num_history_latent_frames = sum(history_sizes) + history_video = None + total_generated_latent_frames = 0 + + if not keep_first_frame: + history_sizes[-1] = history_sizes[-1] + 1 + history_latents = torch.zeros( + batch_size, + num_channels_latents, + num_history_latent_frames, + height // self.vae_scale_factor_spatial, + width // self.vae_scale_factor_spatial, + device=device, + dtype=torch.float32, + ) + if fake_image_latents is not None: + history_latents = torch.cat([history_latents[:, :, :-1, :, :], fake_image_latents], dim=2) + total_generated_latent_frames += 1 + if video_latents is not None: + history_frames = history_latents.shape[2] + video_frames = video_latents.shape[2] + if video_frames < history_frames: + keep_frames = history_frames - video_frames + history_latents = torch.cat([history_latents[:, :, :keep_frames, :, :], video_latents], dim=2) + else: + history_latents = video_latents + total_generated_latent_frames += video_latents.shape[2] + + if keep_first_frame: + indices = torch.arange(0, sum([1, *history_sizes, num_latent_frames_per_chunk])) + ( + indices_prefix, + indices_latents_history_long, + indices_latents_history_mid, + indices_latents_history_1x, + indices_hidden_states, + ) = indices.split([1, *history_sizes, num_latent_frames_per_chunk], dim=0) + indices_latents_history_short = torch.cat([indices_prefix, indices_latents_history_1x], dim=0) + else: + indices = torch.arange(0, sum([*history_sizes, num_latent_frames_per_chunk])) + ( + indices_latents_history_long, + indices_latents_history_mid, + indices_latents_history_short, + indices_hidden_states, + ) = indices.split([*history_sizes, num_latent_frames_per_chunk], dim=0) + indices_hidden_states = indices_hidden_states.unsqueeze(0) + indices_latents_history_short = indices_latents_history_short.unsqueeze(0) + indices_latents_history_mid = indices_latents_history_mid.unsqueeze(0) + indices_latents_history_long = indices_latents_history_long.unsqueeze(0) + + # 6. Denoising loop + patch_size = self.transformer.config.patch_size + image_seq_len = ( + num_latent_frames_per_chunk + * (height // self.vae_scale_factor_spatial) + * (width // self.vae_scale_factor_spatial) + // (patch_size[0] * patch_size[1] * patch_size[2]) + ) + sigmas = np.linspace(0.999, 0.0, num_inference_steps + 1)[:-1] if sigmas is None else sigmas + mu = calculate_shift( + image_seq_len, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.5), + self.scheduler.config.get("max_shift", 1.15), + ) + + for k in range(num_latent_chunk): + is_first_chunk = k == 0 + is_second_chunk = k == 1 + if keep_first_frame: + latents_history_long, latents_history_mid, latents_history_1x = history_latents[ + :, :, -num_history_latent_frames: + ].split(history_sizes, dim=2) + if image_latents is None and is_first_chunk: + latents_prefix = torch.zeros( + ( + batch_size, + num_channels_latents, + 1, + latents_history_1x.shape[-2], + latents_history_1x.shape[-1], + ), + device=device, + dtype=latents_history_1x.dtype, + ) + else: + latents_prefix = image_latents + latents_history_short = torch.cat([latents_prefix, latents_history_1x], dim=2) + else: + latents_history_long, latents_history_mid, latents_history_short = history_latents[ + :, :, -num_history_latent_frames: + ].split(history_sizes, dim=2) + + latents = self.prepare_latents( + batch_size, + num_channels_latents, + height, + width, + window_num_frames, + dtype=torch.float32, + device=device, + generator=generator, + latents=None, + ) + + self.scheduler.set_timesteps(num_inference_steps, device=device, sigmas=sigmas, mu=mu) + timesteps = self.scheduler.timesteps + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + self._num_timesteps = len(timesteps) + + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + self._current_timestep = t + timestep = t.expand(latents.shape[0]) + + latent_model_input = latents.to(transformer_dtype) + latents_history_short = latents_history_short.to(transformer_dtype) + latents_history_mid = latents_history_mid.to(transformer_dtype) + latents_history_long = latents_history_long.to(transformer_dtype) + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=prompt_embeds, + indices_hidden_states=indices_hidden_states, + indices_latents_history_short=indices_latents_history_short, + indices_latents_history_mid=indices_latents_history_mid, + indices_latents_history_long=indices_latents_history_long, + latents_history_short=latents_history_short, + latents_history_mid=latents_history_mid, + latents_history_long=latents_history_long, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] + + if self.do_classifier_free_guidance: + with self.transformer.cache_context("uncond"): + noise_uncond = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=negative_prompt_embeds, + indices_hidden_states=indices_hidden_states, + indices_latents_history_short=indices_latents_history_short, + indices_latents_history_mid=indices_latents_history_mid, + indices_latents_history_long=indices_latents_history_long, + latents_history_short=latents_history_short, + latents_history_mid=latents_history_mid, + latents_history_long=latents_history_long, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] + noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond) + + latents = self.scheduler.step( + noise_pred, + t, + latents, + generator=generator, + return_dict=False, + )[0] + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + if keep_first_frame and ( + (is_first_chunk and image_latents is None) or (is_skip_first_chunk and is_second_chunk) + ): + image_latents = latents[:, :, 0:1, :, :] + + total_generated_latent_frames += latents.shape[2] + history_latents = torch.cat([history_latents, latents], dim=2) + real_history_latents = history_latents[:, :, -total_generated_latent_frames:] + current_latents = ( + real_history_latents[:, :, -num_latent_frames_per_chunk:].to(vae_dtype) / latents_std + + latents_mean + ) + current_video = self.vae.decode(current_latents, return_dict=False)[0] + + if history_video is None: + history_video = current_video + else: + history_video = torch.cat([history_video, current_video], dim=2) + + self._current_timestep = None + + if output_type != "latent": + generated_frames = history_video.size(2) + generated_frames = ( + generated_frames - 1 + ) // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1 + history_video = history_video[:, :, :generated_frames] + video = self.video_processor.postprocess_video(history_video, output_type=output_type) + else: + video = real_history_latents + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (video,) + + return HeliosPipelineOutput(frames=video) diff --git a/src/diffusers/pipelines/helios/pipeline_helios_pyramid.py b/src/diffusers/pipelines/helios/pipeline_helios_pyramid.py new file mode 100644 index 000000000000..d8f317a9a6f1 --- /dev/null +++ b/src/diffusers/pipelines/helios/pipeline_helios_pyramid.py @@ -0,0 +1,1085 @@ +# Copyright 2025 The Helios Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import html +import math +from typing import Any, Callable + +import regex as re +import torch +import torch.nn.functional as F +from transformers import AutoTokenizer, UMT5EncoderModel + +from ...callbacks import MultiPipelineCallbacks, PipelineCallback +from ...image_processor import PipelineImageInput +from ...loaders import HeliosLoraLoaderMixin +from ...models import AutoencoderKLWan, HeliosTransformer3DModel +from ...schedulers import HeliosDMDScheduler, HeliosScheduler +from ...utils import is_ftfy_available, is_torch_xla_available, logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor +from ..pipeline_utils import DiffusionPipeline +from .pipeline_output import HeliosPipelineOutput + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +if is_ftfy_available(): + import ftfy + + +EXAMPLE_DOC_STRING = """ + Examples: + ```python + >>> import torch + >>> from diffusers.utils import export_to_video + >>> from diffusers import AutoencoderKLWan, HeliosPyramidPipeline + + >>> # Available models: BestWishYsh/Helios-Base, BestWishYsh/Helios-Mid, BestWishYsh/Helios-Distilled + >>> model_id = "BestWishYsh/Helios-Base" + >>> vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32) + >>> pipe = HeliosPyramidPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16) + >>> pipe.to("cuda") + + >>> prompt = "A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window." + >>> negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" + + >>> output = pipe( + ... prompt=prompt, + ... negative_prompt=negative_prompt, + ... height=384, + ... width=640, + ... num_frames=132, + ... guidance_scale=5.0, + ... ).frames[0] + >>> export_to_video(output, "output.mp4", fps=24) + ``` +""" + + +def optimized_scale(positive_flat, negative_flat): + positive_flat = positive_flat.float() + negative_flat = negative_flat.float() + # Calculate dot production + dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True) + # Squared norm of uncondition + squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8 + # st_star = v_cond^T * v_uncond / ||v_uncond||^2 + st_star = dot_product / squared_norm + return st_star + + +def basic_clean(text): + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + return text.strip() + + +def whitespace_clean(text): + text = re.sub(r"\s+", " ", text) + text = text.strip() + return text + + +def prompt_clean(text): + text = whitespace_clean(basic_clean(text)) + return text + + +# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +class HeliosPyramidPipeline(DiffusionPipeline, HeliosLoraLoaderMixin): + r""" + Pipeline for text-to-video / image-to-video / video-to-video generation using Helios. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods + implemented for all pipelines (downloading, saving, running on a particular device, etc.). + + Args: + tokenizer ([`T5Tokenizer`]): + Tokenizer from [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5Tokenizer), + specifically the [google/umt5-xxl](https://huggingface.co/google/umt5-xxl) variant. + text_encoder ([`T5EncoderModel`]): + [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically + the [google/umt5-xxl](https://huggingface.co/google/umt5-xxl) variant. + transformer ([`HeliosTransformer3DModel`]): + Conditional Transformer to denoise the input latents. + scheduler ([`HeliosScheduler`, `HeliosDMDScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + vae ([`AutoencoderKLWan`]): + Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. + """ + + model_cpu_offload_seq = "text_encoder->transformer->vae" + _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] + _optional_components = ["transformer"] + + def __init__( + self, + tokenizer: AutoTokenizer, + text_encoder: UMT5EncoderModel, + vae: AutoencoderKLWan, + scheduler: HeliosScheduler | HeliosDMDScheduler, + transformer: HeliosTransformer3DModel, + is_cfg_zero_star: bool = False, + is_distilled: bool = False, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + scheduler=scheduler, + ) + self.register_to_config(is_cfg_zero_star=is_cfg_zero_star) + self.register_to_config(is_distilled=is_distilled) + self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if getattr(self, "vae", None) else 4 + self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if getattr(self, "vae", None) else 8 + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + + # Copied from diffusers.pipelines.helios.pipeline_helios.HeliosPipeline._get_t5_prompt_embeds + def _get_t5_prompt_embeds( + self, + prompt: str | list[str] = None, + num_videos_per_prompt: int = 1, + max_sequence_length: int = 226, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + prompt = [prompt_clean(u) for u in prompt] + batch_size = len(prompt) + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask + seq_lens = mask.gt(0).sum(dim=1).long() + + prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + prompt_embeds = torch.stack( + [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0 + ) + + # duplicate text embeddings for each generation per prompt, using mps friendly method + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) + + return prompt_embeds, text_inputs.attention_mask.bool() + + def encode_prompt( + self, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, + do_classifier_free_guidance: bool = True, + num_videos_per_prompt: int = 1, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + max_sequence_length: int = 226, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `list[str]`, *optional*): + prompt to be encoded + negative_prompt (`str` or `list[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): + Whether to use classifier free guidance or not. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos that should be generated per prompt. torch device to place the resulting embeddings on + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + device: (`torch.device`, *optional*): + torch device + dtype: (`torch.dtype`, *optional*): + torch dtype + """ + device = device or self._execution_device + + prompt = [prompt] if isinstance(prompt, str) else prompt + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds, _ = self._get_t5_prompt_embeds( + prompt=prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + + if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + + if prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + + negative_prompt_embeds, _ = self._get_t5_prompt_embeds( + prompt=negative_prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + + return prompt_embeds, negative_prompt_embeds + + def check_inputs( + self, + prompt, + negative_prompt, + height, + width, + prompt_embeds=None, + negative_prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + image=None, + video=None, + guidance_scale=None, + ): + if height % 16 != 0 or width % 16 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.") + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`: {negative_prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + elif negative_prompt is not None and ( + not isinstance(negative_prompt, str) and not isinstance(negative_prompt, list) + ): + raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}") + + if image is not None and video is not None: + raise ValueError("image and video cannot be provided simultaneously") + + if guidance_scale > 1.0 and self.config.is_distilled: + logger.warning(f"Guidance scale {guidance_scale} is ignored for step-wise distilled models.") + + def prepare_latents( + self, + batch_size: int, + num_channels_latents: int = 16, + height: int = 384, + width: int = 640, + num_frames: int = 33, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + + num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + shape = ( + batch_size, + num_channels_latents, + num_latent_frames, + int(height) // self.vae_scale_factor_spatial, + int(width) // self.vae_scale_factor_spatial, + ) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + return latents + + def prepare_image_latents( + self, + image: torch.Tensor, + latents_mean: torch.Tensor, + latents_std: torch.Tensor, + num_latent_frames_per_chunk: int, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + fake_latents: torch.Tensor | None = None, + ) -> torch.Tensor: + device = device or self._execution_device + if latents is None: + image = image.unsqueeze(2).to(device=device, dtype=self.vae.dtype) + latents = self.vae.encode(image).latent_dist.sample(generator=generator) + latents = (latents - latents_mean) * latents_std + if fake_latents is None: + min_frames = (num_latent_frames_per_chunk - 1) * self.vae_scale_factor_temporal + 1 + fake_video = image.repeat(1, 1, min_frames, 1, 1).to(device=device, dtype=self.vae.dtype) + fake_latents_full = self.vae.encode(fake_video).latent_dist.sample(generator=generator) + fake_latents_full = (fake_latents_full - latents_mean) * latents_std + fake_latents = fake_latents_full[:, :, -1:, :, :] + return latents.to(device=device, dtype=dtype), fake_latents.to(device=device, dtype=dtype) + + def prepare_video_latents( + self, + video: torch.Tensor, + latents_mean: torch.Tensor, + latents_std: torch.Tensor, + num_latent_frames_per_chunk: int, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + device = device or self._execution_device + video = video.to(device=device, dtype=self.vae.dtype) + if latents is None: + num_frames = video.shape[2] + min_frames = (num_latent_frames_per_chunk - 1) * self.vae_scale_factor_temporal + 1 + num_chunks = num_frames // min_frames + if num_chunks == 0: + raise ValueError( + f"Video must have at least {min_frames} frames " + f"(got {num_frames} frames). " + f"Required: (num_latent_frames_per_chunk - 1) * {self.vae_scale_factor_temporal} + 1 = ({num_latent_frames_per_chunk} - 1) * {self.vae_scale_factor_temporal} + 1 = {min_frames}" + ) + total_valid_frames = num_chunks * min_frames + start_frame = num_frames - total_valid_frames + + first_frame = video[:, :, 0:1, :, :] + first_frame_latent = self.vae.encode(first_frame).latent_dist.sample(generator=generator) + first_frame_latent = (first_frame_latent - latents_mean) * latents_std + + latents_chunks = [] + for i in range(num_chunks): + chunk_start = start_frame + i * min_frames + chunk_end = chunk_start + min_frames + video_chunk = video[:, :, chunk_start:chunk_end, :, :] + chunk_latents = self.vae.encode(video_chunk).latent_dist.sample(generator=generator) + chunk_latents = (chunk_latents - latents_mean) * latents_std + latents_chunks.append(chunk_latents) + latents = torch.cat(latents_chunks, dim=2) + return first_frame_latent.to(device=device, dtype=dtype), latents.to(device=device, dtype=dtype) + + def sample_block_noise( + self, + batch_size, + channel, + num_frames, + height, + width, + patch_size: tuple[int, ...] = (1, 2, 2), + device: torch.device | None = None, + generator: torch.Generator | None = None, + ): + # NOTE: A generator must be provided to ensure correct and reproducible results. + # Creating a default generator here is a fallback only — without a fixed seed, + # the output will be non-deterministic and may produce incorrect results in CP context. + if generator is None: + generator = torch.Generator(device=device) + + gamma = self.scheduler.config.gamma + _, ph, pw = patch_size + block_size = ph * pw + + cov = ( + torch.eye(block_size, device=device) * (1 + gamma) + - torch.ones(block_size, block_size, device=device) * gamma + ) + cov += torch.eye(block_size, device=device) * 1e-8 + cov = cov.float() # Upcast to fp32 for numerical stability — cholesky is unreliable in fp16/bf16. + + L = torch.linalg.cholesky(cov) + block_number = batch_size * channel * num_frames * (height // ph) * (width // pw) + z = torch.randn(block_number, block_size, device=device, generator=generator) + noise = z @ L.T + + noise = noise.view(batch_size, channel, num_frames, height // ph, width // pw, ph, pw) + noise = noise.permute(0, 1, 2, 3, 5, 4, 6).reshape(batch_size, channel, num_frames, height, width) + + return noise + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, + height: int = 384, + width: int = 640, + num_frames: int = 132, + sigmas: list[float] = None, + guidance_scale: float = 5.0, + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "np", + return_dict: bool = True, + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + max_sequence_length: int = 512, + # ------------ I2V ------------ + image: PipelineImageInput | None = None, + image_latents: torch.Tensor | None = None, + fake_image_latents: torch.Tensor | None = None, + add_noise_to_image_latents: bool = True, + image_noise_sigma_min: float = 0.111, + image_noise_sigma_max: float = 0.135, + # ------------ V2V ------------ + video: PipelineImageInput | None = None, + video_latents: torch.Tensor | None = None, + add_noise_to_video_latents: bool = True, + video_noise_sigma_min: float = 0.111, + video_noise_sigma_max: float = 0.135, + # ------------ Stage 1 ------------ + history_sizes: list = [16, 2, 1], + num_latent_frames_per_chunk: int = 9, + keep_first_frame: bool = True, + is_skip_first_chunk: bool = False, + # ------------ Stage 2 ------------ + pyramid_num_inference_steps_list: list = [10, 10, 10], + # ------------ CFG Zero ------------ + use_zero_init: bool | None = True, + zero_steps: int | None = 1, + # ------------ DMD ------------ + is_amplify_first_chunk: bool = False, + ): + r""" + The call function to the pipeline for generation. + + Args: + prompt (`str` or `list[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, pass `prompt_embeds` instead. + negative_prompt (`str` or `list[str]`, *optional*): + The prompt or prompts to avoid during image generation. If not defined, pass `negative_prompt_embeds` + instead. Ignored when not using guidance (`guidance_scale` < `1`). + height (`int`, defaults to `384`): + The height in pixels of the generated image. + width (`int`, defaults to `640`): + The width in pixels of the generated image. + num_frames (`int`, defaults to `132`): + The number of frames in the generated video. + guidance_scale (`float`, defaults to `5.0`): + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. + of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting + `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to + the text `prompt`, usually at the expense of lower image quality. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make + generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor is generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not + provided, text embeddings are generated from the `prompt` input argument. + output_type (`str`, *optional*, defaults to `"np"`): + The output format of the generated image. Choose between `PIL.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`HeliosPipelineOutput`] instead of a plain tuple. + attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): + A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of + each denoising step during the inference. with the following arguments: `callback_on_step_end(self: + DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a + list of all tensors as specified by `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`list`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int`, defaults to `512`): + The maximum sequence length of the text encoder. If the prompt is longer than this, it will be + truncated. If the prompt is shorter, it will be padded to this length. + + Examples: + + Returns: + [`~HeliosPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`HeliosPipelineOutput`] is returned, otherwise a `tuple` is returned where + the first element is a list with the generated images and the second element is a list of `bool`s + indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content. + """ + + history_sizes = sorted(history_sizes, reverse=True) # From big to small + pyramid_num_stages = len(pyramid_num_inference_steps_list) + + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + negative_prompt, + height, + width, + prompt_embeds, + negative_prompt_embeds, + callback_on_step_end_tensor_inputs, + image, + video, + guidance_scale, + ) + + num_frames = max(num_frames, 1) + + self._guidance_scale = guidance_scale + self._attention_kwargs = attention_kwargs + self._current_timestep = None + self._interrupt = False + + device = self._execution_device + vae_dtype = self.vae.dtype + + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(device, self.vae.dtype) + ) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( + device, self.vae.dtype + ) + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + # 3. Encode input prompt + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + num_videos_per_prompt=num_videos_per_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + max_sequence_length=max_sequence_length, + device=device, + ) + + transformer_dtype = self.transformer.dtype + prompt_embeds = prompt_embeds.to(transformer_dtype) + if negative_prompt_embeds is not None: + negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype) + + # 4. Prepare image or video + if image is not None: + image = self.video_processor.preprocess(image, height=height, width=width) + image_latents, fake_image_latents = self.prepare_image_latents( + image, + latents_mean=latents_mean, + latents_std=latents_std, + num_latent_frames_per_chunk=num_latent_frames_per_chunk, + dtype=torch.float32, + device=device, + generator=generator, + latents=image_latents, + fake_latents=fake_image_latents, + ) + + if image_latents is not None and add_noise_to_image_latents: + image_noise_sigma = ( + torch.rand(1, device=device, generator=generator) * (image_noise_sigma_max - image_noise_sigma_min) + + image_noise_sigma_min + ) + image_latents = ( + image_noise_sigma * randn_tensor(image_latents.shape, generator=generator, device=device) + + (1 - image_noise_sigma) * image_latents + ) + fake_image_noise_sigma = ( + torch.rand(1, device=device, generator=generator) * (video_noise_sigma_max - video_noise_sigma_min) + + video_noise_sigma_min + ) + fake_image_latents = ( + fake_image_noise_sigma * randn_tensor(fake_image_latents.shape, generator=generator, device=device) + + (1 - fake_image_noise_sigma) * fake_image_latents + ) + + if video is not None: + video = self.video_processor.preprocess_video(video, height=height, width=width) + image_latents, video_latents = self.prepare_video_latents( + video, + latents_mean=latents_mean, + latents_std=latents_std, + num_latent_frames_per_chunk=num_latent_frames_per_chunk, + dtype=torch.float32, + device=device, + generator=generator, + latents=video_latents, + ) + + if video_latents is not None and add_noise_to_video_latents: + image_noise_sigma = ( + torch.rand(1, device=device, generator=generator) * (image_noise_sigma_max - image_noise_sigma_min) + + image_noise_sigma_min + ) + image_latents = ( + image_noise_sigma * randn_tensor(image_latents.shape, generator=generator, device=device) + + (1 - image_noise_sigma) * image_latents + ) + + noisy_latents_chunks = [] + num_latent_chunks = video_latents.shape[2] // num_latent_frames_per_chunk + for i in range(num_latent_chunks): + chunk_start = i * num_latent_frames_per_chunk + chunk_end = chunk_start + num_latent_frames_per_chunk + latent_chunk = video_latents[:, :, chunk_start:chunk_end, :, :] + + chunk_frames = latent_chunk.shape[2] + frame_sigmas = ( + torch.rand(chunk_frames, device=device, generator=generator) + * (video_noise_sigma_max - video_noise_sigma_min) + + video_noise_sigma_min + ) + frame_sigmas = frame_sigmas.view(1, 1, chunk_frames, 1, 1) + + noisy_chunk = ( + frame_sigmas * randn_tensor(latent_chunk.shape, generator=generator, device=device) + + (1 - frame_sigmas) * latent_chunk + ) + noisy_latents_chunks.append(noisy_chunk) + video_latents = torch.cat(noisy_latents_chunks, dim=2) + + # 5. Prepare latent variables + num_channels_latents = self.transformer.config.in_channels + window_num_frames = (num_latent_frames_per_chunk - 1) * self.vae_scale_factor_temporal + 1 + num_latent_chunk = max(1, (num_frames + window_num_frames - 1) // window_num_frames) + num_history_latent_frames = sum(history_sizes) + history_video = None + total_generated_latent_frames = 0 + + if not keep_first_frame: + history_sizes[-1] = history_sizes[-1] + 1 + history_latents = torch.zeros( + batch_size, + num_channels_latents, + num_history_latent_frames, + height // self.vae_scale_factor_spatial, + width // self.vae_scale_factor_spatial, + device=device, + dtype=torch.float32, + ) + if fake_image_latents is not None: + history_latents = torch.cat([history_latents[:, :, :-1, :, :], fake_image_latents], dim=2) + total_generated_latent_frames += 1 + if video_latents is not None: + history_frames = history_latents.shape[2] + video_frames = video_latents.shape[2] + if video_frames < history_frames: + keep_frames = history_frames - video_frames + history_latents = torch.cat([history_latents[:, :, :keep_frames, :, :], video_latents], dim=2) + else: + history_latents = video_latents + total_generated_latent_frames += video_latents.shape[2] + + if keep_first_frame: + indices = torch.arange(0, sum([1, *history_sizes, num_latent_frames_per_chunk])) + ( + indices_prefix, + indices_latents_history_long, + indices_latents_history_mid, + indices_latents_history_1x, + indices_hidden_states, + ) = indices.split([1, *history_sizes, num_latent_frames_per_chunk], dim=0) + indices_latents_history_short = torch.cat([indices_prefix, indices_latents_history_1x], dim=0) + else: + indices = torch.arange(0, sum([*history_sizes, num_latent_frames_per_chunk])) + ( + indices_latents_history_long, + indices_latents_history_mid, + indices_latents_history_short, + indices_hidden_states, + ) = indices.split([*history_sizes, num_latent_frames_per_chunk], dim=0) + indices_hidden_states = indices_hidden_states.unsqueeze(0) + indices_latents_history_short = indices_latents_history_short.unsqueeze(0) + indices_latents_history_mid = indices_latents_history_mid.unsqueeze(0) + indices_latents_history_long = indices_latents_history_long.unsqueeze(0) + + # 6. Denoising loop + for k in range(num_latent_chunk): + is_first_chunk = k == 0 + is_second_chunk = k == 1 + if keep_first_frame: + latents_history_long, latents_history_mid, latents_history_1x = history_latents[ + :, :, -num_history_latent_frames: + ].split(history_sizes, dim=2) + if image_latents is None and is_first_chunk: + latents_prefix = torch.zeros( + ( + batch_size, + num_channels_latents, + 1, + latents_history_1x.shape[-2], + latents_history_1x.shape[-1], + ), + device=device, + dtype=latents_history_1x.dtype, + ) + else: + latents_prefix = image_latents + latents_history_short = torch.cat([latents_prefix, latents_history_1x], dim=2) + else: + latents_history_long, latents_history_mid, latents_history_short = history_latents[ + :, :, -num_history_latent_frames: + ].split(history_sizes, dim=2) + + latents = self.prepare_latents( + batch_size, + num_channels_latents, + height, + width, + window_num_frames, + dtype=torch.float32, + device=device, + generator=generator, + latents=None, + ) + + num_inference_steps = ( + sum(pyramid_num_inference_steps_list) * 2 + if is_amplify_first_chunk and self.config.is_distilled and is_first_chunk + else sum(pyramid_num_inference_steps_list) + ) + + with self.progress_bar(total=num_inference_steps) as progress_bar: + _, _, _, pyramid_height, pyramid_width = latents.shape + latents = latents.permute(0, 2, 1, 3, 4).reshape( + batch_size * num_latent_frames_per_chunk, num_channels_latents, pyramid_height, pyramid_width + ) + for _ in range(pyramid_num_stages - 1): + pyramid_height //= 2 + pyramid_width //= 2 + latents = ( + F.interpolate( + latents, + size=(pyramid_height, pyramid_width), + mode="bilinear", + ) + * 2 + ) + latents = latents.reshape( + batch_size, num_latent_frames_per_chunk, num_channels_latents, pyramid_height, pyramid_width + ).permute(0, 2, 1, 3, 4) + + start_point_list = None + if self.config.is_distilled: + start_point_list = [latents] + + for stage_idx in range(pyramid_num_stages): + patch_size = self.transformer.config.patch_size + image_seq_len = (latents.shape[-1] * latents.shape[-2] * latents.shape[-3]) // ( + patch_size[0] * patch_size[1] * patch_size[2] + ) + mu = calculate_shift( + image_seq_len, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.5), + self.scheduler.config.get("max_shift", 1.15), + ) + self.scheduler.set_timesteps( + pyramid_num_inference_steps_list[stage_idx], + stage_idx, + device=device, + mu=mu, + is_amplify_first_chunk=is_amplify_first_chunk and is_first_chunk, + ) + timesteps = self.scheduler.timesteps + num_warmup_steps = 0 + self._num_timesteps = len(timesteps) + + if stage_idx > 0: + pyramid_height *= 2 + pyramid_width *= 2 + num_frames = latents.shape[2] + latents = latents.permute(0, 2, 1, 3, 4).reshape( + batch_size * num_latent_frames_per_chunk, + num_channels_latents, + pyramid_height // 2, + pyramid_width // 2, + ) + latents = F.interpolate(latents, size=(pyramid_height, pyramid_width), mode="nearest") + latents = latents.reshape( + batch_size, + num_latent_frames_per_chunk, + num_channels_latents, + pyramid_height, + pyramid_width, + ).permute(0, 2, 1, 3, 4) + # Fix the stage + ori_sigma = 1 - self.scheduler.ori_start_sigmas[stage_idx] # the original coeff of signal + gamma = self.scheduler.config.gamma + alpha = 1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma) + beta = alpha * (1 - ori_sigma) / math.sqrt(gamma) + + batch_size, channel, num_frames, pyramid_height, pyramid_width = latents.shape + noise = self.sample_block_noise( + batch_size, + channel, + num_frames, + pyramid_height, + pyramid_width, + patch_size, + device, + generator, + ) + noise = noise.to(device=device, dtype=transformer_dtype) + latents = alpha * latents + beta * noise # To fix the block artifact + + if self.config.is_distilled: + start_point_list.append(latents) + + for i, t in enumerate(timesteps): + timestep = t.expand(latents.shape[0]).to(torch.int64) + + latent_model_input = latents.to(transformer_dtype) + latents_history_short = latents_history_short.to(transformer_dtype) + latents_history_mid = latents_history_mid.to(transformer_dtype) + latents_history_long = latents_history_long.to(transformer_dtype) + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=prompt_embeds, + indices_hidden_states=indices_hidden_states, + indices_latents_history_short=indices_latents_history_short, + indices_latents_history_mid=indices_latents_history_mid, + indices_latents_history_long=indices_latents_history_long, + latents_history_short=latents_history_short, + latents_history_mid=latents_history_mid, + latents_history_long=latents_history_long, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] + + if self.do_classifier_free_guidance: + with self.transformer.cache_context("uncond"): + noise_uncond = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=negative_prompt_embeds, + indices_hidden_states=indices_hidden_states, + indices_latents_history_short=indices_latents_history_short, + indices_latents_history_mid=indices_latents_history_mid, + indices_latents_history_long=indices_latents_history_long, + latents_history_short=latents_history_short, + latents_history_mid=latents_history_mid, + latents_history_long=latents_history_long, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] + + if self.config.is_cfg_zero_star: + noise_pred_text = noise_pred + positive_flat = noise_pred_text.view(batch_size, -1) + negative_flat = noise_uncond.view(batch_size, -1) + + alpha = optimized_scale(positive_flat, negative_flat) + alpha = alpha.view(batch_size, *([1] * (len(noise_pred_text.shape) - 1))) + alpha = alpha.to(noise_pred_text.dtype) + + if (stage_idx == 0 and i <= zero_steps) and use_zero_init: + noise_pred = noise_pred_text * 0.0 + else: + noise_pred = noise_uncond * alpha + guidance_scale * ( + noise_pred_text - noise_uncond * alpha + ) + else: + noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond) + + extra_kwargs = ( + { + "cur_sampling_step": i, + "dmd_noisy_tensor": start_point_list[stage_idx] + if start_point_list is not None + else None, + "dmd_sigmas": self.scheduler.sigmas, + "dmd_timesteps": self.scheduler.timesteps, + "all_timesteps": timesteps, + } + if self.config.is_distilled + else {} + ) + + latents = self.scheduler.step( + noise_pred, + t, + latents, + generator=generator, + return_dict=False, + **extra_kwargs, + )[0] + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop( + "negative_prompt_embeds", negative_prompt_embeds + ) + + if i == len(timesteps) - 1 or ( + (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0 + ): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + if keep_first_frame and ( + (is_first_chunk and image_latents is None) or (is_skip_first_chunk and is_second_chunk) + ): + image_latents = latents[:, :, 0:1, :, :] + + total_generated_latent_frames += latents.shape[2] + history_latents = torch.cat([history_latents, latents], dim=2) + real_history_latents = history_latents[:, :, -total_generated_latent_frames:] + current_latents = ( + real_history_latents[:, :, -num_latent_frames_per_chunk:].to(vae_dtype) / latents_std + + latents_mean + ) + current_video = self.vae.decode(current_latents, return_dict=False)[0] + + if history_video is None: + history_video = current_video + else: + history_video = torch.cat([history_video, current_video], dim=2) + + self._current_timestep = None + + if output_type != "latent": + generated_frames = history_video.size(2) + generated_frames = ( + generated_frames - 1 + ) // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1 + history_video = history_video[:, :, :generated_frames] + video = self.video_processor.postprocess_video(history_video, output_type=output_type) + else: + video = real_history_latents + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (video,) + + return HeliosPipelineOutput(frames=video) diff --git a/src/diffusers/pipelines/helios/pipeline_output.py b/src/diffusers/pipelines/helios/pipeline_output.py new file mode 100644 index 000000000000..08546289ef4c --- /dev/null +++ b/src/diffusers/pipelines/helios/pipeline_output.py @@ -0,0 +1,20 @@ +from dataclasses import dataclass + +import torch + +from diffusers.utils import BaseOutput + + +@dataclass +class HeliosPipelineOutput(BaseOutput): + r""" + Output class for Helios pipelines. + + Args: + frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): + List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape + `(batch_size, num_frames, channels, height, width)`. + """ + + frames: torch.Tensor diff --git a/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py b/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py index b6af23bca8fd..8e5e078cc2af 100644 --- a/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py +++ b/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py @@ -14,7 +14,7 @@ import inspect import math -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch from transformers import ( @@ -53,7 +53,6 @@ >>> from transformers import AutoTokenizer, LlamaForCausalLM >>> from diffusers import HiDreamImagePipeline - >>> tokenizer_4 = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B-Instruct") >>> text_encoder_4 = LlamaForCausalLM.from_pretrained( ... "meta-llama/Meta-Llama-3.1-8B-Instruct", @@ -100,10 +99,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -118,15 +117,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -202,10 +201,10 @@ def __init__( def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, max_sequence_length: int = 128, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder_3.dtype @@ -241,10 +240,10 @@ def _get_clip_prompt_embeds( self, tokenizer, text_encoder, - prompt: Union[str, List[str]], + prompt: str | list[str], max_sequence_length: int = 128, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or text_encoder.dtype @@ -276,10 +275,10 @@ def _get_clip_prompt_embeds( def _get_llama3_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, max_sequence_length: int = 128, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder_4.dtype @@ -320,26 +319,26 @@ def _get_llama3_prompt_embeds( def encode_prompt( self, - prompt: Optional[Union[str, List[str]]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - prompt_3: Optional[Union[str, List[str]]] = None, - prompt_4: Optional[Union[str, List[str]]] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt: str | list[str] | None = None, + prompt_2: str | list[str] | None = None, + prompt_3: str | list[str] | None = None, + prompt_4: str | list[str] | None = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt_3: Optional[Union[str, List[str]]] = None, - negative_prompt_4: Optional[Union[str, List[str]]] = None, - prompt_embeds_t5: Optional[List[torch.FloatTensor]] = None, - prompt_embeds_llama3: Optional[List[torch.FloatTensor]] = None, - negative_prompt_embeds_t5: Optional[List[torch.FloatTensor]] = None, - negative_prompt_embeds_llama3: Optional[List[torch.FloatTensor]] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_3: str | list[str] | None = None, + negative_prompt_4: str | list[str] | None = None, + prompt_embeds_t5: list[torch.FloatTensor] | None = None, + prompt_embeds_llama3: list[torch.FloatTensor] | None = None, + negative_prompt_embeds_t5: list[torch.FloatTensor] | None = None, + negative_prompt_embeds_llama3: list[torch.FloatTensor] | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 128, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): prompt = [prompt] if isinstance(prompt, str) else prompt if prompt is not None: @@ -729,33 +728,33 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - prompt_3: Optional[Union[str, List[str]]] = None, - prompt_4: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + prompt_3: str | list[str] | None = None, + prompt_4: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt_3: Optional[Union[str, List[str]]] = None, - negative_prompt_4: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds_t5: Optional[torch.FloatTensor] = None, - prompt_embeds_llama3: Optional[torch.FloatTensor] = None, - negative_prompt_embeds_t5: Optional[torch.FloatTensor] = None, - negative_prompt_embeds_llama3: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_3: str | list[str] | None = None, + negative_prompt_4: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds_t5: torch.FloatTensor | None = None, + prompt_embeds_llama3: torch.FloatTensor | None = None, + negative_prompt_embeds_t5: torch.FloatTensor | None = None, + negative_prompt_embeds_llama3: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 128, **kwargs, ): @@ -763,16 +762,16 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead. - prompt_3 (`str` or `List[str]`, *optional*): + prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is will be used instead. - prompt_4 (`str` or `List[str]`, *optional*): + prompt_4 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_4` and `text_encoder_4`. If not defined, `prompt` is will be used instead. height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): @@ -782,7 +781,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -792,22 +791,22 @@ def __call__( Guidance-distilled models approximates true classifer-free guidance for `guidance_scale` > 1. Refer to the [paper](https://huggingface.co/papers/2210.03142) to learn more. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is not greater than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders. - negative_prompt_3 (`str` or `List[str]`, *optional*): + negative_prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `negative_prompt` is used in all the text-encoders. - negative_prompt_4 (`str` or `List[str]`, *optional*): + negative_prompt_4 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_4` and `text_encoder_4`. If not defined, `negative_prompt` is used in all the text-encoders. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -842,7 +841,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -965,14 +964,18 @@ def __call__( # 5. Prepare timesteps mu = calculate_shift(self.transformer.max_seq) scheduler_kwargs = {"mu": mu} + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device if isinstance(self.scheduler, UniPCMultistepScheduler): - self.scheduler.set_timesteps(num_inference_steps, device=device) # , shift=math.exp(mu)) + self.scheduler.set_timesteps(num_inference_steps, device=timestep_device) # , shift=math.exp(mu)) timesteps = self.scheduler.timesteps else: timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, - device, + timestep_device, sigmas=sigmas, **scheduler_kwargs, ) diff --git a/src/diffusers/pipelines/hidream_image/pipeline_output.py b/src/diffusers/pipelines/hidream_image/pipeline_output.py index 66f0f1260d18..1802c7220691 100644 --- a/src/diffusers/pipelines/hidream_image/pipeline_output.py +++ b/src/diffusers/pipelines/hidream_image/pipeline_output.py @@ -13,7 +13,6 @@ # limitations under the License. from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -27,9 +26,9 @@ class HiDreamImagePipelineOutput(BaseOutput): Output class for HiDreamImage pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image] | np.ndarray diff --git a/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage.py b/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage.py index 658935ccd886..50239e9afa22 100644 --- a/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage.py +++ b/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage.py @@ -12,9 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import inspect import re -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -98,10 +100,10 @@ def extract_glyph_text(prompt: str): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -116,15 +118,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -193,8 +195,8 @@ def __init__( text_encoder_2: T5EncoderModel, tokenizer_2: ByT5Tokenizer, transformer: HunyuanImageTransformer2DModel, - guider: Optional[AdaptiveProjectedMixGuidance] = None, - ocr_guider: Optional[AdaptiveProjectedMixGuidance] = None, + guider: AdaptiveProjectedMixGuidance | None = None, + ocr_guider: AdaptiveProjectedMixGuidance | None = None, ): super().__init__() @@ -222,9 +224,9 @@ def _get_qwen_prompt_embeds( self, tokenizer: Qwen2Tokenizer, text_encoder: Qwen2_5_VLForConditionalGeneration, - prompt: Union[str, List[str]] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt: str | list[str] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, tokenizer_max_length: int = 1000, template: str = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>", drop_idx: int = 34, @@ -260,8 +262,8 @@ def _get_byt5_prompt_embeds( tokenizer: ByT5Tokenizer, text_encoder: T5EncoderModel, prompt: str, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, tokenizer_max_length: int = 128, ): device = device or self._execution_device @@ -293,19 +295,19 @@ def _get_byt5_prompt_embeds( def encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + device: torch.device | None = None, batch_size: int = 1, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, - prompt_embeds_2: Optional[torch.Tensor] = None, - prompt_embeds_mask_2: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + prompt_embeds_2: torch.Tensor | None = None, + prompt_embeds_mask_2: torch.Tensor | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -503,38 +505,38 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - distilled_guidance_scale: Optional[float] = 3.25, - sigmas: Optional[List[float]] = None, + distilled_guidance_scale: float | None = 3.25, + sigmas: list[float] | None = None, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_mask: Optional[torch.Tensor] = None, - prompt_embeds_2: Optional[torch.Tensor] = None, - prompt_embeds_mask_2: Optional[torch.Tensor] = None, - negative_prompt_embeds_2: Optional[torch.Tensor] = None, - negative_prompt_embeds_mask_2: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_mask: torch.Tensor | None = None, + prompt_embeds_2: torch.Tensor | None = None, + prompt_embeds_mask_2: torch.Tensor | None = None, + negative_prompt_embeds_2: torch.Tensor | None = None, + negative_prompt_embeds_mask_2: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined and negative_prompt_embeds is not provided, will use an empty negative prompt. Ignored when not using guidance. ). height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): @@ -544,7 +546,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -558,7 +560,7 @@ def __call__( ignored. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage_refiner.py b/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage_refiner.py index f38f53d9a562..93e4deb2974a 100644 --- a/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage_refiner.py +++ b/src/diffusers/pipelines/hunyuan_image/pipeline_hunyuanimage_refiner.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -62,10 +64,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -80,15 +82,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -121,7 +123,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -161,7 +163,7 @@ def __init__( text_encoder: Qwen2_5_VLForConditionalGeneration, tokenizer: Qwen2Tokenizer, transformer: HunyuanImageTransformer2DModel, - guider: Optional[AdaptiveProjectedMixGuidance] = None, + guider: AdaptiveProjectedMixGuidance | None = None, ): super().__init__() @@ -187,9 +189,9 @@ def _get_qwen_prompt_embeds( self, tokenizer: Qwen2Tokenizer, text_encoder: Qwen2_5_VLForConditionalGeneration, - prompt: Union[str, List[str]] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt: str | list[str] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, tokenizer_max_length: int = 1000, template: str = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>", drop_idx: int = 34, @@ -222,17 +224,17 @@ def _get_qwen_prompt_embeds( def encode_prompt( self, - prompt: Optional[Union[str, List[str]]] = None, - device: Optional[torch.device] = None, + prompt: str | list[str] | None = None, + device: torch.device | None = None, batch_size: int = 1, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -435,35 +437,35 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, - distilled_guidance_scale: Optional[float] = 3.25, - image: Optional[PipelineImageInput] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, + distilled_guidance_scale: float | None = 3.25, + image: PipelineImageInput | None = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 4, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, will use an empty negative prompt. Ignored when not using guidance. distilled_guidance_scale (`float`, *optional*, defaults to None): @@ -482,13 +484,13 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/hunyuan_image/pipeline_output.py b/src/diffusers/pipelines/hunyuan_image/pipeline_output.py index 1e76892a0e81..169436b7d86f 100644 --- a/src/diffusers/pipelines/hunyuan_image/pipeline_output.py +++ b/src/diffusers/pipelines/hunyuan_image/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -13,9 +12,9 @@ class HunyuanImagePipelineOutput(BaseOutput): Output class for HunyuanImage pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) + images (`list[PIL.Image.Image]` or `np.ndarray`) List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image, np.ndarray] diff --git a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_skyreels_image2video.py b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_skyreels_image2video.py index b50a6ae3ed27..1a7cae256d63 100644 --- a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_skyreels_image2video.py +++ b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_skyreels_image2video.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import torch @@ -95,10 +95,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -113,15 +113,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -154,7 +154,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -225,14 +225,14 @@ def __init__( # Copied from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video.HunyuanVideoPipeline._get_llama_prompt_embeds def _get_llama_prompt_embeds( self, - prompt: Union[str, List[str]], - prompt_template: Dict[str, Any], + prompt: str | list[str], + prompt_template: dict[str, Any], num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, max_sequence_length: int = 256, num_hidden_layers_to_skip: int = 2, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -292,10 +292,10 @@ def _get_llama_prompt_embeds( # Copied from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video.HunyuanVideoPipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, max_sequence_length: int = 77, ) -> torch.Tensor: device = device or self._execution_device @@ -332,15 +332,15 @@ def _get_clip_prompt_embeds( # Copied from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video.HunyuanVideoPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]] = None, - prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE, + prompt: str | list[str], + prompt_2: str | list[str] = None, + prompt_template: dict[str, Any] = DEFAULT_PROMPT_TEMPLATE, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, max_sequence_length: int = 256, ): if prompt_embeds is None: @@ -422,10 +422,10 @@ def prepare_latents( height: int = 544, width: int = 960, num_frames: int = 97, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( @@ -536,51 +536,49 @@ def interrupt(self): def __call__( self, image: PipelineImageInput, - prompt: Union[str, List[str]] = None, - prompt_2: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, - negative_prompt_2: Union[str, List[str]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] = None, + negative_prompt: str | list[str] = None, + negative_prompt_2: str | list[str] = None, height: int = 544, width: int = 960, num_frames: int = 97, num_inference_steps: int = 50, - sigmas: List[float] = None, + sigmas: list[float] = None, true_cfg_scale: float = 6.0, guidance_scale: float = 1.0, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], - prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE, + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + prompt_template: dict[str, Any] = DEFAULT_PROMPT_TEMPLATE, max_sequence_length: int = 256, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is not greater than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders. height (`int`, defaults to `720`): @@ -592,7 +590,7 @@ def __call__( num_inference_steps (`int`, defaults to `50`): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -608,7 +606,7 @@ def __call__( conditional latent is not applied. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -645,7 +643,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -728,7 +726,13 @@ def __call__( # 4. Prepare timesteps sigmas = np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1] if sigmas is None else sigmas - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, sigmas=sigmas) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, sigmas=sigmas + ) # 5. Prepare latent variables vae_dtype = self.vae.dtype diff --git a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video.py b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video.py index 5c8e295eaf4c..3c6ec39398ef 100644 --- a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video.py +++ b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import torch @@ -84,10 +84,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -102,15 +102,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -198,14 +198,14 @@ def __init__( def _get_llama_prompt_embeds( self, - prompt: Union[str, List[str]], - prompt_template: Dict[str, Any], + prompt: str | list[str], + prompt_template: dict[str, Any], num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, max_sequence_length: int = 256, num_hidden_layers_to_skip: int = 2, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -264,10 +264,10 @@ def _get_llama_prompt_embeds( def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, max_sequence_length: int = 77, ) -> torch.Tensor: device = device or self._execution_device @@ -303,15 +303,15 @@ def _get_clip_prompt_embeds( def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]] = None, - prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE, + prompt: str | list[str], + prompt_2: str | list[str] = None, + prompt_template: dict[str, Any] = DEFAULT_PROMPT_TEMPLATE, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, max_sequence_length: int = 256, ): if prompt_embeds is None: @@ -391,10 +391,10 @@ def prepare_latents( height: int = 720, width: int = 1280, num_frames: int = 129, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: if latents is not None: return latents.to(device=device, dtype=dtype) @@ -492,51 +492,49 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, - negative_prompt_2: Union[str, List[str]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] = None, + negative_prompt: str | list[str] = None, + negative_prompt_2: str | list[str] = None, height: int = 720, width: int = 1280, num_frames: int = 129, num_inference_steps: int = 50, - sigmas: List[float] = None, + sigmas: list[float] = None, true_cfg_scale: float = 1.0, guidance_scale: float = 6.0, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], - prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE, + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + prompt_template: dict[str, Any] = DEFAULT_PROMPT_TEMPLATE, max_sequence_length: int = 256, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is not greater than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders. height (`int`, defaults to `720`): @@ -548,7 +546,7 @@ def __call__( num_inference_steps (`int`, defaults to `50`): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -563,7 +561,7 @@ def __call__( the [paper](https://huggingface.co/papers/2210.03142) to learn more. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -600,7 +598,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -683,7 +681,13 @@ def __call__( # 4. Prepare timesteps sigmas = np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1] if sigmas is None else sigmas - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, sigmas=sigmas) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, sigmas=sigmas + ) # 5. Prepare latent variables num_channels_latents = self.transformer.config.in_channels diff --git a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_framepack.py b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_framepack.py index 8006514f47ea..f82f26eea5b9 100644 --- a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_framepack.py +++ b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_framepack.py @@ -15,7 +15,7 @@ import inspect import math from enum import Enum -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import torch @@ -178,10 +178,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -196,15 +196,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -302,14 +302,14 @@ def __init__( # Copied from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video.HunyuanVideoPipeline._get_llama_prompt_embeds def _get_llama_prompt_embeds( self, - prompt: Union[str, List[str]], - prompt_template: Dict[str, Any], + prompt: str | list[str], + prompt_template: dict[str, Any], num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, max_sequence_length: int = 256, num_hidden_layers_to_skip: int = 2, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -369,10 +369,10 @@ def _get_llama_prompt_embeds( # Copied from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video.HunyuanVideoPipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, max_sequence_length: int = 77, ) -> torch.Tensor: device = device or self._execution_device @@ -409,15 +409,15 @@ def _get_clip_prompt_embeds( # Copied from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video.HunyuanVideoPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]] = None, - prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE, + prompt: str | list[str], + prompt_2: str | list[str] = None, + prompt_template: dict[str, Any] = DEFAULT_PROMPT_TEMPLATE, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, max_sequence_length: int = 256, ): if prompt_embeds is None: @@ -443,9 +443,7 @@ def encode_prompt( return prompt_embeds, pooled_prompt_embeds, prompt_attention_mask - def encode_image( - self, image: torch.Tensor, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None - ): + def encode_image(self, image: torch.Tensor, device: torch.device | None = None, dtype: torch.dtype | None = None): device = device or self._execution_device image = (image + 1) / 2.0 # [-1, 1] -> [0, 1] image = self.feature_extractor(images=image, return_tensors="pt", do_rescale=False).to( @@ -528,10 +526,10 @@ def prepare_latents( height: int = 720, width: int = 1280, num_frames: int = 129, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: if latents is not None: return latents.to(device=device, dtype=dtype) @@ -553,10 +551,10 @@ def prepare_latents( def prepare_image_latents( self, image: torch.Tensor, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: device = device or self._execution_device if latents is None: @@ -643,37 +641,35 @@ def interrupt(self): def __call__( self, image: PipelineImageInput, - last_image: Optional[PipelineImageInput] = None, - prompt: Union[str, List[str]] = None, - prompt_2: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, - negative_prompt_2: Union[str, List[str]] = None, + last_image: PipelineImageInput | None = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] = None, + negative_prompt: str | list[str] = None, + negative_prompt_2: str | list[str] = None, height: int = 720, width: int = 1280, num_frames: int = 129, latent_window_size: int = 9, num_inference_steps: int = 50, - sigmas: List[float] = None, + sigmas: list[float] = None, true_cfg_scale: float = 1.0, guidance_scale: float = 6.0, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - image_latents: Optional[torch.Tensor] = None, - last_image_latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + image_latents: torch.Tensor | None = None, + last_image_latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], - prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE, + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + prompt_template: dict[str, Any] = DEFAULT_PROMPT_TEMPLATE, max_sequence_length: int = 256, sampling_type: FramepackSamplingType = FramepackSamplingType.INVERTED_ANTI_DRIFTING, ): @@ -686,17 +682,17 @@ def __call__( last_image (`PIL.Image.Image` or `np.ndarray` or `torch.Tensor`, *optional*): The optional last image to be used as the ending point for the video generation. This is useful for generating transitions between two images. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is not greater than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders. height (`int`, defaults to `720`): @@ -708,7 +704,7 @@ def __call__( num_inference_steps (`int`, defaults to `50`): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -724,7 +720,7 @@ def __call__( conditional latent is not applied. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. image_latents (`torch.Tensor`, *optional*): @@ -761,7 +757,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_image2video.py b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_image2video.py index aa04e6509730..c599488c2379 100644 --- a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_image2video.py +++ b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_image2video.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -147,10 +147,10 @@ def _expand_input_ids_with_image_tokens( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -165,15 +165,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -206,7 +206,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -279,15 +279,15 @@ def __init__( def _get_llama_prompt_embeds( self, image: torch.Tensor, - prompt: Union[str, List[str]], - prompt_template: Dict[str, Any], + prompt: str | list[str], + prompt_template: dict[str, Any], num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, max_sequence_length: int = 256, num_hidden_layers_to_skip: int = 2, image_embed_interleave: int = 2, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -417,10 +417,10 @@ def _get_llama_prompt_embeds( def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, max_sequence_length: int = 77, ) -> torch.Tensor: device = device or self._execution_device @@ -451,18 +451,18 @@ def _get_clip_prompt_embeds( def encode_prompt( self, image: torch.Tensor, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]] = None, - prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE, + prompt: str | list[str], + prompt_2: str | list[str] = None, + prompt_template: dict[str, Any] = DEFAULT_PROMPT_TEMPLATE, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, max_sequence_length: int = 256, image_embed_interleave: int = 2, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if prompt_embeds is None: prompt_embeds, prompt_attention_mask = self._get_llama_prompt_embeds( image, @@ -552,10 +552,10 @@ def prepare_latents( height: int = 720, width: int = 1280, num_frames: int = 129, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, image_condition_type: str = "latent_concat", ) -> torch.Tensor: if isinstance(generator, list) and len(generator) != batch_size: @@ -671,52 +671,50 @@ def interrupt(self): def __call__( self, image: PIL.Image.Image, - prompt: Union[str, List[str]] = None, - prompt_2: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, - negative_prompt_2: Union[str, List[str]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] = None, + negative_prompt: str | list[str] = None, + negative_prompt_2: str | list[str] = None, height: int = 720, width: int = 1280, num_frames: int = 129, num_inference_steps: int = 50, - sigmas: List[float] = None, + sigmas: list[float] = None, true_cfg_scale: float = 1.0, guidance_scale: float = 1.0, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], - prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE, + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + prompt_template: dict[str, Any] = DEFAULT_PROMPT_TEMPLATE, max_sequence_length: int = 256, - image_embed_interleave: Optional[int] = None, + image_embed_interleave: int | None = None, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is not greater than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders. height (`int`, defaults to `720`): @@ -728,7 +726,7 @@ def __call__( num_inference_steps (`int`, defaults to `50`): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -744,7 +742,7 @@ def __call__( conditional latent is not applied. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -781,7 +779,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/hunyuan_video/pipeline_output.py b/src/diffusers/pipelines/hunyuan_video/pipeline_output.py index fae0370a53b7..4cf966262770 100644 --- a/src/diffusers/pipelines/hunyuan_video/pipeline_output.py +++ b/src/diffusers/pipelines/hunyuan_video/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -14,8 +13,8 @@ class HunyuanVideoPipelineOutput(BaseOutput): Output class for HunyuanVideo pipelines. Args: - frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): - List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): + list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape `(batch_size, num_frames, channels, height, width)`. """ @@ -29,11 +28,11 @@ class HunyuanVideoFramepackPipelineOutput(BaseOutput): Output class for HunyuanVideo pipelines. Args: - frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): - List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): + list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape `(batch_size, num_frames, channels, height, width)`. Or, a list of torch tensors where each tensor corresponds to a latent that decodes to multiple frames. """ - frames: Union[torch.Tensor, np.ndarray, List[List[PIL.Image.Image]], List[torch.Tensor]] + frames: torch.Tensor | np.ndarray | list[list[PIL.Image.Image]] | list[torch.Tensor] diff --git a/src/diffusers/pipelines/hunyuan_video1_5/pipeline_hunyuan_video1_5.py b/src/diffusers/pipelines/hunyuan_video1_5/pipeline_hunyuan_video1_5.py index 00a703939004..a0adff493ac0 100644 --- a/src/diffusers/pipelines/hunyuan_video1_5/pipeline_hunyuan_video1_5.py +++ b/src/diffusers/pipelines/hunyuan_video1_5/pipeline_hunyuan_video1_5.py @@ -14,7 +14,7 @@ import inspect import re -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import numpy as np import torch @@ -61,16 +61,16 @@ """ -def format_text_input(prompt: List[str], system_message: str) -> List[Dict[str, Any]]: +def format_text_input(prompt: list[str], system_message: str) -> list[dict[str, Any]]: """ Apply text to template. Args: - prompt (List[str]): Input text. + prompt (list[str]): Input text. system_message (str): System message. Returns: - List[Dict[str, Any]]: List of chat conversation. + list[dict[str, Any]]: List of chat conversation. """ template = [ @@ -80,7 +80,7 @@ def format_text_input(prompt: List[str], system_message: str) -> List[Dict[str, return template -def extract_glyph_texts(prompt: str) -> List[str]: +def extract_glyph_texts(prompt: str) -> list[str]: """ Extract glyph texts from prompt using regex pattern. @@ -106,10 +106,10 @@ def extract_glyph_texts(prompt: str) -> List[str]: # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -124,15 +124,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -241,7 +241,7 @@ def __init__( def _get_mllm_prompt_embeds( text_encoder: Qwen2_5_VLTextModel, tokenizer: Qwen2Tokenizer, - prompt: Union[str, List[str]], + prompt: str | list[str], device: torch.device, tokenizer_max_length: int = 1000, num_hidden_layers_to_skip: int = 2, @@ -254,7 +254,7 @@ def _get_mllm_prompt_embeds( 5. camera angles, movements, and transitions used in the video.", # fmt: on crop_start: int = 108, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: prompt = [prompt] if isinstance(prompt, str) else prompt prompt = format_text_input(prompt, system_message) @@ -289,7 +289,7 @@ def _get_mllm_prompt_embeds( def _get_byt5_prompt_embeds( tokenizer: ByT5Tokenizer, text_encoder: T5EncoderModel, - prompt: Union[str, List[str]], + prompt: str | list[str], device: torch.device, tokenizer_max_length: int = 256, ): @@ -333,20 +333,20 @@ def _get_byt5_prompt_embeds( def encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt: str | list[str], + device: torch.device | None = None, + dtype: torch.dtype | None = None, batch_size: int = 1, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, - prompt_embeds_2: Optional[torch.Tensor] = None, - prompt_embeds_mask_2: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + prompt_embeds_2: torch.Tensor | None = None, + prompt_embeds_mask_2: torch.Tensor | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -481,10 +481,10 @@ def prepare_latents( height: int = 720, width: int = 1280, num_frames: int = 129, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: if latents is not None: return latents.to(device=device, dtype=dtype) @@ -505,7 +505,7 @@ def prepare_latents( latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) return latents - def prepare_cond_latents_and_mask(self, latents, dtype: Optional[torch.dtype], device: Optional[torch.device]): + def prepare_cond_latents_and_mask(self, latents, dtype: torch.dtype | None, device: torch.device | None): """ Prepare conditional latents and mask for t2v generation. @@ -543,36 +543,36 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_frames: int = 121, num_inference_steps: int = 50, - sigmas: List[float] = None, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_mask: Optional[torch.Tensor] = None, - prompt_embeds_2: Optional[torch.Tensor] = None, - prompt_embeds_mask_2: Optional[torch.Tensor] = None, - negative_prompt_embeds_2: Optional[torch.Tensor] = None, - negative_prompt_embeds_mask_2: Optional[torch.Tensor] = None, - output_type: Optional[str] = "np", + sigmas: list[float] = None, + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_mask: torch.Tensor | None = None, + prompt_embeds_2: torch.Tensor | None = None, + prompt_embeds_mask_2: torch.Tensor | None = None, + negative_prompt_embeds_2: torch.Tensor | None = None, + negative_prompt_embeds_mask_2: torch.Tensor | None = None, + output_type: str | None = "np", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, + attention_kwargs: dict[str, Any] | None = None, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds` instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. height (`int`, *optional*): @@ -584,13 +584,13 @@ def __call__( num_inference_steps (`int`, defaults to `50`): The number of denoising steps. More denoising steps usually lead to a higher quality video at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of videos to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/hunyuan_video1_5/pipeline_hunyuan_video1_5_image2video.py b/src/diffusers/pipelines/hunyuan_video1_5/pipeline_hunyuan_video1_5_image2video.py index 8c555eabba11..791dec073524 100644 --- a/src/diffusers/pipelines/hunyuan_video1_5/pipeline_hunyuan_video1_5_image2video.py +++ b/src/diffusers/pipelines/hunyuan_video1_5/pipeline_hunyuan_video1_5_image2video.py @@ -14,7 +14,7 @@ import inspect import re -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import numpy as np import PIL @@ -73,16 +73,16 @@ # Copied from diffusers.pipelines.hunyuan_video1_5.pipeline_hunyuan_video1_5.format_text_input -def format_text_input(prompt: List[str], system_message: str) -> List[Dict[str, Any]]: +def format_text_input(prompt: list[str], system_message: str) -> list[dict[str, Any]]: """ Apply text to template. Args: - prompt (List[str]): Input text. + prompt (list[str]): Input text. system_message (str): System message. Returns: - List[Dict[str, Any]]: List of chat conversation. + list[dict[str, Any]]: List of chat conversation. """ template = [ @@ -93,7 +93,7 @@ def format_text_input(prompt: List[str], system_message: str) -> List[Dict[str, # Copied from diffusers.pipelines.hunyuan_video1_5.pipeline_hunyuan_video1_5.extract_glyph_texts -def extract_glyph_texts(prompt: str) -> List[str]: +def extract_glyph_texts(prompt: str) -> list[str]: """ Extract glyph texts from prompt using regex pattern. @@ -118,7 +118,7 @@ def extract_glyph_texts(prompt: str) -> List[str]: # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -133,10 +133,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -151,15 +151,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -280,7 +280,7 @@ def __init__( def _get_mllm_prompt_embeds( text_encoder: Qwen2_5_VLTextModel, tokenizer: Qwen2Tokenizer, - prompt: Union[str, List[str]], + prompt: str | list[str], device: torch.device, tokenizer_max_length: int = 1000, num_hidden_layers_to_skip: int = 2, @@ -293,7 +293,7 @@ def _get_mllm_prompt_embeds( 5. camera angles, movements, and transitions used in the video.", # fmt: on crop_start: int = 108, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: prompt = [prompt] if isinstance(prompt, str) else prompt prompt = format_text_input(prompt, system_message) @@ -329,7 +329,7 @@ def _get_mllm_prompt_embeds( def _get_byt5_prompt_embeds( tokenizer: ByT5Tokenizer, text_encoder: T5EncoderModel, - prompt: Union[str, List[str]], + prompt: str | list[str], device: torch.device, tokenizer_max_length: int = 256, ): @@ -421,20 +421,20 @@ def encode_image( # Copied from diffusers.pipelines.hunyuan_video1_5.pipeline_hunyuan_video1_5.HunyuanVideo15Pipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt: str | list[str], + device: torch.device | None = None, + dtype: torch.dtype | None = None, batch_size: int = 1, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, - prompt_embeds_2: Optional[torch.Tensor] = None, - prompt_embeds_mask_2: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + prompt_embeds_2: torch.Tensor | None = None, + prompt_embeds_mask_2: torch.Tensor | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -567,10 +567,10 @@ def prepare_latents( height: int = 720, width: int = 1280, num_frames: int = 129, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: if latents is not None: return latents.to(device=device, dtype=dtype) @@ -652,25 +652,25 @@ def interrupt(self): def __call__( self, image: PIL.Image.Image, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, num_frames: int = 121, num_inference_steps: int = 50, - sigmas: List[float] = None, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_mask: Optional[torch.Tensor] = None, - prompt_embeds_2: Optional[torch.Tensor] = None, - prompt_embeds_mask_2: Optional[torch.Tensor] = None, - negative_prompt_embeds_2: Optional[torch.Tensor] = None, - negative_prompt_embeds_mask_2: Optional[torch.Tensor] = None, - output_type: Optional[str] = "np", + sigmas: list[float] = None, + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_mask: torch.Tensor | None = None, + prompt_embeds_2: torch.Tensor | None = None, + prompt_embeds_mask_2: torch.Tensor | None = None, + negative_prompt_embeds_2: torch.Tensor | None = None, + negative_prompt_embeds_mask_2: torch.Tensor | None = None, + output_type: str | None = "np", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, + attention_kwargs: dict[str, Any] | None = None, ): r""" The call function to the pipeline for generation. @@ -678,10 +678,10 @@ def __call__( Args: image (`PIL.Image.Image`): The input image to condition video generation on. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds` instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the video generation. If not defined, one has to pass `negative_prompt_embeds` instead. num_frames (`int`, defaults to `121`): @@ -689,13 +689,13 @@ def __call__( num_inference_steps (`int`, defaults to `50`): The number of denoising steps. More denoising steps usually lead to a higher quality video at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of videos to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/hunyuan_video1_5/pipeline_output.py b/src/diffusers/pipelines/hunyuan_video1_5/pipeline_output.py index 441164db5a09..ab8767299e55 100644 --- a/src/diffusers/pipelines/hunyuan_video1_5/pipeline_output.py +++ b/src/diffusers/pipelines/hunyuan_video1_5/pipeline_output.py @@ -11,7 +11,7 @@ class HunyuanVideo15PipelineOutput(BaseOutput): Output class for HunyuanVideo1.5 pipelines. Args: - frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape `(batch_size, num_frames, channels, height, width)`. diff --git a/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py b/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py index 052c7b473915..b908dd5dfe83 100644 --- a/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py +++ b/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Callable, Dict, List, Optional, Tuple, Union +from typing import Callable import numpy as np import torch @@ -160,10 +160,10 @@ class HunyuanDiTPipeline(DiffusionPipeline): vae ([`AutoencoderKL`]): Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. We use `sdxl-vae-fp16-fix`. - text_encoder (Optional[`~transformers.BertModel`, `~transformers.CLIPTextModel`]): + text_encoder (`~transformers.BertModel`, `~transformers.CLIPTextModel` | None): Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). HunyuanDiT uses a fine-tuned [bilingual CLIP]. - tokenizer (Optional[`~transformers.BertTokenizer`, `~transformers.CLIPTokenizer`]): + tokenizer (`~transformers.BertTokenizer`, `~transformers.CLIPTokenizer` | None): A `BertTokenizer` or `CLIPTokenizer` to tokenize text. transformer ([`HunyuanDiT2DModel`]): The HunyuanDiT model designed by Tencent Hunyuan. @@ -203,8 +203,8 @@ def __init__( safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, requires_safety_checker: bool = True, - text_encoder_2: Optional[T5EncoderModel] = None, - tokenizer_2: Optional[T5Tokenizer] = None, + text_encoder_2: T5EncoderModel | None = None, + tokenizer_2: T5Tokenizer | None = None, ): super().__init__() @@ -252,19 +252,19 @@ def encode_prompt( dtype: torch.dtype = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - max_sequence_length: Optional[int] = None, + negative_prompt: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + max_sequence_length: int | None = None, text_encoder_index: int = 0, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -274,7 +274,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -363,7 +363,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -569,41 +569,39 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: Optional[int] = 50, - guidance_scale: Optional[float] = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - eta: Optional[float] = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_2: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_2: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - prompt_attention_mask_2: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask_2: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, + num_inference_steps: int | None = 50, + guidance_scale: float | None = 5.0, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + eta: float | None = 0.0, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_2: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_2: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + prompt_attention_mask_2: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask_2: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], guidance_rescale: float = 0.0, - original_size: Optional[Tuple[int, int]] = (1024, 1024), - target_size: Optional[Tuple[int, int]] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), + original_size: tuple[int, int] | None = (1024, 1024), + target_size: tuple[int, int] | None = None, + crops_coords_top_left: tuple[int, int] = (0, 0), use_resolution_binning: bool = True, ): r""" The call function to the pipeline for generation with HunyuanDiT. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`): The height in pixels of the generated image. @@ -615,7 +613,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -623,7 +621,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. prompt_embeds (`torch.Tensor`, *optional*): @@ -651,19 +649,19 @@ def __call__( return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a plain tuple. - callback_on_step_end (`Callable[[int, int, Dict], None]`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): + callback_on_step_end (`Callable[[int, int], None]`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): A callback function or a list of callback functions to be called at the end of each denoising step. - callback_on_step_end_tensor_inputs (`List[str]`, *optional*): + callback_on_step_end_tensor_inputs (`list[str]`, *optional*): A list of tensor inputs that should be passed to the callback function. If not defined, all tensor inputs will be passed. guidance_rescale (`float`, *optional*, defaults to 0.0): Rescale the noise_cfg according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). See Section 3.4 - original_size (`Tuple[int, int]`, *optional*, defaults to `(1024, 1024)`): + original_size (`tuple[int, int]`, *optional*, defaults to `(1024, 1024)`): The original size of the image. Used to calculate the time ids. - target_size (`Tuple[int, int]`, *optional*): + target_size (`tuple[int, int]`, *optional*): The target size of the image. Used to calculate the time ids. - crops_coords_top_left (`Tuple[int, int]`, *optional*, defaults to `(0, 0)`): + crops_coords_top_left (`tuple[int, int]`, *optional*, defaults to `(0, 0)`): The top left coordinates of the crop. Used to calculate the time ids. use_resolution_binning (`bool`, *optional*, defaults to `True`): Whether to use resolution binning or not. If `True`, the input resolution will be mapped to the closest diff --git a/src/diffusers/pipelines/i2vgen_xl/pipeline_i2vgen_xl.py b/src/diffusers/pipelines/i2vgen_xl/pipeline_i2vgen_xl.py index c6cc724a71f0..731ac27a0ff5 100644 --- a/src/diffusers/pipelines/i2vgen_xl/pipeline_i2vgen_xl.py +++ b/src/diffusers/pipelines/i2vgen_xl/pipeline_i2vgen_xl.py @@ -14,7 +14,7 @@ import inspect from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import numpy as np import PIL @@ -86,14 +86,14 @@ class I2VGenXLPipelineOutput(BaseOutput): Output class for image-to-video pipeline. Args: - frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): - List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): + list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape `(batch_size, num_frames, channels, height, width)` """ - frames: Union[torch.Tensor, np.ndarray, List[List[PIL.Image.Image]]] + frames: torch.Tensor | np.ndarray | list[list[PIL.Image.Image]] class I2VGenXLPipeline( @@ -165,15 +165,15 @@ def encode_prompt( device, num_videos_per_prompt, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -181,7 +181,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -263,7 +263,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if self.do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -445,7 +445,7 @@ def check_inputs( and not isinstance(image, list) ): raise ValueError( - "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is" + "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]` but is" f" {type(image)}" ) @@ -511,34 +511,34 @@ def prepare_latents( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, image: PipelineImageInput = None, - height: Optional[int] = 704, - width: Optional[int] = 1280, - target_fps: Optional[int] = 16, + height: int | None = 704, + width: int | None = 1280, + target_fps: int | None = 16, num_frames: int = 16, num_inference_steps: int = 50, guidance_scale: float = 9.0, - negative_prompt: Optional[Union[str, List[str]]] = None, + negative_prompt: str | list[str] | None = None, eta: float = 0.0, - num_videos_per_prompt: Optional[int] = 1, - decode_chunk_size: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + num_videos_per_prompt: int | None = 1, + decode_chunk_size: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = 1, + cross_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = 1, ): r""" The call function to the pipeline for image-to-video generation with [`I2VGenXLPipeline`]. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.Tensor`): + image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.Tensor`): Image or images to guide image generation. If you provide a tensor, it needs to be compatible with [`CLIPImageProcessor`](https://huggingface.co/lambdalabs/sd-image-variations-diffusers/blob/main/feature_extractor/preprocessor_config.json). height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): @@ -555,7 +555,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). eta (`float`, *optional*): @@ -567,7 +567,7 @@ def __call__( The number of frames to decode at a time. The higher the chunk size, the higher the temporal consistency between frames, but also the higher the memory consumption. By default, the decoder will decode all frames at once for maximal quality. Reduce `decode_chunk_size` to reduce memory usage. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -745,7 +745,7 @@ def __call__( # https://github.com/ali-vilab/i2vgen-xl/blob/main/utils/transforms.py. -def _convert_pt_to_pil(image: Union[torch.Tensor, List[torch.Tensor]]): +def _convert_pt_to_pil(image: torch.Tensor | list[torch.Tensor]): if isinstance(image, list) and isinstance(image[0], torch.Tensor): image = torch.cat(image, 0) @@ -761,7 +761,7 @@ def _convert_pt_to_pil(image: Union[torch.Tensor, List[torch.Tensor]]): def _resize_bilinear( - image: Union[torch.Tensor, List[torch.Tensor], PIL.Image.Image, List[PIL.Image.Image]], resolution: Tuple[int, int] + image: torch.Tensor | list[torch.Tensor] | PIL.Image.Image | list[PIL.Image.Image], resolution: tuple[int, int] ): # First convert the images to PIL in case they are float tensors (only relevant for tests now). image = _convert_pt_to_pil(image) @@ -774,7 +774,7 @@ def _resize_bilinear( def _center_crop_wide( - image: Union[torch.Tensor, List[torch.Tensor], PIL.Image.Image, List[PIL.Image.Image]], resolution: Tuple[int, int] + image: torch.Tensor | list[torch.Tensor] | PIL.Image.Image | list[PIL.Image.Image], resolution: tuple[int, int] ): # First convert the images to PIL in case they are float tensors (only relevant for tests now). image = _convert_pt_to_pil(image) diff --git a/src/diffusers/pipelines/kandinsky/pipeline_kandinsky.py b/src/diffusers/pipelines/kandinsky/pipeline_kandinsky.py index 33529f5d0954..d635057f2b05 100644 --- a/src/diffusers/pipelines/kandinsky/pipeline_kandinsky.py +++ b/src/diffusers/pipelines/kandinsky/pipeline_kandinsky.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Callable, List, Optional, Union +from typing import Callable import torch from transformers import ( @@ -90,7 +90,7 @@ class KandinskyPipeline(DiffusionPipeline): Frozen text-encoder. tokenizer ([`XLMRobertaTokenizer`]): Tokenizer of class - scheduler (Union[`DDIMScheduler`,`DDPMScheduler`]): + scheduler (`DDIMScheduler` | `DDPMScheduler`): A scheduler to be used in combination with `unet` to generate image latents. unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the image embedding. @@ -105,7 +105,7 @@ def __init__( text_encoder: MultilingualCLIP, tokenizer: XLMRobertaTokenizer, unet: UNet2DConditionModel, - scheduler: Union[DDIMScheduler, DDPMScheduler], + scheduler: DDIMScheduler | DDPMScheduler, movq: VQModel, ): super().__init__() @@ -173,7 +173,7 @@ def _encode_prompt( text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0) if do_classifier_free_guidance: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -237,19 +237,19 @@ def _encode_prompt( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]], - image_embeds: Union[torch.Tensor, List[torch.Tensor]], - negative_image_embeds: Union[torch.Tensor, List[torch.Tensor]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + image_embeds: torch.Tensor | list[torch.Tensor], + negative_image_embeds: torch.Tensor | list[torch.Tensor], + negative_prompt: str | list[str] | None = None, height: int = 512, width: int = 512, num_inference_steps: int = 100, guidance_scale: float = 4.0, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, return_dict: bool = True, ): @@ -257,13 +257,13 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. - image_embeds (`torch.Tensor` or `List[torch.Tensor]`): + image_embeds (`torch.Tensor` or `list[torch.Tensor]`): The clip image embeddings for text prompt, that will be used to condition the image generation. - negative_image_embeds (`torch.Tensor` or `List[torch.Tensor]`): + negative_image_embeds (`torch.Tensor` or `list[torch.Tensor]`): The clip image embeddings for negative text prompt, will be used to condition the image generation. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). height (`int`, *optional*, defaults to 512): @@ -281,7 +281,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py b/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py index 7286bcbee17b..21d5c98efa90 100644 --- a/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py +++ b/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Callable, List, Optional, Union +from typing import Callable import PIL.Image import torch @@ -122,7 +122,7 @@ class KandinskyCombinedPipeline(DiffusionPipeline): Frozen text-encoder. tokenizer ([`XLMRobertaTokenizer`]): Tokenizer of class - scheduler (Union[`DDIMScheduler`,`DDPMScheduler`]): + scheduler (`DDIMScheduler` | `DDPMScheduler`): A scheduler to be used in combination with `unet` to generate image latents. unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the image embedding. @@ -150,7 +150,7 @@ def __init__( text_encoder: MultilingualCLIP, tokenizer: XLMRobertaTokenizer, unet: UNet2DConditionModel, - scheduler: Union[DDIMScheduler, DDPMScheduler], + scheduler: DDIMScheduler | DDPMScheduler, movq: VQModel, prior_prior: PriorTransformer, prior_image_encoder: CLIPVisionModelWithProjection, @@ -190,10 +190,10 @@ def __init__( movq=movq, ) - def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None): + def enable_xformers_memory_efficient_attention(self, attention_op: Callable | None = None): self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op) - def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None): + def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None): r""" Offloads all models (`unet`, `text_encoder`, `vae`, and `safety checker` state dicts) to CPU using 🤗 Accelerate, significantly reducing memory usage. Models are moved to a `torch.device('meta')` and loaded on a @@ -216,8 +216,8 @@ def set_progress_bar_config(self, **kwargs): @replace_example_docstring(TEXT2IMAGE_EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, num_inference_steps: int = 100, guidance_scale: float = 4.0, num_images_per_prompt: int = 1, @@ -225,10 +225,10 @@ def __call__( width: int = 512, prior_guidance_scale: float = 4.0, prior_num_inference_steps: int = 25, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, return_dict: bool = True, ): @@ -236,9 +236,9 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -265,7 +265,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -340,7 +340,7 @@ class KandinskyImg2ImgCombinedPipeline(DiffusionPipeline): Frozen text-encoder. tokenizer ([`XLMRobertaTokenizer`]): Tokenizer of class - scheduler (Union[`DDIMScheduler`,`DDPMScheduler`]): + scheduler (`DDIMScheduler` | `DDPMScheduler`): A scheduler to be used in combination with `unet` to generate image latents. unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the image embedding. @@ -368,7 +368,7 @@ def __init__( text_encoder: MultilingualCLIP, tokenizer: XLMRobertaTokenizer, unet: UNet2DConditionModel, - scheduler: Union[DDIMScheduler, DDPMScheduler], + scheduler: DDIMScheduler | DDPMScheduler, movq: VQModel, prior_prior: PriorTransformer, prior_image_encoder: CLIPVisionModelWithProjection, @@ -408,10 +408,10 @@ def __init__( movq=movq, ) - def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None): + def enable_xformers_memory_efficient_attention(self, attention_op: Callable | None = None): self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op) - def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None): + def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None): r""" Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet, text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a @@ -435,9 +435,9 @@ def set_progress_bar_config(self, **kwargs): @replace_example_docstring(IMAGE2IMAGE_EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]], - image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + image: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image], + negative_prompt: str | list[str] | None = None, num_inference_steps: int = 100, guidance_scale: float = 4.0, num_images_per_prompt: int = 1, @@ -446,10 +446,10 @@ def __call__( width: int = 512, prior_guidance_scale: float = 4.0, prior_num_inference_steps: int = 25, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, return_dict: bool = True, ): @@ -457,13 +457,13 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, or tensor representing an image batch, that will be used as the starting point for the process. Can also accept image latents as `image`, if passing latents directly, it will not be encoded again. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -496,7 +496,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -581,7 +581,7 @@ class KandinskyInpaintCombinedPipeline(DiffusionPipeline): Frozen text-encoder. tokenizer ([`XLMRobertaTokenizer`]): Tokenizer of class - scheduler (Union[`DDIMScheduler`,`DDPMScheduler`]): + scheduler (`DDIMScheduler` | `DDPMScheduler`): A scheduler to be used in combination with `unet` to generate image latents. unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the image embedding. @@ -609,7 +609,7 @@ def __init__( text_encoder: MultilingualCLIP, tokenizer: XLMRobertaTokenizer, unet: UNet2DConditionModel, - scheduler: Union[DDIMScheduler, DDPMScheduler], + scheduler: DDIMScheduler | DDPMScheduler, movq: VQModel, prior_prior: PriorTransformer, prior_image_encoder: CLIPVisionModelWithProjection, @@ -649,10 +649,10 @@ def __init__( movq=movq, ) - def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None): + def enable_xformers_memory_efficient_attention(self, attention_op: Callable | None = None): self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op) - def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None): + def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None): r""" Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet, text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a @@ -676,10 +676,10 @@ def set_progress_bar_config(self, **kwargs): @replace_example_docstring(INPAINT_EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]], - image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]], - mask_image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + image: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image], + mask_image: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image], + negative_prompt: str | list[str] | None = None, num_inference_steps: int = 100, guidance_scale: float = 4.0, num_images_per_prompt: int = 1, @@ -687,10 +687,10 @@ def __call__( width: int = 512, prior_guidance_scale: float = 4.0, prior_num_inference_steps: int = 25, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, return_dict: bool = True, ): @@ -698,9 +698,9 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, or tensor representing an image batch, that will be used as the starting point for the process. Can also accept image latents as `image`, if passing latents directly, it will not be encoded again. @@ -709,7 +709,7 @@ def __call__( black pixels will be preserved. If `mask_image` is a PIL image, it will be converted to a single channel (luminance) before use. If it's a tensor, it should contain one color channel (L) instead of 3, so the expected shape would be `(B, H, W, 1)`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -736,7 +736,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_img2img.py b/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_img2img.py index f5e41d499dc3..f33317447b49 100644 --- a/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_img2img.py +++ b/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_img2img.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Callable, List, Optional, Union +from typing import Callable import PIL.Image import torch @@ -207,7 +207,7 @@ def _encode_prompt( text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0) if do_classifier_free_guidance: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -298,20 +298,20 @@ def add_noise( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]], - image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]], + prompt: str | list[str], + image: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image], image_embeds: torch.Tensor, negative_image_embeds: torch.Tensor, - negative_prompt: Optional[Union[str, List[str]]] = None, + negative_prompt: str | list[str] | None = None, height: int = 512, width: int = 512, num_inference_steps: int = 100, strength: float = 0.3, guidance_scale: float = 7.0, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - output_type: Optional[str] = "pil", - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + output_type: str | None = "pil", + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, return_dict: bool = True, ): @@ -319,16 +319,16 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. image (`torch.Tensor`, `PIL.Image.Image`): `Image`, or tensor representing an image batch, that will be used as the starting point for the process. - image_embeds (`torch.Tensor` or `List[torch.Tensor]`): + image_embeds (`torch.Tensor` or `list[torch.Tensor]`): The clip image embeddings for text prompt, that will be used to condition the image generation. - negative_image_embeds (`torch.Tensor` or `List[torch.Tensor]`): + negative_image_embeds (`torch.Tensor` or `list[torch.Tensor]`): The clip image embeddings for negative text prompt, will be used to condition the image generation. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). height (`int`, *optional*, defaults to 512): @@ -352,7 +352,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. output_type (`str`, *optional*, defaults to `"pil"`): diff --git a/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py b/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py index 731fce499859..4dba85446db5 100644 --- a/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py +++ b/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py @@ -13,7 +13,7 @@ # limitations under the License. from copy import deepcopy -from typing import Callable, List, Optional, Union +from typing import Callable import numpy as np import PIL.Image @@ -134,7 +134,7 @@ def prepare_mask_and_masked_image(image, mask, height, width): binarized (``mask > 0.5``) and cast to ``torch.float32`` too. Args: - image (Union[np.array, PIL.Image, torch.Tensor]): The image to inpaint. + image (np.array | PIL.Image | torch.Tensor): The image to inpaint. It can be a ``PIL.Image``, or a ``height x width x 3`` ``np.array`` or a ``channels x height x width`` ``torch.Tensor`` or a ``batch x channels x height x width`` ``torch.Tensor``. mask (_type_): The mask to apply to the image, i.e. regions to inpaint. @@ -338,7 +338,7 @@ def _encode_prompt( text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0) if do_classifier_free_guidance: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -402,21 +402,21 @@ def _encode_prompt( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]], - image: Union[torch.Tensor, PIL.Image.Image], - mask_image: Union[torch.Tensor, PIL.Image.Image, np.ndarray], + prompt: str | list[str], + image: torch.Tensor | PIL.Image.Image, + mask_image: torch.Tensor | PIL.Image.Image | np.ndarray, image_embeds: torch.Tensor, negative_image_embeds: torch.Tensor, - negative_prompt: Optional[Union[str, List[str]]] = None, + negative_prompt: str | list[str] | None = None, height: int = 512, width: int = 512, num_inference_steps: int = 100, guidance_scale: float = 4.0, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, return_dict: bool = True, ): @@ -424,7 +424,7 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. image (`torch.Tensor`, `PIL.Image.Image` or `np.ndarray`): `Image`, or tensor representing an image batch, that will be used as the starting point for the @@ -437,11 +437,11 @@ def __call__( image or numpy array, mask should also be a either PIL image or numpy array. If it is a PIL image, it will be converted to a single channel (luminance) before use. If it is a nummpy array, the expected shape is `(H, W)`. - image_embeds (`torch.Tensor` or `List[torch.Tensor]`): + image_embeds (`torch.Tensor` or `list[torch.Tensor]`): The clip image embeddings for text prompt, that will be used to condition the image generation. - negative_image_embeds (`torch.Tensor` or `List[torch.Tensor]`): + negative_image_embeds (`torch.Tensor` or `list[torch.Tensor]`): The clip image embeddings for negative text prompt, will be used to condition the image generation. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). height (`int`, *optional*, defaults to 512): @@ -459,7 +459,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.py b/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.py index 10ea8005c90d..523fd010eb7f 100644 --- a/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.py +++ b/src/diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.py @@ -13,7 +13,6 @@ # limitations under the License. from dataclasses import dataclass -from typing import List, Optional, Union import numpy as np import PIL.Image @@ -126,12 +125,12 @@ class KandinskyPriorPipelineOutput(BaseOutput): Args: image_embeds (`torch.Tensor`) clip image embeddings for text prompt - negative_image_embeds (`List[PIL.Image.Image]` or `np.ndarray`) + negative_image_embeds (`list[PIL.Image.Image]` or `np.ndarray`) clip image embeddings for unconditional tokens """ - image_embeds: Union[torch.Tensor, np.ndarray] - negative_image_embeds: Union[torch.Tensor, np.ndarray] + image_embeds: torch.Tensor | np.ndarray + negative_image_embeds: torch.Tensor | np.ndarray class KandinskyPriorPipeline(DiffusionPipeline): @@ -182,13 +181,13 @@ def __init__( @replace_example_docstring(EXAMPLE_INTERPOLATE_DOC_STRING) def interpolate( self, - images_and_prompts: List[Union[str, PIL.Image.Image, torch.Tensor]], - weights: List[float], + images_and_prompts: list[str | PIL.Image.Image | torch.Tensor], + weights: list[float], num_images_per_prompt: int = 1, num_inference_steps: int = 25, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - negative_prior_prompt: Optional[str] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + negative_prior_prompt: str | None = None, negative_prompt: str = "", guidance_scale: float = 4.0, device=None, @@ -197,16 +196,16 @@ def interpolate( Function invoked when using the prior pipeline for interpolation. Args: - images_and_prompts (`List[Union[str, PIL.Image.Image, torch.Tensor]]`): + images_and_prompts (`list[str | PIL.Image.Image | torch.Tensor]`): list of prompts and images to guide the image generation. - weights: (`List[float]`): + weights: (`list[float]`): list of weights for each condition in `images_and_prompts` num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. num_inference_steps (`int`, *optional*, defaults to 25): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -216,7 +215,7 @@ def interpolate( negative_prior_prompt (`str`, *optional*): The prompt not to guide the prior diffusion process. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). guidance_scale (`float`, *optional*, defaults to 4.0): @@ -346,7 +345,7 @@ def _encode_prompt( text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0) if do_classifier_free_guidance: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -407,23 +406,23 @@ def _encode_prompt( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, num_images_per_prompt: int = 1, num_inference_steps: int = 25, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, guidance_scale: float = 4.0, - output_type: Optional[str] = "pt", + output_type: str | None = "pt", return_dict: bool = True, ): """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -431,7 +430,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 25): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/kandinsky/text_encoder.py b/src/diffusers/pipelines/kandinsky/text_encoder.py index caa0029f00ca..58cc9ac4d3ed 100644 --- a/src/diffusers/pipelines/kandinsky/text_encoder.py +++ b/src/diffusers/pipelines/kandinsky/text_encoder.py @@ -20,6 +20,8 @@ def __init__(self, config, *args, **kwargs): self.LinearTransformation = torch.nn.Linear( in_features=config.transformerDimensions, out_features=config.numDims ) + if hasattr(self, "post_init"): + self.post_init() def forward(self, input_ids, attention_mask): embs = self.transformer(input_ids=input_ids, attention_mask=attention_mask)[0] diff --git a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py index 429253e99898..5129b3f548e8 100644 --- a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py +++ b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import torch @@ -77,7 +77,7 @@ class KandinskyV22Pipeline(DiffusionPipeline): library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.) Args: - scheduler (Union[`DDIMScheduler`,`DDPMScheduler`]): + scheduler (`DDIMScheduler` | `DDPMScheduler`): A scheduler to be used in combination with `unet` to generate image latents. unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the image embedding. @@ -131,28 +131,28 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - image_embeds: Union[torch.Tensor, List[torch.Tensor]], - negative_image_embeds: Union[torch.Tensor, List[torch.Tensor]], + image_embeds: torch.Tensor | list[torch.Tensor], + negative_image_embeds: torch.Tensor | list[torch.Tensor], height: int = 512, width: int = 512, num_inference_steps: int = 100, guidance_scale: float = 4.0, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): """ Function invoked when calling the pipeline for generation. Args: - image_embeds (`torch.Tensor` or `List[torch.Tensor]`): + image_embeds (`torch.Tensor` or `list[torch.Tensor]`): The clip image embeddings for text prompt, that will be used to condition the image generation. - negative_image_embeds (`torch.Tensor` or `List[torch.Tensor]`): + negative_image_embeds (`torch.Tensor` or `list[torch.Tensor]`): The clip image embeddings for negative text prompt, will be used to condition the image generation. height (`int`, *optional*, defaults to 512): The height in pixels of the generated image. @@ -169,7 +169,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -186,7 +186,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py index fc2083247bb0..26e163a70142 100644 --- a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py +++ b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import PIL.Image import torch @@ -112,7 +112,7 @@ class KandinskyV22CombinedPipeline(DiffusionPipeline): library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.) Args: - scheduler (Union[`DDIMScheduler`,`DDPMScheduler`]): + scheduler (`DDIMScheduler` | `DDPMScheduler`): A scheduler to be used in combination with `unet` to generate image latents. unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the image embedding. @@ -176,10 +176,10 @@ def __init__( movq=movq, ) - def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None): + def enable_xformers_memory_efficient_attention(self, attention_op: Callable | None = None): self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op) - def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None): + def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None): r""" Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet, text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a @@ -203,8 +203,8 @@ def set_progress_bar_config(self, **kwargs): @replace_example_docstring(TEXT2IMAGE_EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, num_inference_steps: int = 100, guidance_scale: float = 4.0, num_images_per_prompt: int = 1, @@ -212,24 +212,24 @@ def __call__( width: int = 512, prior_guidance_scale: float = 4.0, prior_num_inference_steps: int = 25, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, return_dict: bool = True, - prior_callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - prior_callback_on_step_end_tensor_inputs: List[str] = ["latents"], - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + prior_callback_on_step_end: Callable[[int, int], None] | None = None, + prior_callback_on_step_end_tensor_inputs: list[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], ): """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -256,7 +256,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -272,7 +272,7 @@ def __call__( A function that calls at the end of each denoising steps during the inference of the prior pipeline. The function is called with the following arguments: `prior_callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. - prior_callback_on_step_end_tensor_inputs (`List`, *optional*): + prior_callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `prior_callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your prior pipeline class. @@ -281,7 +281,7 @@ def __call__( The function is called with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -340,7 +340,7 @@ class KandinskyV22Img2ImgCombinedPipeline(DiffusionPipeline): library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.) Args: - scheduler (Union[`DDIMScheduler`,`DDPMScheduler`]): + scheduler (`DDIMScheduler` | `DDPMScheduler`): A scheduler to be used in combination with `unet` to generate image latents. unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the image embedding. @@ -404,10 +404,10 @@ def __init__( movq=movq, ) - def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None): + def enable_xformers_memory_efficient_attention(self, attention_op: Callable | None = None): self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op) - def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None): + def enable_model_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None): r""" Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward` @@ -417,7 +417,7 @@ def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[t self.prior_pipe.enable_model_cpu_offload(gpu_id=gpu_id, device=device) self.decoder_pipe.enable_model_cpu_offload(gpu_id=gpu_id, device=device) - def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None): + def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None): r""" Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet, text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a @@ -441,9 +441,9 @@ def set_progress_bar_config(self, **kwargs): @replace_example_docstring(IMAGE2IMAGE_EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]], - image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + image: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image], + negative_prompt: str | list[str] | None = None, num_inference_steps: int = 100, guidance_scale: float = 4.0, strength: float = 0.3, @@ -452,28 +452,28 @@ def __call__( width: int = 512, prior_guidance_scale: float = 4.0, prior_num_inference_steps: int = 25, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, return_dict: bool = True, - prior_callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - prior_callback_on_step_end_tensor_inputs: List[str] = ["latents"], - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + prior_callback_on_step_end: Callable[[int, int], None] | None = None, + prior_callback_on_step_end_tensor_inputs: list[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], ): """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, or tensor representing an image batch, that will be used as the starting point for the process. Can also accept image latents as `image`, if passing latents directly, it will not be encoded again. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -506,7 +506,7 @@ def __call__( prior_num_inference_steps (`int`, *optional*, defaults to 100): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -589,7 +589,7 @@ class KandinskyV22InpaintCombinedPipeline(DiffusionPipeline): library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.) Args: - scheduler (Union[`DDIMScheduler`,`DDPMScheduler`]): + scheduler (`DDIMScheduler` | `DDPMScheduler`): A scheduler to be used in combination with `unet` to generate image latents. unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the image embedding. @@ -653,10 +653,10 @@ def __init__( movq=movq, ) - def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None): + def enable_xformers_memory_efficient_attention(self, attention_op: Callable | None = None): self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op) - def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None): + def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None): r""" Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet, text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a @@ -680,10 +680,10 @@ def set_progress_bar_config(self, **kwargs): @replace_example_docstring(INPAINT_EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]], - image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]], - mask_image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + image: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image], + mask_image: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image], + negative_prompt: str | list[str] | None = None, num_inference_steps: int = 100, guidance_scale: float = 4.0, num_images_per_prompt: int = 1, @@ -691,23 +691,23 @@ def __call__( width: int = 512, prior_guidance_scale: float = 4.0, prior_num_inference_steps: int = 25, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - prior_callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - prior_callback_on_step_end_tensor_inputs: List[str] = ["latents"], - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + prior_callback_on_step_end: Callable[[int, int], None] | None = None, + prior_callback_on_step_end_tensor_inputs: list[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, or tensor representing an image batch, that will be used as the starting point for the process. Can also accept image latents as `image`, if passing latents directly, it will not be encoded again. @@ -716,7 +716,7 @@ def __call__( black pixels will be preserved. If `mask_image` is a PIL image, it will be converted to a single channel (luminance) before use. If it's a tensor, it should contain one color channel (L) instead of 3, so the expected shape would be `(B, H, W, 1)`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -743,7 +743,7 @@ def __call__( prior_num_inference_steps (`int`, *optional*, defaults to 100): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -759,7 +759,7 @@ def __call__( A function that calls at the end of each denoising steps during the inference. The function is called with the following arguments: `prior_callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. - prior_callback_on_step_end_tensor_inputs (`List`, *optional*): + prior_callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `prior_callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -768,7 +768,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py index c5faae82796b..01001e0c9eba 100644 --- a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py +++ b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Callable, List, Optional, Union +from typing import Callable import torch @@ -160,18 +160,18 @@ def prepare_latents(self, shape, dtype, device, generator, latents, scheduler): @torch.no_grad() def __call__( self, - image_embeds: Union[torch.Tensor, List[torch.Tensor]], - negative_image_embeds: Union[torch.Tensor, List[torch.Tensor]], + image_embeds: torch.Tensor | list[torch.Tensor], + negative_image_embeds: torch.Tensor | list[torch.Tensor], hint: torch.Tensor, height: int = 512, width: int = 512, num_inference_steps: int = 100, guidance_scale: float = 4.0, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, return_dict: bool = True, ): @@ -179,15 +179,15 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. hint (`torch.Tensor`): The controlnet condition. - image_embeds (`torch.Tensor` or `List[torch.Tensor]`): + image_embeds (`torch.Tensor` or `list[torch.Tensor]`): The clip image embeddings for text prompt, that will be used to condition the image generation. - negative_image_embeds (`torch.Tensor` or `List[torch.Tensor]`): + negative_image_embeds (`torch.Tensor` or `list[torch.Tensor]`): The clip image embeddings for negative text prompt, will be used to condition the image generation. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). height (`int`, *optional*, defaults to 512): @@ -205,7 +205,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet_img2img.py b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet_img2img.py index 54154c6ec1f2..891235558330 100644 --- a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet_img2img.py +++ b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet_img2img.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Callable, List, Optional, Union +from typing import Callable import PIL.Image import torch @@ -200,9 +200,9 @@ def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dt @torch.no_grad() def __call__( self, - image_embeds: Union[torch.Tensor, List[torch.Tensor]], - image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]], - negative_image_embeds: Union[torch.Tensor, List[torch.Tensor]], + image_embeds: torch.Tensor | list[torch.Tensor], + image: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image], + negative_image_embeds: torch.Tensor | list[torch.Tensor], hint: torch.Tensor, height: int = 512, width: int = 512, @@ -210,9 +210,9 @@ def __call__( guidance_scale: float = 4.0, strength: float = 0.3, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - output_type: Optional[str] = "pil", - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + output_type: str | None = "pil", + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, return_dict: bool = True, ): @@ -220,9 +220,9 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - image_embeds (`torch.Tensor` or `List[torch.Tensor]`): + image_embeds (`torch.Tensor` or `list[torch.Tensor]`): The clip image embeddings for text prompt, that will be used to condition the image generation. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, or tensor representing an image batch, that will be used as the starting point for the process. Can also accept image latents as `image`, if passing latents directly, it will not be encoded again. @@ -234,7 +234,7 @@ def __call__( `num_inference_steps`. A value of 1, therefore, essentially ignores `image`. hint (`torch.Tensor`): The controlnet condition. - negative_image_embeds (`torch.Tensor` or `List[torch.Tensor]`): + negative_image_embeds (`torch.Tensor` or `list[torch.Tensor]`): The clip image embeddings for negative text prompt, will be used to condition the image generation. height (`int`, *optional*, defaults to 512): The height in pixels of the generated image. @@ -251,7 +251,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. output_type (`str`, *optional*, defaults to `"pil"`): diff --git a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py index 3b2509098fd1..a43c085a8921 100644 --- a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py +++ b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import PIL.Image import torch @@ -183,29 +183,29 @@ def num_timesteps(self): @torch.no_grad() def __call__( self, - image_embeds: Union[torch.Tensor, List[torch.Tensor]], - image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]], - negative_image_embeds: Union[torch.Tensor, List[torch.Tensor]], + image_embeds: torch.Tensor | list[torch.Tensor], + image: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image], + negative_image_embeds: torch.Tensor | list[torch.Tensor], height: int = 512, width: int = 512, num_inference_steps: int = 100, guidance_scale: float = 4.0, strength: float = 0.3, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): """ Function invoked when calling the pipeline for generation. Args: - image_embeds (`torch.Tensor` or `List[torch.Tensor]`): + image_embeds (`torch.Tensor` or `list[torch.Tensor]`): The clip image embeddings for text prompt, that will be used to condition the image generation. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, or tensor representing an image batch, that will be used as the starting point for the process. Can also accept image latents as `image`, if passing latents directly, it will not be encoded again. @@ -215,7 +215,7 @@ def __call__( denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will be maximum and the denoising process will run for the full number of iterations specified in `num_inference_steps`. A value of 1, therefore, essentially ignores `image`. - negative_image_embeds (`torch.Tensor` or `List[torch.Tensor]`): + negative_image_embeds (`torch.Tensor` or `list[torch.Tensor]`): The clip image embeddings for negative text prompt, will be used to condition the image generation. height (`int`, *optional*, defaults to 512): The height in pixels of the generated image. @@ -232,7 +232,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. output_type (`str`, *optional*, defaults to `"pil"`): @@ -245,7 +245,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py index a61673293e1f..31bd88103a06 100644 --- a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py +++ b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py @@ -13,7 +13,7 @@ # limitations under the License. from copy import deepcopy -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import numpy as np import PIL.Image @@ -132,7 +132,7 @@ def prepare_mask_and_masked_image(image, mask, height, width): binarized (``mask > 0.5``) and cast to ``torch.float32`` too. Args: - image (Union[np.array, PIL.Image, torch.Tensor]): The image to inpaint. + image (np.array | PIL.Image | torch.Tensor): The image to inpaint. It can be a ``PIL.Image``, or a ``height x width x 3`` ``np.array`` or a ``channels x height x width`` ``torch.Tensor`` or a ``batch x channels x height x width`` ``torch.Tensor``. mask (_type_): The mask to apply to the image, i.e. regions to inpaint. @@ -302,28 +302,28 @@ def num_timesteps(self): @torch.no_grad() def __call__( self, - image_embeds: Union[torch.Tensor, List[torch.Tensor]], - image: Union[torch.Tensor, PIL.Image.Image], - mask_image: Union[torch.Tensor, PIL.Image.Image, np.ndarray], - negative_image_embeds: Union[torch.Tensor, List[torch.Tensor]], + image_embeds: torch.Tensor | list[torch.Tensor], + image: torch.Tensor | PIL.Image.Image, + mask_image: torch.Tensor | PIL.Image.Image | np.ndarray, + negative_image_embeds: torch.Tensor | list[torch.Tensor], height: int = 512, width: int = 512, num_inference_steps: int = 100, guidance_scale: float = 4.0, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): """ Function invoked when calling the pipeline for generation. Args: - image_embeds (`torch.Tensor` or `List[torch.Tensor]`): + image_embeds (`torch.Tensor` or `list[torch.Tensor]`): The clip image embeddings for text prompt, that will be used to condition the image generation. image (`PIL.Image.Image`): `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will @@ -333,7 +333,7 @@ def __call__( black pixels will be preserved. If `mask_image` is a PIL image, it will be converted to a single channel (luminance) before use. If it's a tensor, it should contain one color channel (L) instead of 3, so the expected shape would be `(B, H, W, 1)`. - negative_image_embeds (`torch.Tensor` or `List[torch.Tensor]`): + negative_image_embeds (`torch.Tensor` or `list[torch.Tensor]`): The clip image embeddings for negative text prompt, will be used to condition the image generation. height (`int`, *optional*, defaults to 512): The height in pixels of the generated image. @@ -350,7 +350,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -367,7 +367,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py index bc67847831a5..41f4474c3906 100644 --- a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py +++ b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py @@ -1,4 +1,4 @@ -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import PIL.Image import torch @@ -137,13 +137,13 @@ def __init__( @replace_example_docstring(EXAMPLE_INTERPOLATE_DOC_STRING) def interpolate( self, - images_and_prompts: List[Union[str, PIL.Image.Image, torch.Tensor]], - weights: List[float], + images_and_prompts: list[str | PIL.Image.Image | torch.Tensor], + weights: list[float], num_images_per_prompt: int = 1, num_inference_steps: int = 25, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - negative_prior_prompt: Optional[str] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + negative_prior_prompt: str | None = None, negative_prompt: str = "", guidance_scale: float = 4.0, device=None, @@ -152,16 +152,16 @@ def interpolate( Function invoked when using the prior pipeline for interpolation. Args: - images_and_prompts (`List[Union[str, PIL.Image.Image, torch.Tensor]]`): + images_and_prompts (`list[str | PIL.Image.Image | torch.Tensor]`): list of prompts and images to guide the image generation. - weights: (`List[float]`): + weights: (`list[float]`): list of weights for each condition in `images_and_prompts` num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. num_inference_steps (`int`, *optional*, defaults to 100): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -171,7 +171,7 @@ def interpolate( negative_prior_prompt (`str`, *optional*): The prompt not to guide the prior diffusion process. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). guidance_scale (`float`, *optional*, defaults to 4.0): @@ -303,7 +303,7 @@ def _encode_prompt( text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0) if do_classifier_free_guidance: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -376,25 +376,25 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, num_images_per_prompt: int = 1, num_inference_steps: int = 25, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, guidance_scale: float = 4.0, - output_type: Optional[str] = "pt", # pt only + output_type: str | None = "pt", # pt only return_dict: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], ): """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -402,7 +402,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 100): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -425,7 +425,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py index b586d166118b..adbc3a5badc5 100644 --- a/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py +++ b/src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py @@ -1,5 +1,3 @@ -from typing import List, Optional, Union - import PIL.Image import torch from transformers import CLIPImageProcessor, CLIPTextModelWithProjection, CLIPTokenizer, CLIPVisionModelWithProjection @@ -161,13 +159,13 @@ def get_timesteps(self, num_inference_steps, strength, device): @replace_example_docstring(EXAMPLE_INTERPOLATE_DOC_STRING) def interpolate( self, - images_and_prompts: List[Union[str, PIL.Image.Image, torch.Tensor]], - weights: List[float], + images_and_prompts: list[str | PIL.Image.Image | torch.Tensor], + weights: list[float], num_images_per_prompt: int = 1, num_inference_steps: int = 25, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - negative_prior_prompt: Optional[str] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + negative_prior_prompt: str | None = None, negative_prompt: str = "", guidance_scale: float = 4.0, device=None, @@ -176,16 +174,16 @@ def interpolate( Function invoked when using the prior pipeline for interpolation. Args: - images_and_prompts (`List[Union[str, PIL.Image.Image, torch.Tensor]]`): + images_and_prompts (`list[str | PIL.Image.Image | torch.Tensor]`): list of prompts and images to guide the image generation. - weights: (`List[float]`): + weights: (`list[float]`): list of weights for each condition in `images_and_prompts` num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. num_inference_steps (`int`, *optional*, defaults to 100): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -195,7 +193,7 @@ def interpolate( negative_prior_prompt (`str`, *optional*): The prompt not to guide the prior diffusion process. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). guidance_scale (`float`, *optional*, defaults to 4.0): @@ -249,7 +247,7 @@ def interpolate( def _encode_image( self, - image: Union[torch.Tensor, List[PIL.Image.Image]], + image: torch.Tensor | list[PIL.Image.Image], device, num_images_per_prompt, ): @@ -341,7 +339,7 @@ def _encode_prompt( text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0) if do_classifier_free_guidance: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -402,22 +400,22 @@ def _encode_prompt( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]], - image: Union[torch.Tensor, List[torch.Tensor], PIL.Image.Image, List[PIL.Image.Image]], + prompt: str | list[str], + image: torch.Tensor | list[torch.Tensor] | PIL.Image.Image | list[PIL.Image.Image], strength: float = 0.3, - negative_prompt: Optional[Union[str, List[str]]] = None, + negative_prompt: str | list[str] | None = None, num_images_per_prompt: int = 1, num_inference_steps: int = 25, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + generator: torch.Generator | list[torch.Generator] | None = None, guidance_scale: float = 4.0, - output_type: Optional[str] = "pt", # pt only + output_type: str | None = "pt", # pt only return_dict: bool = True, ): """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. strength (`float`, *optional*, defaults to 0.8): Conceptually, indicates how much to transform the reference `emb`. Must be between 0 and 1. `image` @@ -425,7 +423,7 @@ def __call__( denoising steps depends on the amount of noise initially added. emb (`torch.Tensor`): The image embedding. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -433,7 +431,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 100): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. guidance_scale (`float`, *optional*, defaults to 4.0): @@ -480,7 +478,7 @@ def __call__( prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt ) - if not isinstance(image, List): + if not isinstance(image, list): image = [image] if isinstance(image[0], torch.Tensor): diff --git a/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py b/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py index 57cc0270442d..97353c95c9c7 100644 --- a/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py +++ b/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py @@ -1,4 +1,4 @@ -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import torch from transformers import T5EncoderModel, T5Tokenizer @@ -96,17 +96,17 @@ def encode_prompt( num_images_per_prompt=1, device=None, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, _cut_context=False, - attention_mask: Optional[torch.Tensor] = None, - negative_attention_mask: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, + negative_attention_mask: torch.Tensor | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`, *optional*): torch device to place the resulting embeddings on @@ -114,7 +114,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -181,7 +181,7 @@ def encode_prompt( attention_mask = attention_mask.repeat(num_images_per_prompt, 1) # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size @@ -335,36 +335,36 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_inference_steps: int = 25, guidance_scale: float = 3.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - height: Optional[int] = 1024, - width: Optional[int] = 1024, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - negative_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + height: int | None = 1024, + width: int | None = 1024, + generator: torch.Generator | list[torch.Generator] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + negative_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, latents=None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. num_inference_steps (`int`, *optional*, defaults to 25): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps` timesteps are used. Must be in descending order. guidance_scale (`float`, *optional*, defaults to 3.0): @@ -373,7 +373,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -386,7 +386,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. prompt_embeds (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py b/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py index 73c268897502..beb4caafb6d3 100644 --- a/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py +++ b/src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py @@ -1,5 +1,5 @@ import inspect -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import PIL import PIL.Image @@ -111,17 +111,17 @@ def encode_prompt( num_images_per_prompt=1, device=None, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, _cut_context=True, - attention_mask: Optional[torch.Tensor] = None, - negative_attention_mask: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, + negative_attention_mask: torch.Tensor | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`, *optional*): torch device to place the resulting embeddings on @@ -129,7 +129,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -196,7 +196,7 @@ def encode_prompt( attention_mask = attention_mask.repeat(num_images_per_prompt, 1) # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size @@ -401,32 +401,32 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - image: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]] = None, + prompt: str | list[str] = None, + image: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image] = None, strength: float = 0.3, num_inference_steps: int = 25, guidance_scale: float = 3.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - negative_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + negative_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, or tensor representing an image batch, that will be used as the starting point for the process. strength (`float`, *optional*, defaults to 0.8): @@ -444,13 +444,13 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. prompt_embeds (`torch.Tensor`, *optional*): @@ -474,7 +474,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py index 2b666f0ec697..1c94a8219e2a 100644 --- a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py +++ b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py @@ -13,7 +13,7 @@ # limitations under the License. import html -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import regex as re import torch @@ -312,10 +312,10 @@ def get_sparse_params(self, sample, device): def _encode_prompt_qwen( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + device: torch.device | None = None, max_sequence_length: int = 256, - dtype: Optional[torch.dtype] = None, + dtype: torch.dtype | None = None, ): """ Encode prompt using Qwen2.5-VL text encoder. @@ -324,14 +324,14 @@ def _encode_prompt_qwen( video generation. Args: - prompt (Union[str, List[str]]): Input prompt or list of prompts + prompt (str | list[str]): Input prompt or list of prompts device (torch.device): Device to run encoding on num_videos_per_prompt (int): Number of videos to generate per prompt max_sequence_length (int): Maximum sequence length for tokenization dtype (torch.dtype): Data type for embeddings Returns: - Tuple[torch.Tensor, torch.Tensor]: Text embeddings and cumulative sequence lengths + tuple[torch.Tensor, torch.Tensor]: Text embeddings and cumulative sequence lengths """ device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -382,9 +382,9 @@ def _encode_prompt_qwen( def _encode_prompt_clip( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt: str | list[str], + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): """ Encode prompt using CLIP text encoder. @@ -393,7 +393,7 @@ def _encode_prompt_clip( semantic information. Args: - prompt (Union[str, List[str]]): Input prompt or list of prompts + prompt (str | list[str]): Input prompt or list of prompts device (torch.device): Device to run encoding on num_videos_per_prompt (int): Number of videos to generate per prompt dtype (torch.dtype): Data type for embeddings @@ -419,11 +419,11 @@ def _encode_prompt_clip( def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_videos_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes a single prompt (positive or negative) into text encoder hidden states. @@ -432,7 +432,7 @@ def encode_prompt( representations for video generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): Prompt to be encoded. num_videos_per_prompt (`int`, *optional*, defaults to 1): Number of videos to generate per prompt. @@ -444,7 +444,7 @@ def encode_prompt( Torch dtype. Returns: - Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - Qwen text embeddings of shape (batch_size * num_videos_per_prompt, sequence_length, embedding_dim) - CLIP pooled embeddings of shape (batch_size * num_videos_per_prompt, clip_embedding_dim) - Cumulative sequence lengths (`cu_seqlens`) for Qwen embeddings of shape (batch_size * @@ -603,10 +603,10 @@ def prepare_latents( height: int = 480, width: int = 832, num_frames: int = 81, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: """ Prepare initial latent variables for video generation. @@ -683,37 +683,35 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] | None = None, height: int = 512, width: int = 768, num_frames: int = 121, num_inference_steps: int = 50, guidance_scale: float = 5.0, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds_qwen: Optional[torch.Tensor] = None, - prompt_embeds_clip: Optional[torch.Tensor] = None, - negative_prompt_embeds_qwen: Optional[torch.Tensor] = None, - negative_prompt_embeds_clip: Optional[torch.Tensor] = None, - prompt_cu_seqlens: Optional[torch.Tensor] = None, - negative_prompt_cu_seqlens: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds_qwen: torch.Tensor | None = None, + prompt_embeds_clip: torch.Tensor | None = None, + negative_prompt_embeds_qwen: torch.Tensor | None = None, + negative_prompt_embeds_clip: torch.Tensor | None = None, + prompt_cu_seqlens: torch.Tensor | None = None, + negative_prompt_cu_seqlens: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the video generation. If not defined, pass `prompt_embeds` instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to avoid during video generation. If not defined, pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale` < `1`). height (`int`, defaults to `512`): @@ -728,7 +726,7 @@ def __call__( Guidance scale as defined in classifier-free guidance. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of videos to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A torch generator to make generation deterministic. latents (`torch.Tensor`, *optional*): Pre-generated noisy latents. @@ -742,7 +740,7 @@ def __call__( Whether or not to return a [`KandinskyPipelineOutput`]. callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): A function that is called at the end of each denoising step. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. max_sequence_length (`int`, defaults to `512`): The maximum sequence length for text encoding. diff --git a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2i.py b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2i.py index f965cdad8f3e..244db7300767 100644 --- a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2i.py +++ b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2i.py @@ -13,7 +13,7 @@ # limitations under the License. import html -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import numpy as np import regex as re @@ -183,11 +183,11 @@ def __init__( def _encode_prompt_qwen( self, - prompt: List[str], - image: Optional[PipelineImageInput] = None, - device: Optional[torch.device] = None, + prompt: list[str], + image: PipelineImageInput | None = None, + device: torch.device | None = None, max_sequence_length: int = 1024, - dtype: Optional[torch.dtype] = None, + dtype: torch.dtype | None = None, ): """ Encode prompt using Qwen2.5-VL text encoder. @@ -196,14 +196,14 @@ def _encode_prompt_qwen( image generation. Args: - prompt List[str]: Input list of prompts + prompt list[str]: Input list of prompts image (PipelineImageInput): Input list of images to condition the generation on device (torch.device): Device to run encoding on max_sequence_length (int): Maximum sequence length for tokenization dtype (torch.dtype): Data type for embeddings Returns: - Tuple[torch.Tensor, torch.Tensor]: Text embeddings and cumulative sequence lengths + tuple[torch.Tensor, torch.Tensor]: Text embeddings and cumulative sequence lengths """ device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -258,9 +258,9 @@ def _encode_prompt_qwen( def _encode_prompt_clip( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt: str | list[str], + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): """ Encode prompt using CLIP text encoder. @@ -269,7 +269,7 @@ def _encode_prompt_clip( semantic information. Args: - prompt (Union[str, List[str]]): Input prompt or list of prompts + prompt (str | list[str]): Input prompt or list of prompts device (torch.device): Device to run encoding on dtype (torch.dtype): Data type for embeddings @@ -294,12 +294,12 @@ def _encode_prompt_clip( def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], image: torch.Tensor, num_images_per_prompt: int = 1, max_sequence_length: int = 1024, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes a single prompt (positive or negative) into text encoder hidden states. @@ -308,7 +308,7 @@ def encode_prompt( representations for image generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): Prompt to be encoded. num_images_per_prompt (`int`, *optional*, defaults to 1): Number of images to generate per prompt. @@ -320,7 +320,7 @@ def encode_prompt( Torch dtype. Returns: - Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - Qwen text embeddings of shape (batch_size * num_images_per_prompt, sequence_length, embedding_dim) - CLIP pooled embeddings of shape (batch_size * num_images_per_prompt, clip_embedding_dim) - Cumulative sequence lengths (`cu_seqlens`) for Qwen embeddings of shape (batch_size * @@ -486,10 +486,10 @@ def prepare_latents( num_channels_latents: int = 16, height: int = 1024, width: int = 1024, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: """ Prepare initial latent variables for image-to-image generation. @@ -568,27 +568,25 @@ def interrupt(self): def __call__( self, image: PipelineImageInput, - prompt: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 3.5, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds_qwen: Optional[torch.Tensor] = None, - prompt_embeds_clip: Optional[torch.Tensor] = None, - negative_prompt_embeds_qwen: Optional[torch.Tensor] = None, - negative_prompt_embeds_clip: Optional[torch.Tensor] = None, - prompt_cu_seqlens: Optional[torch.Tensor] = None, - negative_prompt_cu_seqlens: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds_qwen: torch.Tensor | None = None, + prompt_embeds_clip: torch.Tensor | None = None, + negative_prompt_embeds_qwen: torch.Tensor | None = None, + negative_prompt_embeds_clip: torch.Tensor | None = None, + prompt_cu_seqlens: torch.Tensor | None = None, + negative_prompt_cu_seqlens: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int, None], PipelineCallback | MultiPipelineCallbacks] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 1024, ): r""" @@ -597,9 +595,9 @@ def __call__( Args: image (`PipelineImageInput`): The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, pass `prompt_embeds` instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to avoid during image generation. If not defined, pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale` < `1`). height (`int`): @@ -612,7 +610,7 @@ def __call__( Guidance scale as defined in classifier-free guidance. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A torch generator to make generation deterministic. latents (`torch.Tensor`, *optional*): Pre-generated noisy latents. diff --git a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2v.py b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2v.py index d457c9b69657..ad4bb182d248 100644 --- a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2v.py +++ b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_i2v.py @@ -13,7 +13,7 @@ # limitations under the License. import html -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import regex as re import torch @@ -309,10 +309,10 @@ def get_sparse_params(self, sample, device): def _encode_prompt_qwen( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + device: torch.device | None = None, max_sequence_length: int = 256, - dtype: Optional[torch.dtype] = None, + dtype: torch.dtype | None = None, ): """ Encode prompt using Qwen2.5-VL text encoder. @@ -321,13 +321,13 @@ def _encode_prompt_qwen( video generation. Args: - prompt (Union[str, List[str]]): Input prompt or list of prompts + prompt (str | list[str]): Input prompt or list of prompts device (torch.device): Device to run encoding on max_sequence_length (int): Maximum sequence length for tokenization dtype (torch.dtype): Data type for embeddings Returns: - Tuple[torch.Tensor, torch.Tensor]: Text embeddings and cumulative sequence lengths + tuple[torch.Tensor, torch.Tensor]: Text embeddings and cumulative sequence lengths """ device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -377,9 +377,9 @@ def _encode_prompt_qwen( def _encode_prompt_clip( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt: str | list[str], + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): """ Encode prompt using CLIP text encoder. @@ -388,7 +388,7 @@ def _encode_prompt_clip( semantic information. Args: - prompt (Union[str, List[str]]): Input prompt or list of prompts + prompt (str | list[str]): Input prompt or list of prompts device (torch.device): Device to run encoding on dtype (torch.dtype): Data type for embeddings @@ -453,11 +453,11 @@ def normalize_first_frame(self, latents, reference_frames=5, clump_values=False) def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_videos_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes a single prompt (positive or negative) into text encoder hidden states. @@ -466,7 +466,7 @@ def encode_prompt( representations for video generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): Prompt to be encoded. num_videos_per_prompt (`int`, *optional*, defaults to 1): Number of videos to generate per prompt. @@ -478,7 +478,7 @@ def encode_prompt( Torch dtype. Returns: - Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - Qwen text embeddings of shape (batch_size * num_videos_per_prompt, sequence_length, embedding_dim) - CLIP pooled embeddings of shape (batch_size * num_videos_per_prompt, clip_embedding_dim) - Cumulative sequence lengths (`cu_seqlens`) for Qwen embeddings of shape (batch_size * @@ -641,10 +641,10 @@ def prepare_latents( height: int = 480, width: int = 832, num_frames: int = 81, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: """ Prepare initial latent variables for image-to-video generation. @@ -750,28 +750,26 @@ def interrupt(self): def __call__( self, image: PipelineImageInput, - prompt: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] | None = None, height: int = 512, width: int = 768, num_frames: int = 121, num_inference_steps: int = 50, guidance_scale: float = 5.0, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds_qwen: Optional[torch.Tensor] = None, - prompt_embeds_clip: Optional[torch.Tensor] = None, - negative_prompt_embeds_qwen: Optional[torch.Tensor] = None, - negative_prompt_embeds_clip: Optional[torch.Tensor] = None, - prompt_cu_seqlens: Optional[torch.Tensor] = None, - negative_prompt_cu_seqlens: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds_qwen: torch.Tensor | None = None, + prompt_embeds_clip: torch.Tensor | None = None, + negative_prompt_embeds_qwen: torch.Tensor | None = None, + negative_prompt_embeds_clip: torch.Tensor | None = None, + prompt_cu_seqlens: torch.Tensor | None = None, + negative_prompt_cu_seqlens: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int, None], PipelineCallback | MultiPipelineCallbacks] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" @@ -780,9 +778,9 @@ def __call__( Args: image (`PipelineImageInput`): The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the video generation. If not defined, pass `prompt_embeds` instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to avoid during video generation. If not defined, pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale` < `1`). height (`int`, defaults to `512`): @@ -797,7 +795,7 @@ def __call__( Guidance scale as defined in classifier-free guidance. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of videos to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A torch generator to make generation deterministic. latents (`torch.Tensor`, *optional*): Pre-generated noisy latents. diff --git a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_t2i.py b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_t2i.py index bb5c40327b4e..2a58d4bed33a 100644 --- a/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_t2i.py +++ b/src/diffusers/pipelines/kandinsky5/pipeline_kandinsky_t2i.py @@ -13,7 +13,7 @@ # limitations under the License. import html -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import numpy as np import regex as re @@ -184,10 +184,10 @@ def __init__( def _encode_prompt_qwen( self, - prompt: List[str], - device: Optional[torch.device] = None, + prompt: list[str], + device: torch.device | None = None, max_sequence_length: int = 512, - dtype: Optional[torch.dtype] = None, + dtype: torch.dtype | None = None, ): """ Encode prompt using Qwen2.5-VL text encoder. @@ -196,13 +196,13 @@ def _encode_prompt_qwen( image generation. Args: - prompt List[str]: Input list of prompts + prompt list[str]: Input list of prompts device (torch.device): Device to run encoding on max_sequence_length (int): Maximum sequence length for tokenization dtype (torch.dtype): Data type for embeddings Returns: - Tuple[torch.Tensor, torch.Tensor]: Text embeddings and cumulative sequence lengths + tuple[torch.Tensor, torch.Tensor]: Text embeddings and cumulative sequence lengths """ device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -252,9 +252,9 @@ def _encode_prompt_qwen( def _encode_prompt_clip( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt: str | list[str], + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): """ Encode prompt using CLIP text encoder. @@ -263,7 +263,7 @@ def _encode_prompt_clip( semantic information. Args: - prompt (Union[str, List[str]]): Input prompt or list of prompts + prompt (str | list[str]): Input prompt or list of prompts device (torch.device): Device to run encoding on dtype (torch.dtype): Data type for embeddings @@ -288,11 +288,11 @@ def _encode_prompt_clip( def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes a single prompt (positive or negative) into text encoder hidden states. @@ -301,7 +301,7 @@ def encode_prompt( representations for image generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): Prompt to be encoded. num_images_per_prompt (`int`, *optional*, defaults to 1): Number of images to generate per prompt. @@ -313,7 +313,7 @@ def encode_prompt( Torch dtype. Returns: - Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - Qwen text embeddings of shape (batch_size * num_images_per_prompt, sequence_length, embedding_dim) - CLIP pooled embeddings of shape (batch_size * num_images_per_prompt, clip_embedding_dim) - Cumulative sequence lengths (`cu_seqlens`) for Qwen embeddings of shape (batch_size * @@ -472,10 +472,10 @@ def prepare_latents( num_channels_latents: int = 16, height: int = 1024, width: int = 1024, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: """ Prepare initial latent variables for text-to-image generation. @@ -535,36 +535,34 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] | None = None, height: int = 1024, width: int = 1024, num_inference_steps: int = 50, guidance_scale: float = 3.5, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds_qwen: Optional[torch.Tensor] = None, - prompt_embeds_clip: Optional[torch.Tensor] = None, - negative_prompt_embeds_qwen: Optional[torch.Tensor] = None, - negative_prompt_embeds_clip: Optional[torch.Tensor] = None, - prompt_cu_seqlens: Optional[torch.Tensor] = None, - negative_prompt_cu_seqlens: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds_qwen: torch.Tensor | None = None, + prompt_embeds_clip: torch.Tensor | None = None, + negative_prompt_embeds_qwen: torch.Tensor | None = None, + negative_prompt_embeds_clip: torch.Tensor | None = None, + prompt_cu_seqlens: torch.Tensor | None = None, + negative_prompt_cu_seqlens: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int, None], PipelineCallback | MultiPipelineCallbacks] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" The call function to the pipeline for text-to-image generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, pass `prompt_embeds` instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to avoid during image generation. If not defined, pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale` < `1`). height (`int`, defaults to `1024`): @@ -577,7 +575,7 @@ def __call__( Guidance scale as defined in classifier-free guidance. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A torch generator to make generation deterministic. latents (`torch.Tensor`, *optional*): Pre-generated noisy latents. diff --git a/src/diffusers/pipelines/kandinsky5/pipeline_output.py b/src/diffusers/pipelines/kandinsky5/pipeline_output.py index 2172ddff7e22..9b41d2f14424 100644 --- a/src/diffusers/pipelines/kandinsky5/pipeline_output.py +++ b/src/diffusers/pipelines/kandinsky5/pipeline_output.py @@ -11,8 +11,8 @@ class KandinskyPipelineOutput(BaseOutput): Output class for kandinsky video pipelines. Args: - frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): - List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): + list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape `(batch_size, num_frames, channels, height, width)`. """ @@ -26,7 +26,7 @@ class KandinskyImagePipelineOutput(BaseOutput): Output class for kandinsky image pipelines. Args: - image (`torch.Tensor`, `np.ndarray`, or List[PIL.Image.Image]): + image (`torch.Tensor`, `np.ndarray`, or list[PIL.Image.Image]): List of image outputs - It can be a nested list of length `batch_size,` with each sub-list containing denoised PIL image. It can also be a NumPy array or Torch tensor of shape `(batch_size, channels, height, width)`. diff --git a/src/diffusers/pipelines/kolors/pipeline_kolors.py b/src/diffusers/pipelines/kolors/pipeline_kolors.py index 7c8468bcb109..1e11faf8b9b6 100644 --- a/src/diffusers/pipelines/kolors/pipeline_kolors.py +++ b/src/diffusers/pipelines/kolors/pipeline_kolors.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection @@ -63,10 +63,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -81,15 +81,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -199,21 +199,21 @@ def __init__( def encode_prompt( self, prompt, - device: Optional[torch.device] = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, negative_prompt=None, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 256, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -221,7 +221,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -291,7 +291,7 @@ def encode_prompt( if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt: negative_prompt_embeds = torch.zeros_like(prompt_embeds) elif do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -648,45 +648,43 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_end: Optional[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, + denoising_end: float | None = None, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - original_size: Optional[Tuple[int, int]] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Optional[Tuple[int, int]] = None, - negative_original_size: Optional[Tuple[int, int]] = None, - negative_crops_coords_top_left: Tuple[int, int] = (0, 0), - negative_target_size: Optional[Tuple[int, int]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + cross_attention_kwargs: dict[str, Any] | None = None, + original_size: tuple[int, int] | None = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] | None = None, + negative_original_size: tuple[int, int] | None = None, + negative_crops_coords_top_left: tuple[int, int] = (0, 0), + negative_target_size: tuple[int, int] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 256, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): @@ -702,11 +700,11 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -723,7 +721,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -732,7 +730,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -754,7 +752,7 @@ def __call__( weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -768,31 +766,31 @@ def __call__( A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a specific image resolution. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a target image resolution. It should be as same as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more @@ -802,7 +800,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -877,8 +875,12 @@ def __call__( ) # 4. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) # 5. Prepare latent variables diff --git a/src/diffusers/pipelines/kolors/pipeline_kolors_img2img.py b/src/diffusers/pipelines/kolors/pipeline_kolors_img2img.py index 10a7962c258c..d9b519267216 100644 --- a/src/diffusers/pipelines/kolors/pipeline_kolors_img2img.py +++ b/src/diffusers/pipelines/kolors/pipeline_kolors_img2img.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import PIL.Image import torch @@ -67,7 +67,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -82,10 +82,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -100,15 +100,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -219,21 +219,21 @@ def __init__( def encode_prompt( self, prompt, - device: Optional[torch.device] = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, negative_prompt=None, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 256, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -241,7 +241,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -311,7 +311,7 @@ def encode_prompt( if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt: negative_prompt_embeds = torch.zeros_like(prompt_embeds) elif do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -780,51 +780,49 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, image: PipelineImageInput = None, strength: float = 0.3, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_start: Optional[float] = None, - denoising_end: Optional[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, + denoising_start: float | None = None, + denoising_end: float | None = None, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - original_size: Optional[Tuple[int, int]] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Optional[Tuple[int, int]] = None, - negative_original_size: Optional[Tuple[int, int]] = None, - negative_crops_coords_top_left: Tuple[int, int] = (0, 0), - negative_target_size: Optional[Tuple[int, int]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + cross_attention_kwargs: dict[str, Any] | None = None, + original_size: tuple[int, int] | None = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] | None = None, + negative_original_size: tuple[int, int] | None = None, + negative_crops_coords_top_left: tuple[int, int] = (0, 0), + negative_target_size: tuple[int, int] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 256, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - image (`torch.Tensor` or `PIL.Image.Image` or `np.ndarray` or `List[torch.Tensor]` or `List[PIL.Image.Image]` or `List[np.ndarray]`): + image (`torch.Tensor` or `PIL.Image.Image` or `np.ndarray` or `list[torch.Tensor]` or `list[PIL.Image.Image]` or `list[np.ndarray]`): The image(s) to modify with the pipeline. strength (`float`, *optional*, defaults to 0.3): Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1. `image` @@ -846,11 +844,11 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -874,7 +872,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -883,7 +881,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -905,7 +903,7 @@ def __call__( weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -919,31 +917,31 @@ def __call__( A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a specific image resolution. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a target image resolution. It should be as same as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more @@ -953,7 +951,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1034,8 +1032,12 @@ def __call__( def denoising_value_valid(dnv): return isinstance(dnv, float) and 0 < dnv < 1 + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) timesteps, num_inference_steps = self.get_timesteps( diff --git a/src/diffusers/pipelines/kolors/pipeline_output.py b/src/diffusers/pipelines/kolors/pipeline_output.py index 310ee7e8a89b..0418191cc6e5 100644 --- a/src/diffusers/pipelines/kolors/pipeline_output.py +++ b/src/diffusers/pipelines/kolors/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -13,9 +12,9 @@ class KolorsPipelineOutput(BaseOutput): Output class for Kolors pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image] | np.ndarray diff --git a/src/diffusers/pipelines/kolors/text_encoder.py b/src/diffusers/pipelines/kolors/text_encoder.py index 6fd17156a116..434f4fed6fbb 100644 --- a/src/diffusers/pipelines/kolors/text_encoder.py +++ b/src/diffusers/pipelines/kolors/text_encoder.py @@ -13,7 +13,6 @@ # limitations under the License. import math -from typing import List, Optional, Tuple import torch import torch.nn.functional as F @@ -234,7 +233,7 @@ def split_tensor_along_last_dim( tensor: torch.Tensor, num_partitions: int, contiguous_split_chunks: bool = False, -) -> List[torch.Tensor]: +) -> list[torch.Tensor]: """Split a tensor along its last dimension. Arguments: @@ -566,8 +565,8 @@ def forward( attention_mask, rotary_pos_emb, kv_caches=None, - use_cache: Optional[bool] = True, - output_hidden_states: Optional[bool] = False, + use_cache: bool | None = True, + output_hidden_states: bool | None = False, ): if not kv_caches: kv_caches = [None for _ in range(self.num_layers)] @@ -782,6 +781,9 @@ def __init__(self, config: ChatGLMConfig, device=None, empty_init=True): self.prefix_encoder = PrefixEncoder(config) self.dropout = torch.nn.Dropout(0.1) + if hasattr(self, "post_init"): + self.post_init() + def get_input_embeddings(self): return self.embedding.word_embeddings @@ -799,19 +801,19 @@ def get_prompt(self, batch_size, device, dtype=torch.half): def forward( self, input_ids, - position_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.BoolTensor] = None, - full_attention_mask: Optional[torch.BoolTensor] = None, - past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, - inputs_embeds: Optional[torch.Tensor] = None, - use_cache: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.BoolTensor | None = None, + full_attention_mask: torch.BoolTensor | None = None, + past_key_values: tuple[tuple[torch.Tensor, torch.Tensor], ...] | None = None, + inputs_embeds: torch.Tensor | None = None, + use_cache: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, ): output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) - use_cache = use_cache if use_cache is not None else self.config.use_cache + use_cache = use_cache if use_cache is not None else getattr(self.config, "use_cache", None) return_dict = return_dict if return_dict is not None else self.config.use_return_dict batch_size, seq_length = input_ids.shape diff --git a/src/diffusers/pipelines/kolors/tokenizer.py b/src/diffusers/pipelines/kolors/tokenizer.py index b824ba12e079..c0a1831bdeee 100644 --- a/src/diffusers/pipelines/kolors/tokenizer.py +++ b/src/diffusers/pipelines/kolors/tokenizer.py @@ -15,7 +15,6 @@ import json import os import re -from typing import Dict, List, Optional, Union from sentencepiece import SentencePieceProcessor from transformers import PreTrainedTokenizer @@ -61,7 +60,7 @@ def tokenize(self, s: str, encode_special_tokens=False): else: return self.sp_model.EncodeAsPieces(s) - def encode(self, s: str, bos: bool = False, eos: bool = False) -> List[int]: + def encode(self, s: str, bos: bool = False, eos: bool = False) -> list[int]: assert isinstance(s, str) t = self.sp_model.encode(s) if bos: @@ -70,7 +69,7 @@ def encode(self, s: str, bos: bool = False, eos: bool = False) -> List[int]: t = t + [self.eos_id] return t - def decode(self, t: List[int]) -> str: + def decode(self, t: list[int]) -> str: text, buffer = "", [] for token in t: if token in self.index_special_tokens: @@ -84,7 +83,7 @@ def decode(self, t: List[int]) -> str: text += self.sp_model.decode(buffer) return text - def decode_tokens(self, tokens: List[str]) -> str: + def decode_tokens(self, tokens: list[str]) -> str: text = self.sp_model.DecodePieces(tokens) return text @@ -192,7 +191,7 @@ def _convert_id_to_token(self, index): """Converts an index (integer) in a token (str) using the vocab.""" return self.tokenizer.convert_id_to_token(index) - def convert_tokens_to_string(self, tokens: List[str]) -> str: + def convert_tokens_to_string(self, tokens: list[str]) -> str: return self.tokenizer.decode_tokens(tokens) def save_vocabulary(self, save_directory, filename_prefix=None): @@ -206,7 +205,7 @@ def save_vocabulary(self, save_directory, filename_prefix=None): An optional prefix to add to the named of the saved files. Returns: - `Tuple(str)`: Paths to the files saved. + `tuple(str)`: Paths to the files saved. """ if os.path.isdir(save_directory): vocab_file = os.path.join(save_directory, self.vocab_files_names["vocab_file"]) @@ -246,8 +245,8 @@ def build_chat_input(self, query, history=None, role="user"): return self.batch_encode_plus([input_ids], return_tensors="pt", is_split_into_words=True) def build_inputs_with_special_tokens( - self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None - ) -> List[int]: + self, token_ids_0: list[int], token_ids_1: list[int] | None = None + ) -> list[int]: """ Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and adding special tokens. A BERT sequence has the following format: @@ -256,13 +255,13 @@ def build_inputs_with_special_tokens( - pair of sequences: `[CLS] A [SEP] B [SEP]` Args: - token_ids_0 (`List[int]`): - List of IDs to which the special tokens will be added. - token_ids_1 (`List[int]`, *optional*): + token_ids_0 (`list[int]`): + list of IDs to which the special tokens will be added. + token_ids_1 (`list[int]`, *optional*): Optional second list of IDs for sequence pairs. Returns: - `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens. + `list[int]`: list of [input IDs](../glossary#input-ids) with the appropriate special tokens. """ prefix_tokens = self.get_prefix_tokens() token_ids_0 = prefix_tokens + token_ids_0 @@ -272,19 +271,19 @@ def build_inputs_with_special_tokens( def _pad( self, - encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding], - max_length: Optional[int] = None, + encoded_inputs: dict[str, EncodedInput] | BatchEncoding, + max_length: int | None = None, padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, - pad_to_multiple_of: Optional[int] = None, - return_attention_mask: Optional[bool] = None, - padding_side: Optional[bool] = None, + pad_to_multiple_of: int | None = None, + return_attention_mask: bool | None = None, + padding_side: bool | None = None, ) -> dict: """ Pad encoded inputs (on left/right and up to predefined length or max length in the batch) Args: encoded_inputs: - Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`). + Dictionary of tokenized inputs (`list[int]`) or batch of tokenized inputs (`list[list[int]]`). max_length: maximum length of the returned list and optionally padding length (see below). Will truncate by taking into account the special tokens. padding_strategy: PaddingStrategy to use for padding. diff --git a/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py b/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py index 59f733a498ed..6861438e4c63 100644 --- a/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py +++ b/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py @@ -16,7 +16,7 @@ # and https://github.com/hojonathanho/diffusion import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import PIL.Image import torch @@ -53,7 +53,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -68,10 +68,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -86,15 +86,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -208,7 +208,7 @@ def __init__( scheduler: LCMScheduler, safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, - image_encoder: Optional[CLIPVisionModelWithProjection] = None, + image_encoder: CLIPVisionModelWithProjection | None = None, requires_safety_checker: bool = True, ): super().__init__() @@ -245,16 +245,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -262,7 +262,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -361,7 +361,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -635,10 +635,10 @@ def get_timesteps(self, num_inference_steps, strength, device): def check_inputs( self, - prompt: Union[str, List[str]], + prompt: str | list[str], strength: float, callback_steps: int, - prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, ip_adapter_image=None, ip_adapter_image_embeds=None, callback_on_step_end_tensor_inputs=None, @@ -710,32 +710,32 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, image: PipelineImageInput = None, num_inference_steps: int = 4, strength: float = 0.8, original_inference_steps: int = None, - timesteps: List[int] = None, + timesteps: list[int] = None, guidance_scale: float = 8.5, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + cross_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -749,7 +749,7 @@ def __call__( we will draw `num_inference_steps` evenly spaced timesteps from as our final timestep schedule, following the Skipping-Step method in the paper (see Section 4.3). If not set this will default to the scheduler's `original_inference_steps` attribute. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps` timesteps on the original LCM training/distillation timestep schedule are used. Must be in descending order. @@ -761,7 +761,7 @@ def __call__( 0`). num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -773,7 +773,7 @@ def __call__( provided, text embeddings are generated from the `prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -794,7 +794,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -881,10 +881,14 @@ def __call__( image = self.image_processor.preprocess(image) # 5. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, - device, + timestep_device, timesteps, original_inference_steps=original_inference_steps, strength=strength, diff --git a/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py b/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py index e463884618f5..60f59ec7f9d3 100644 --- a/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py +++ b/src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py @@ -16,7 +16,7 @@ # and https://github.com/hojonathanho/diffusion import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection @@ -73,10 +73,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -91,15 +91,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -187,7 +187,7 @@ def __init__( scheduler: LCMScheduler, safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, - image_encoder: Optional[CLIPVisionModelWithProjection] = None, + image_encoder: CLIPVisionModelWithProjection | None = None, requires_safety_checker: bool = True, ): super().__init__() @@ -230,16 +230,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -247,7 +247,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -346,7 +346,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -565,11 +565,11 @@ def prepare_extra_step_kwargs(self, generator, eta): # Currently StableDiffusionPipeline.check_inputs with negative prompt stuff removed def check_inputs( self, - prompt: Union[str, List[str]], + prompt: str | list[str], height: int, width: int, callback_steps: int, - prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, ip_adapter_image=None, ip_adapter_image_embeds=None, callback_on_step_end_tensor_inputs=None, @@ -641,32 +641,32 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 4, original_inference_steps: int = None, - timesteps: List[int] = None, + timesteps: list[int] = None, guidance_scale: float = 8.5, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + cross_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -680,7 +680,7 @@ def __call__( we will draw `num_inference_steps` evenly spaced timesteps from as our final timestep schedule, following the Skipping-Step method in the paper (see Section 4.3). If not set this will default to the scheduler's `original_inference_steps` attribute. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps` timesteps on the original LCM training/distillation timestep schedule are used. Must be in descending order. @@ -692,7 +692,7 @@ def __call__( 0`). num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -704,7 +704,7 @@ def __call__( provided, text embeddings are generated from the `prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -725,7 +725,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -815,8 +815,16 @@ def __call__( ) # 4. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, original_inference_steps=original_inference_steps + self.scheduler, + num_inference_steps, + timestep_device, + timesteps, + original_inference_steps=original_inference_steps, ) # 5. Prepare latent variable diff --git a/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py b/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py index f1bf4701e31f..ec43988f9389 100644 --- a/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py +++ b/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py @@ -13,7 +13,6 @@ # limitations under the License. import inspect -from typing import List, Optional, Tuple, Union import torch import torch.nn as nn @@ -62,11 +61,11 @@ class LDMTextToImagePipeline(DiffusionPipeline): def __init__( self, - vqvae: Union[VQModel, AutoencoderKL], + vqvae: VQModel | AutoencoderKL, bert: PreTrainedModel, tokenizer: PreTrainedTokenizer, - unet: Union[UNet2DModel, UNet2DConditionModel], - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], + unet: UNet2DModel | UNet2DConditionModel, + scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler, ): super().__init__() self.register_modules(vqvae=vqvae, bert=bert, tokenizer=tokenizer, unet=unet, scheduler=scheduler) @@ -75,23 +74,23 @@ def __init__( @torch.no_grad() def __call__( self, - prompt: Union[str, List[str]], - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: Optional[int] = 50, - guidance_scale: Optional[float] = 1.0, - eta: Optional[float] = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + prompt: str | list[str], + height: int | None = None, + width: int | None = None, + num_inference_steps: int | None = 50, + guidance_scale: float | None = 1.0, + eta: float | None = 0.0, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, **kwargs, - ) -> Union[Tuple, ImagePipelineOutput]: + ) -> tuple | ImagePipelineOutput: r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -236,7 +235,7 @@ def __call__( logger = logging.get_logger(__name__) -LDMBERT_PRETRAINED_MODEL_ARCHIVE_LIST = [ +LDMBERT_PRETRAINED_MODEL_ARCHIVE_list = [ "ldm-bert", # See all LDMBert models at https://huggingface.co/models?filter=ldmbert ] @@ -297,7 +296,7 @@ def __init__( super().__init__(pad_token_id=pad_token_id, **kwargs) -def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): +def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: int | None = None): """ Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`. """ @@ -345,12 +344,12 @@ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): def forward( self, hidden_states: torch.Tensor, - key_value_states: Optional[torch.Tensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, - attention_mask: Optional[torch.Tensor] = None, - layer_head_mask: Optional[torch.Tensor] = None, + key_value_states: torch.Tensor | None = None, + past_key_value: tuple[torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + layer_head_mask: torch.Tensor | None = None, output_attentions: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: """Input shape: Batch x Time x Channel""" # if key_value_states are provided this layer is used as a cross-attention layer @@ -382,10 +381,10 @@ def forward( value_states = self._shape(self.v_proj(hidden_states), -1, bsz) if self.is_decoder: - # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. + # if cross_attention save tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. # Further calls to cross_attention layer can then reuse all cross-attention # key/value_states (first "if" case) - # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of + # if uni-directional self-attention (decoder) save tuple(torch.Tensor, torch.Tensor) of # all previous decoder key/value_states. Further calls to uni-directional self-attention # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case) # if encoder bi-directional self-attention `past_key_value` is always `None` @@ -479,8 +478,8 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.Tensor, layer_head_mask: torch.Tensor, - output_attentions: Optional[bool] = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + output_attentions: bool | None = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: """ Args: hidden_states (`torch.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)` @@ -591,14 +590,14 @@ def set_input_embeddings(self, value): def forward( self, input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - head_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple, BaseModelOutput]: + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + head_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> tuple | BaseModelOutput: r""" Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): diff --git a/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion_superresolution.py b/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion_superresolution.py index 631539e5c667..18cb8274f9b5 100644 --- a/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion_superresolution.py +++ b/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion_superresolution.py @@ -1,5 +1,4 @@ import inspect -from typing import List, Optional, Tuple, Union import numpy as np import PIL.Image @@ -59,14 +58,12 @@ def __init__( self, vqvae: VQModel, unet: UNet2DModel, - scheduler: Union[ - DDIMScheduler, - PNDMScheduler, - LMSDiscreteScheduler, - EulerDiscreteScheduler, - EulerAncestralDiscreteScheduler, - DPMSolverMultistepScheduler, - ], + scheduler: DDIMScheduler + | PNDMScheduler + | LMSDiscreteScheduler + | EulerDiscreteScheduler + | EulerAncestralDiscreteScheduler + | DPMSolverMultistepScheduler, ): super().__init__() self.register_modules(vqvae=vqvae, unet=unet, scheduler=scheduler) @@ -74,14 +71,14 @@ def __init__( @torch.no_grad() def __call__( self, - image: Union[torch.Tensor, PIL.Image.Image] = None, - batch_size: Optional[int] = 1, - num_inference_steps: Optional[int] = 100, - eta: Optional[float] = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - output_type: Optional[str] = "pil", + image: torch.Tensor | PIL.Image.Image = None, + batch_size: int | None = 1, + num_inference_steps: int | None = 100, + eta: float | None = 0.0, + generator: torch.Generator | list[torch.Generator] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - ) -> Union[Tuple, ImagePipelineOutput]: + ) -> tuple | ImagePipelineOutput: r""" The call function to the pipeline for generation. @@ -96,7 +93,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. output_type (`str`, *optional*, defaults to `"pil"`): diff --git a/src/diffusers/pipelines/latte/pipeline_latte.py b/src/diffusers/pipelines/latte/pipeline_latte.py index 4d42a7049ec9..eed7762cebf1 100644 --- a/src/diffusers/pipelines/latte/pipeline_latte.py +++ b/src/diffusers/pipelines/latte/pipeline_latte.py @@ -18,7 +18,7 @@ import re import urllib.parse as ul from dataclasses import dataclass -from typing import Callable, Dict, List, Optional, Tuple, Union +from typing import Callable import torch from transformers import T5EncoderModel, T5Tokenizer @@ -80,10 +80,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -98,15 +98,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -205,13 +205,13 @@ def mask_text_embeddings(self, emb, mask): # Adapted from diffusers.pipelines.deepfloyd_if.pipeline_if.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], do_classifier_free_guidance: bool = True, negative_prompt: str = "", num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, + device: torch.device | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, clean_caption: bool = False, mask_feature: bool = True, dtype=None, @@ -220,9 +220,9 @@ def encode_prompt( Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide the video generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For Latte, this should be "". @@ -614,46 +614,44 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, negative_prompt: str = "", num_inference_steps: int = 50, - timesteps: Optional[List[int]] = None, + timesteps: list[int] | None = None, guidance_scale: float = 7.5, num_images_per_prompt: int = 1, video_length: int = 16, height: int = 512, width: int = 512, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, output_type: str = "pil", return_dict: bool = True, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], clean_caption: bool = True, mask_feature: bool = True, enable_temporal_attentions: bool = True, decode_chunk_size: int = 14, - ) -> Union[LattePipelineOutput, Tuple]: + ) -> LattePipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the video generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_inference_steps (`int`, *optional*, defaults to 100): The number of denoising steps. More denoising steps usually lead to a higher quality video at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps` timesteps are used. Must be in descending order. guidance_scale (`float`, *optional*, defaults to 7.0): @@ -673,7 +671,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -691,9 +689,9 @@ def __call__( [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~pipelines.stable_diffusion.IFPipelineOutput`] instead of a plain tuple. - callback_on_step_end (`Callable[[int, int, Dict], None]`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): + callback_on_step_end (`Callable[[int, int], None]`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): A callback function or a list of callback functions to be called at the end of each denoising step. - callback_on_step_end_tensor_inputs (`List[str]`, *optional*): + callback_on_step_end_tensor_inputs (`list[str]`, *optional*): A list of tensor inputs that should be passed to the callback function. If not defined, all tensor inputs will be passed. clean_caption (`bool`, *optional*, defaults to `True`): @@ -767,7 +765,13 @@ def __call__( prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, timesteps + ) self._num_timesteps = len(timesteps) # 5. Prepare latents. diff --git a/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py b/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py index fbf4dc23d043..864f9feeb5aa 100644 --- a/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py +++ b/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py @@ -1,7 +1,7 @@ import inspect import math from itertools import repeat -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch import torch.nn.functional as F @@ -115,7 +115,7 @@ def get_attention(self, step: int): return attention def aggregate_attention( - self, attention_maps, prompts, res: Union[int, Tuple[int]], from_where: List[str], is_cross: bool, select: int + self, attention_maps, prompts, res: int | tuple[int], from_where: list[str], is_cross: bool, select: int ): out = [[] for x in range(self.batch_size)] if isinstance(res, int): @@ -309,7 +309,7 @@ def __init__( text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, unet: UNet2DConditionModel, - scheduler: Union[DDIMScheduler, DPMSolverMultistepScheduler], + scheduler: DDIMScheduler | DPMSolverMultistepScheduler, safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, requires_safety_checker: bool = True, @@ -525,10 +525,10 @@ def encode_prompt( enable_edit_guidance, negative_prompt=None, editing_prompt=None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - editing_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt_embeds: torch.Tensor | None = None, + editing_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. @@ -540,11 +540,11 @@ def encode_prompt( number of images that should be generated per prompt enable_edit_guidance (`bool`): whether to perform any editing or reconstruct the input image instead - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - editing_prompt (`str` or `List[str]`, *optional*): + editing_prompt (`str` or `list[str]`, *optional*): Editing prompt(s) to be encoded. If not defined, one has to pass `editing_prompt_embeds` instead. editing_prompt_embeds (`torch.Tensor`, *optional*): Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not @@ -574,7 +574,7 @@ def encode_prompt( num_edit_tokens = None if negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif isinstance(negative_prompt, str): @@ -774,29 +774,29 @@ def disable_vae_tiling(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - negative_prompt: Optional[Union[str, List[str]]] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - output_type: Optional[str] = "pil", + negative_prompt: str | list[str] | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - editing_prompt: Optional[Union[str, List[str]]] = None, - editing_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - reverse_editing_direction: Optional[Union[bool, List[bool]]] = False, - edit_guidance_scale: Optional[Union[float, List[float]]] = 5, - edit_warmup_steps: Optional[Union[int, List[int]]] = 0, - edit_cooldown_steps: Optional[Union[int, List[int]]] = None, - edit_threshold: Optional[Union[float, List[float]]] = 0.9, - user_mask: Optional[torch.Tensor] = None, - sem_guidance: Optional[List[torch.Tensor]] = None, + editing_prompt: str | list[str] | None = None, + editing_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + reverse_editing_direction: bool | list[bool] | None = False, + edit_guidance_scale: float | list[float] | None = 5, + edit_warmup_steps: int | list[int] | None = 0, + edit_cooldown_steps: int | list[int] | None = None, + edit_threshold: float | list[float] | None = 0.9, + user_mask: torch.Tensor | None = None, + sem_guidance: list[torch.Tensor] | None = None, use_cross_attn_mask: bool = False, use_intersect_mask: bool = True, - attn_store_steps: Optional[List[int]] = [], + attn_store_steps: list[int] | None = [], store_averaged_over_steps: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, guidance_rescale: float = 0.0, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" @@ -805,7 +805,7 @@ def __call__( always be performed for the last inverted image(s). Args: - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). generator (`torch.Generator`, *optional*): @@ -817,7 +817,7 @@ def __call__( return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~pipelines.ledits_pp.LEditsPPDiffusionPipelineOutput`] instead of a plain tuple. - editing_prompt (`str` or `List[str]`, *optional*): + editing_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. The image is reconstructed by setting `editing_prompt = None`. Guidance direction of prompt should be specified via `reverse_editing_direction`. @@ -827,25 +827,25 @@ def __call__( negative_prompt_embeds (`torch.Tensor`, *optional*): Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. - reverse_editing_direction (`bool` or `List[bool]`, *optional*, defaults to `False`): + reverse_editing_direction (`bool` or `list[bool]`, *optional*, defaults to `False`): Whether the corresponding prompt in `editing_prompt` should be increased or decreased. - edit_guidance_scale (`float` or `List[float]`, *optional*, defaults to 5): + edit_guidance_scale (`float` or `list[float]`, *optional*, defaults to 5): Guidance scale for guiding the image generation. If provided as list values should correspond to `editing_prompt`. `edit_guidance_scale` is defined as `s_e` of equation 12 of [LEDITS++ Paper](https://huggingface.co/papers/2301.12247). - edit_warmup_steps (`float` or `List[float]`, *optional*, defaults to 10): + edit_warmup_steps (`float` or `list[float]`, *optional*, defaults to 10): Number of diffusion steps (for each prompt) for which guidance will not be applied. - edit_cooldown_steps (`float` or `List[float]`, *optional*, defaults to `None`): + edit_cooldown_steps (`float` or `list[float]`, *optional*, defaults to `None`): Number of diffusion steps (for each prompt) after which guidance will no longer be applied. - edit_threshold (`float` or `List[float]`, *optional*, defaults to 0.9): + edit_threshold (`float` or `list[float]`, *optional*, defaults to 0.9): Masking threshold of guidance. Threshold should be proportional to the image region that is modified. 'edit_threshold' is defined as 'λ' of equation 12 of [LEDITS++ Paper](https://huggingface.co/papers/2301.12247). user_mask (`torch.Tensor`, *optional*): User-provided mask for even better control over the editing process. This is helpful when LEDITS++'s implicit masks do not meet user preferences. - sem_guidance (`List[torch.Tensor]`, *optional*): - List of pre-generated guidance vectors to be applied at generation. Length of the list has to + sem_guidance (`list[torch.Tensor]`, *optional*): + list of pre-generated guidance vectors to be applied at generation. Length of the list has to correspond to `num_inference_steps`. use_cross_attn_mask (`bool`, defaults to `False`): Whether cross-attention masks are used. Cross-attention masks are always used when use_intersect_mask @@ -855,7 +855,7 @@ def __call__( Whether the masking term is calculated as intersection of cross-attention masks and masks derived from the noise estimate. Cross-attention mask are defined as 'M^1' and masks derived from the noise estimate are defined as 'M^2' of equation 12 of [LEDITS++ paper](https://huggingface.co/papers/2311.16711). - attn_store_steps (`List[int]`, *optional*): + attn_store_steps (`list[int]`, *optional*): Steps for which the attention maps are stored in the AttentionStore. Just for visualization purposes. store_averaged_over_steps (`bool`, defaults to `True`): Whether the attention maps for the 'attn_store_steps' are stored averaged over the diffusion steps. If @@ -875,7 +875,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1282,13 +1282,13 @@ def invert( source_guidance_scale: float = 3.5, num_inversion_steps: int = 30, skip: float = 0.15, - generator: Optional[torch.Generator] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - height: Optional[int] = None, - width: Optional[int] = None, - resize_mode: Optional[str] = "default", - crops_coords: Optional[Tuple[int, int, int, int]] = None, + generator: torch.Generator | None = None, + cross_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, + height: int | None = None, + width: int | None = None, + resize_mode: str | None = "default", + crops_coords: tuple[int, int, int, int] | None = None, ): r""" The function to the pipeline for image inversion as described by the [LEDITS++ @@ -1331,7 +1331,7 @@ def invert( image to fit within the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, cropping the excess. Note that resize_mode `fill` and `crop` are only supported for PIL image input. - crops_coords (`List[Tuple[int, int, int, int]]`, *optional*, defaults to `None`): + crops_coords (`list[tuple[int, int, int, int]]`, *optional*, defaults to `None`): The crop coordinates for each image in the batch. If `None`, will not crop the image. Returns: diff --git a/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py b/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py index 993957a052fc..a136770b9f26 100644 --- a/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py +++ b/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py @@ -14,7 +14,7 @@ import inspect import math -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch import torch.nn.functional as F @@ -145,7 +145,7 @@ def get_attention(self, step: int): return attention def aggregate_attention( - self, attention_maps, prompts, res: Union[int, Tuple[int]], from_where: List[str], is_cross: bool, select: int + self, attention_maps, prompts, res: int | tuple[int], from_where: list[str], is_cross: bool, select: int ): out = [[] for x in range(self.batch_size)] if isinstance(res, int): @@ -351,11 +351,11 @@ def __init__( tokenizer: CLIPTokenizer, tokenizer_2: CLIPTokenizer, unet: UNet2DConditionModel, - scheduler: Union[DPMSolverMultistepScheduler, DDIMScheduler], + scheduler: DPMSolverMultistepScheduler | DDIMScheduler, image_encoder: CLIPVisionModelWithProjection = None, feature_extractor: CLIPImageProcessor = None, force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, + add_watermarker: bool | None = None, ): super().__init__() @@ -399,18 +399,18 @@ def __init__( def encode_prompt( self, - device: Optional[torch.device] = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, enable_edit_guidance: bool = True, - editing_prompt: Optional[str] = None, - editing_prompt_embeds: Optional[torch.Tensor] = None, - editing_pooled_prompt_embeds: Optional[torch.Tensor] = None, + editing_prompt: str | None = None, + editing_prompt_embeds: torch.Tensor | None = None, + editing_pooled_prompt_embeds: torch.Tensor | None = None, ) -> object: r""" Encodes the prompt into text encoder hidden states. @@ -420,10 +420,10 @@ def encode_prompt( torch device num_images_per_prompt (`int`): number of images that should be generated per prompt - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders negative_prompt_embeds (`torch.Tensor`, *optional*): @@ -441,7 +441,7 @@ def encode_prompt( the output of the pre-final layer will be used for computing the prompt embeddings. enable_edit_guidance (`bool`): Whether to guide towards an editing prompt or not. - editing_prompt (`str` or `List[str]`, *optional*): + editing_prompt (`str` or `list[str]`, *optional*): Editing prompt(s) to be encoded. If not defined and 'enable_edit_guidance' is True, one has to pass `editing_prompt_embeds` instead. editing_prompt_embeds (`torch.Tensor`, *optional*): @@ -495,7 +495,7 @@ def encode_prompt( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) - uncond_tokens: List[str] + uncond_tokens: list[str] if batch_size != len(negative_prompt): raise ValueError( @@ -837,35 +837,35 @@ def prepare_unet(self, attention_store, PnP: bool = False): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - denoising_end: Optional[float] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - output_type: Optional[str] = "pil", + denoising_end: float | None = None, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, guidance_rescale: float = 0.0, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Optional[Tuple[int, int]] = None, - editing_prompt: Optional[Union[str, List[str]]] = None, - editing_prompt_embeddings: Optional[torch.Tensor] = None, - editing_pooled_prompt_embeds: Optional[torch.Tensor] = None, - reverse_editing_direction: Optional[Union[bool, List[bool]]] = False, - edit_guidance_scale: Optional[Union[float, List[float]]] = 5, - edit_warmup_steps: Optional[Union[int, List[int]]] = 0, - edit_cooldown_steps: Optional[Union[int, List[int]]] = None, - edit_threshold: Optional[Union[float, List[float]]] = 0.9, - sem_guidance: Optional[List[torch.Tensor]] = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] | None = None, + editing_prompt: str | list[str] | None = None, + editing_prompt_embeddings: torch.Tensor | None = None, + editing_pooled_prompt_embeds: torch.Tensor | None = None, + reverse_editing_direction: bool | list[bool] | None = False, + edit_guidance_scale: float | list[float] | None = 5, + edit_warmup_steps: int | list[int] | None = 0, + edit_cooldown_steps: int | list[int] | None = None, + edit_threshold: float | list[float] | None = 0.9, + sem_guidance: list[torch.Tensor] | None = None, use_cross_attn_mask: bool = False, use_intersect_mask: bool = False, - user_mask: Optional[torch.Tensor] = None, - attn_store_steps: Optional[List[int]] = [], + user_mask: torch.Tensor | None = None, + attn_store_steps: list[int] | None = [], store_averaged_over_steps: bool = True, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" @@ -880,11 +880,11 @@ def __call__( still retain a substantial amount of noise as determined by the discrete timesteps selected by the scheduler. The denoising_end parameter should ideally be utilized when this pipeline forms a part of a "Mixture of Denoisers" multi-pipeline setup, as elaborated in [**Refining the Image - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders negative_prompt_embeds (`torch.Tensor`, *optional*): @@ -919,16 +919,16 @@ def __call__( [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when using zero terminal SNR. - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(width, height)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - editing_prompt (`str` or `List[str]`, *optional*): + editing_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. The image is reconstructed by setting `editing_prompt = None`. Guidance direction of prompt should be specified via `reverse_editing_direction`. @@ -939,22 +939,22 @@ def __call__( Pre-generated pooled edit text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, editing_prompt_embeddings will be generated from `editing_prompt` input argument. - reverse_editing_direction (`bool` or `List[bool]`, *optional*, defaults to `False`): + reverse_editing_direction (`bool` or `list[bool]`, *optional*, defaults to `False`): Whether the corresponding prompt in `editing_prompt` should be increased or decreased. - edit_guidance_scale (`float` or `List[float]`, *optional*, defaults to 5): + edit_guidance_scale (`float` or `list[float]`, *optional*, defaults to 5): Guidance scale for guiding the image generation. If provided as list values should correspond to `editing_prompt`. `edit_guidance_scale` is defined as `s_e` of equation 12 of [LEDITS++ Paper](https://huggingface.co/papers/2301.12247). - edit_warmup_steps (`float` or `List[float]`, *optional*, defaults to 10): + edit_warmup_steps (`float` or `list[float]`, *optional*, defaults to 10): Number of diffusion steps (for each prompt) for which guidance is not applied. - edit_cooldown_steps (`float` or `List[float]`, *optional*, defaults to `None`): + edit_cooldown_steps (`float` or `list[float]`, *optional*, defaults to `None`): Number of diffusion steps (for each prompt) after which guidance is no longer applied. - edit_threshold (`float` or `List[float]`, *optional*, defaults to 0.9): + edit_threshold (`float` or `list[float]`, *optional*, defaults to 0.9): Masking threshold of guidance. Threshold should be proportional to the image region that is modified. 'edit_threshold' is defined as 'λ' of equation 12 of [LEDITS++ Paper](https://huggingface.co/papers/2301.12247). - sem_guidance (`List[torch.Tensor]`, *optional*): - List of pre-generated guidance vectors to be applied at generation. Length of the list has to + sem_guidance (`list[torch.Tensor]`, *optional*): + list of pre-generated guidance vectors to be applied at generation. Length of the list has to correspond to `num_inference_steps`. use_cross_attn_mask: Whether cross-attention masks are used. Cross-attention masks are always used when use_intersect_mask @@ -980,7 +980,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1482,14 +1482,14 @@ def invert( negative_prompt_2: str = None, num_inversion_steps: int = 50, skip: float = 0.15, - generator: Optional[torch.Generator] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), + generator: torch.Generator | None = None, + crops_coords_top_left: tuple[int, int] = (0, 0), num_zero_noise_steps: int = 3, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - resize_mode: Optional[str] = "default", - crops_coords: Optional[Tuple[int, int, int, int]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, + height: int | None = None, + width: int | None = None, + resize_mode: str | None = "default", + crops_coords: tuple[int, int, int, int] | None = None, ): r""" The function to the pipeline for image inversion as described by the [LEDITS++ @@ -1505,11 +1505,11 @@ def invert( if the `source_prompt` is `""`. source_guidance_scale (`float`, defaults to `3.5`): Strength of guidance during inversion. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders num_inversion_steps (`int`, defaults to `50`): @@ -1520,7 +1520,7 @@ def invert( generator (`torch.Generator`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make inversion deterministic. - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of diff --git a/src/diffusers/pipelines/ledits_pp/pipeline_output.py b/src/diffusers/pipelines/ledits_pp/pipeline_output.py index 756be82b0069..4a69e7442bdc 100644 --- a/src/diffusers/pipelines/ledits_pp/pipeline_output.py +++ b/src/diffusers/pipelines/ledits_pp/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Optional, Union import numpy as np import PIL.Image @@ -13,16 +12,16 @@ class LEditsPPDiffusionPipelineOutput(BaseOutput): Output class for LEdits++ Diffusion pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, num_channels)`. - nsfw_content_detected (`List[bool]`) - List indicating whether the corresponding generated image contains “not-safe-for-work” (nsfw) content or + nsfw_content_detected (`list[bool]`) + list indicating whether the corresponding generated image contains “not-safe-for-work” (nsfw) content or `None` if safety checking could not be performed. """ - images: Union[List[PIL.Image.Image], np.ndarray] - nsfw_content_detected: Optional[List[bool]] + images: list[PIL.Image.Image] | np.ndarray + nsfw_content_detected: list[bool] | None @dataclass @@ -31,13 +30,13 @@ class LEditsPPInversionPipelineOutput(BaseOutput): Output class for LEdits++ Diffusion pipelines. Args: - input_images (`List[PIL.Image.Image]` or `np.ndarray`) - List of the cropped and resized input images as PIL images of length `batch_size` or NumPy array of shape ` + input_images (`list[PIL.Image.Image]` or `np.ndarray`) + list of the cropped and resized input images as PIL images of length `batch_size` or NumPy array of shape ` (batch_size, height, width, num_channels)`. - vae_reconstruction_images (`List[PIL.Image.Image]` or `np.ndarray`) - List of VAE reconstruction of all input images as PIL images of length `batch_size` or NumPy array of shape + vae_reconstruction_images (`list[PIL.Image.Image]` or `np.ndarray`) + list of VAE reconstruction of all input images as PIL images of length `batch_size` or NumPy array of shape ` (batch_size, height, width, num_channels)`. """ - images: Union[List[PIL.Image.Image], np.ndarray] - vae_reconstruction_images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image] | np.ndarray + vae_reconstruction_images: list[PIL.Image.Image] | np.ndarray diff --git a/src/diffusers/pipelines/longcat_image/pipeline_longcat_image.py b/src/diffusers/pipelines/longcat_image/pipeline_longcat_image.py index a758d545fa4a..19720d7bbab8 100644 --- a/src/diffusers/pipelines/longcat_image/pipeline_longcat_image.py +++ b/src/diffusers/pipelines/longcat_image/pipeline_longcat_image.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect import re -from typing import Any, Dict, List, Optional, Union +from typing import Any import numpy as np import torch @@ -145,10 +145,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -163,15 +163,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -260,10 +260,10 @@ def rewire_prompt(self, prompt, device): text = self.text_processor.apply_chat_template(message, tokenize=False, add_generation_prompt=True) all_text.append(text) - inputs = self.text_processor(text=all_text, padding=True, return_tensors="pt").to(device) + inputs = self.text_processor(text=all_text, padding=True, return_tensors="pt").to(self.text_encoder.device) - self.text_encoder.to(device) generated_ids = self.text_encoder.generate(**inputs, max_new_tokens=self.tokenizer_max_length) + generated_ids.to(device) generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)] output_text = self.text_processor.batch_decode( generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False @@ -271,7 +271,7 @@ def rewire_prompt(self, prompt, device): rewrite_prompt = output_text return rewrite_prompt - def _encode_prompt(self, prompt: List[str]): + def _encode_prompt(self, prompt: list[str]): batch_all_tokens = [] for each_prompt in prompt: @@ -334,9 +334,9 @@ def _encode_prompt(self, prompt: List[str]): def encode_prompt( self, - prompt: Union[str, List[str]] = None, - num_images_per_prompt: Optional[int] = 1, - prompt_embeds: Optional[torch.Tensor] = None, + prompt: str | list[str] = None, + num_images_per_prompt: int | None = 1, + prompt_embeds: torch.Tensor | None = None, ): prompt = [prompt] if isinstance(prompt, str) else prompt batch_size = len(prompt) @@ -472,24 +472,24 @@ def check_inputs( @torch.no_grad() def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 4.5, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - enable_cfg_renorm: Optional[bool] = True, - cfg_renorm_min: Optional[float] = 0.0, - enable_prompt_rewrite: Optional[bool] = True, + joint_attention_kwargs: dict[str, Any] | None = None, + enable_cfg_renorm: bool | None = True, + cfg_renorm_min: float | None = 0.0, + enable_prompt_rewrite: bool | None = True, ): r""" Function invoked when calling the pipeline for generation. diff --git a/src/diffusers/pipelines/longcat_image/pipeline_longcat_image_edit.py b/src/diffusers/pipelines/longcat_image/pipeline_longcat_image_edit.py index e55a2a47f343..69d5d82f18ec 100644 --- a/src/diffusers/pipelines/longcat_image/pipeline_longcat_image_edit.py +++ b/src/diffusers/pipelines/longcat_image/pipeline_longcat_image_edit.py @@ -14,7 +14,7 @@ import inspect import math import re -from typing import Any, Dict, List, Optional, Union +from typing import Any import numpy as np import PIL @@ -144,10 +144,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -162,15 +162,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -203,7 +203,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -343,10 +343,10 @@ def _encode_prompt(self, prompt, image): def encode_prompt( self, - prompt: List[str] = None, - image: Optional[torch.Tensor] = None, - num_images_per_prompt: Optional[int] = 1, - prompt_embeds: Optional[torch.Tensor] = None, + prompt: list[str] = None, + image: torch.Tensor | None = None, + num_images_per_prompt: int | None = 1, + prompt_embeds: torch.Tensor | None = None, ): prompt = [prompt] if isinstance(prompt, str) else prompt batch_size = len(prompt) @@ -532,20 +532,20 @@ def check_inputs( @torch.no_grad() def __call__( self, - image: Optional[PIL.Image.Image] = None, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + image: PIL.Image.Image | None = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 4.5, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, + joint_attention_kwargs: dict[str, Any] | None = None, ): r""" Function invoked when calling the pipeline for generation. diff --git a/src/diffusers/pipelines/longcat_image/pipeline_output.py b/src/diffusers/pipelines/longcat_image/pipeline_output.py index e3c25f1cbfa7..290178eb1a88 100644 --- a/src/diffusers/pipelines/longcat_image/pipeline_output.py +++ b/src/diffusers/pipelines/longcat_image/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -13,9 +12,9 @@ class LongCatImagePipelineOutput(BaseOutput): Output class for Stable Diffusion pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) + images (`list[PIL.Image.Image]` or `np.ndarray`) List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image, np.ndarray] diff --git a/src/diffusers/pipelines/ltx/__init__.py b/src/diffusers/pipelines/ltx/__init__.py index 6001867916b3..05117d35d3b4 100644 --- a/src/diffusers/pipelines/ltx/__init__.py +++ b/src/diffusers/pipelines/ltx/__init__.py @@ -25,6 +25,7 @@ _import_structure["modeling_latent_upsampler"] = ["LTXLatentUpsamplerModel"] _import_structure["pipeline_ltx"] = ["LTXPipeline"] _import_structure["pipeline_ltx_condition"] = ["LTXConditionPipeline"] + _import_structure["pipeline_ltx_i2v_long_multi_prompt"] = ["LTXI2VLongMultiPromptPipeline"] _import_structure["pipeline_ltx_image2video"] = ["LTXImageToVideoPipeline"] _import_structure["pipeline_ltx_latent_upsample"] = ["LTXLatentUpsamplePipeline"] @@ -39,6 +40,7 @@ from .modeling_latent_upsampler import LTXLatentUpsamplerModel from .pipeline_ltx import LTXPipeline from .pipeline_ltx_condition import LTXConditionPipeline + from .pipeline_ltx_i2v_long_multi_prompt import LTXI2VLongMultiPromptPipeline from .pipeline_ltx_image2video import LTXImageToVideoPipeline from .pipeline_ltx_latent_upsample import LTXLatentUpsamplePipeline diff --git a/src/diffusers/pipelines/ltx/modeling_latent_upsampler.py b/src/diffusers/pipelines/ltx/modeling_latent_upsampler.py index 6dce792a2b43..f579cf00dbe7 100644 --- a/src/diffusers/pipelines/ltx/modeling_latent_upsampler.py +++ b/src/diffusers/pipelines/ltx/modeling_latent_upsampler.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional - import torch from ...configuration_utils import ConfigMixin, register_to_config @@ -21,7 +19,7 @@ class ResBlock(torch.nn.Module): - def __init__(self, channels: int, mid_channels: Optional[int] = None, dims: int = 3): + def __init__(self, channels: int, mid_channels: int | None = None, dims: int = 3): super().__init__() if mid_channels is None: mid_channels = channels diff --git a/src/diffusers/pipelines/ltx/pipeline_ltx.py b/src/diffusers/pipelines/ltx/pipeline_ltx.py index 8ca8b4419e18..e2514c3bca24 100644 --- a/src/diffusers/pipelines/ltx/pipeline_ltx.py +++ b/src/diffusers/pipelines/ltx/pipeline_ltx.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -83,10 +83,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -101,15 +101,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -233,11 +233,11 @@ def __init__( def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 128, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -282,25 +282,25 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.mochi.pipeline_mochi.MochiPipeline.encode_prompt with 256->128 def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, max_sequence_length: int = 128, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -478,10 +478,10 @@ def prepare_latents( height: int = 512, width: int = 704, num_frames: int = 161, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[torch.Generator] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: if latents is not None: return latents.to(device=device, dtype=dtype) @@ -536,37 +536,37 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] | None = None, height: int = 512, width: int = 704, num_frames: int = 161, frame_rate: int = 25, num_inference_steps: int = 50, - timesteps: List[int] = None, + timesteps: list[int] = None, guidance_scale: float = 3, guidance_rescale: float = 0.0, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - decode_timestep: Union[float, List[float]] = 0.0, - decode_noise_scale: Optional[Union[float, List[float]]] = None, - output_type: Optional[str] = "pil", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + decode_timestep: float | list[float] = 0.0, + decode_noise_scale: float | list[float] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 128, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. height (`int`, defaults to `512`): @@ -578,7 +578,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. @@ -596,7 +596,7 @@ def __call__( using zero terminal SNR. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of videos to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -631,7 +631,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -726,10 +726,14 @@ def __call__( self.scheduler.config.get("base_shift", 0.5), self.scheduler.config.get("max_shift", 1.15), ) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, - device, + timestep_device, timesteps, sigmas=sigmas, mu=mu, diff --git a/src/diffusers/pipelines/ltx/pipeline_ltx_condition.py b/src/diffusers/pipelines/ltx/pipeline_ltx_condition.py index 48a6f0837c8d..539a28f56e67 100644 --- a/src/diffusers/pipelines/ltx/pipeline_ltx_condition.py +++ b/src/diffusers/pipelines/ltx/pipeline_ltx_condition.py @@ -14,7 +14,7 @@ import inspect from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import PIL.Image import torch @@ -100,7 +100,7 @@ class LTXVideoCondition: Attributes: image (`PIL.Image.Image`): The image to condition the video on. - video (`List[PIL.Image.Image]`): + video (`list[PIL.Image.Image]`): The video to condition the video on. frame_index (`int`): The frame index at which the image or video will conditionally effect the video generation. @@ -108,8 +108,8 @@ class LTXVideoCondition: The strength of the conditioning effect. A value of `1.0` means the conditioning effect is fully applied. """ - image: Optional[PIL.Image.Image] = None - video: Optional[List[PIL.Image.Image]] = None + image: PIL.Image.Image | None = None + video: list[PIL.Image.Image] | None = None frame_index: int = 0 strength: float = 1.0 @@ -151,10 +151,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -169,15 +169,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -210,7 +210,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -319,11 +319,11 @@ def __init__( def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 256, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -368,25 +368,25 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.mochi.pipeline_mochi.MochiPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, max_sequence_length: int = 256, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -671,21 +671,21 @@ def add_noise_to_image_conditioning_latents( def prepare_latents( self, - conditions: Optional[List[torch.Tensor]] = None, - condition_strength: Optional[List[float]] = None, - condition_frame_index: Optional[List[int]] = None, + conditions: list[torch.Tensor] | None = None, + condition_strength: list[float] | None = None, + condition_frame_index: list[int] | None = None, batch_size: int = 1, num_channels_latents: int = 128, height: int = 512, width: int = 704, num_frames: int = 161, num_prefix_latent_frames: int = 2, - sigma: Optional[torch.Tensor] = None, - latents: Optional[torch.Tensor] = None, - generator: Optional[torch.Generator] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + sigma: torch.Tensor | None = None, + latents: torch.Tensor | None = None, + generator: torch.Generator | None = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: num_latent_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 latent_height = height // self.vae_spatial_compression_ratio latent_width = width // self.vae_spatial_compression_ratio @@ -849,61 +849,61 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - conditions: Union[LTXVideoCondition, List[LTXVideoCondition]] = None, - image: Union[PipelineImageInput, List[PipelineImageInput]] = None, - video: List[PipelineImageInput] = None, - frame_index: Union[int, List[int]] = 0, - strength: Union[float, List[float]] = 1.0, + conditions: LTXVideoCondition | list[LTXVideoCondition] = None, + image: PipelineImageInput | list[PipelineImageInput] = None, + video: list[PipelineImageInput] = None, + frame_index: int | list[int] = 0, + strength: float | list[float] = 1.0, denoise_strength: float = 1.0, - prompt: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] | None = None, height: int = 512, width: int = 704, num_frames: int = 161, frame_rate: int = 25, num_inference_steps: int = 50, - timesteps: List[int] = None, + timesteps: list[int] = None, guidance_scale: float = 3, guidance_rescale: float = 0.0, image_cond_noise_scale: float = 0.15, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - decode_timestep: Union[float, List[float]] = 0.0, - decode_noise_scale: Optional[Union[float, List[float]]] = None, - output_type: Optional[str] = "pil", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + decode_timestep: float | list[float] = 0.0, + decode_noise_scale: float | list[float] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 256, ): r""" Function invoked when calling the pipeline for generation. Args: - conditions (`List[LTXVideoCondition], *optional*`): + conditions (`list[LTXVideoCondition], *optional*`): The list of frame-conditioning items for the video generation.If not provided, conditions will be created using `image`, `video`, `frame_index` and `strength`. - image (`PipelineImageInput` or `List[PipelineImageInput]`, *optional*): + image (`PipelineImageInput` or `list[PipelineImageInput]`, *optional*): The image or images to condition the video generation. If not provided, one has to pass `video` or `conditions`. - video (`List[PipelineImageInput]`, *optional*): + video (`list[PipelineImageInput]`, *optional*): The video to condition the video generation. If not provided, one has to pass `image` or `conditions`. - frame_index (`int` or `List[int]`, *optional*): + frame_index (`int` or `list[int]`, *optional*): The frame index or frame indices at which the image or video will conditionally effect the video generation. If not provided, one has to pass `conditions`. - strength (`float` or `List[float]`, *optional*): + strength (`float` or `list[float]`, *optional*): The strength or strengths of the conditioning effect. If not provided, one has to pass `conditions`. denoise_strength (`float`, defaults to `1.0`): The strength of the noise added to the latents for editing. Higher strength leads to more noise added to the latents, therefore leading to more differences between original video and generated video. This is useful for video-to-video editing. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. height (`int`, defaults to `512`): @@ -915,7 +915,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. @@ -933,7 +933,7 @@ def __call__( using zero terminal SNR. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of videos to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -968,7 +968,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1102,11 +1102,24 @@ def __call__( latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 latent_height = height // self.vae_spatial_compression_ratio latent_width = width // self.vae_spatial_compression_ratio + if timesteps is None: sigmas = linear_quadratic_schedule(num_inference_steps) timesteps = sigmas * 1000 - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) + + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + timestep_device, + timesteps, + ) sigmas = self.scheduler.sigmas + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) latent_sigma = None diff --git a/src/diffusers/pipelines/ltx/pipeline_ltx_i2v_long_multi_prompt.py b/src/diffusers/pipelines/ltx/pipeline_ltx_i2v_long_multi_prompt.py new file mode 100644 index 000000000000..838d5afc5c5a --- /dev/null +++ b/src/diffusers/pipelines/ltx/pipeline_ltx_i2v_long_multi_prompt.py @@ -0,0 +1,1410 @@ +# Copyright 2025 Lightricks and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import copy +from typing import Any, Callable + +import numpy as np +import PIL +import torch +from transformers import T5EncoderModel, T5TokenizerFast + +from ...callbacks import MultiPipelineCallbacks, PipelineCallback +from ...loaders import FromSingleFileMixin, LTXVideoLoraLoaderMixin +from ...models.autoencoders import AutoencoderKLLTXVideo +from ...models.transformers import LTXVideoTransformer3DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler, LTXEulerAncestralRFScheduler +from ...utils import is_torch_xla_available, logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor +from ..pipeline_utils import DiffusionPipeline +from .pipeline_output import LTXPipelineOutput + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import LTXEulerAncestralRFScheduler, LTXI2VLongMultiPromptPipeline + + >>> pipe = LTXI2VLongMultiPromptPipeline.from_pretrained("LTX-Video-0.9.8-13B-distilled") + >>> # For ComfyUI parity, swap in the RF scheduler (keeps the original config). + >>> pipe.scheduler = LTXEulerAncestralRFScheduler.from_config(pipe.scheduler.config) + >>> pipe = pipe.to("cuda").to(dtype=torch.bfloat16) + >>> # Example A: get decoded frames (PIL) + >>> out = pipe( + ... prompt="a chimpanzee walks | a chimpanzee eats", + ... num_frames=161, + ... height=512, + ... width=704, + ... temporal_tile_size=80, + ... temporal_overlap=24, + ... output_type="pil", + ... return_dict=True, + ... ) + >>> frames = out.frames[0] # list of PIL.Image.Image + >>> # Example B: get latent video and decode later (saves VRAM during sampling) + >>> out_latent = pipe(prompt="a chimpanzee walking", output_type="latent", return_dict=True).frames + >>> frames = pipe.vae_decode_tiled(out_latent, output_type="pil")[0] + ``` +""" + + +def get_latent_coords( + latent_num_frames, latent_height, latent_width, batch_size, device, rope_interpolation_scale, latent_idx +): + """ + Compute latent patch top-left coordinates in (t, y, x) order. + + Args: + latent_num_frames: int. Number of latent frames (T_lat). + latent_height: int. Latent height (H_lat). + latent_width: int. Latent width (W_lat). + batch_size: int. Batch dimension (B). + device: torch.device for the resulting tensor. + rope_interpolation_scale: + tuple[int|float, int|float, int|float]. Scale per (t, y, x) latent step to pixel coords. + latent_idx: int | None. When not None, shifts the time coordinate to align segments: + - <= 0 uses step multiples of rope_interpolation_scale[0] + - > 0 starts at 1 then increments by rope_interpolation_scale[0] + + Returns: + Tensor of shape [B, 3, T_lat * H_lat * W_lat] containing top-left coordinates per latent patch, repeated for each + batch element. + """ + latent_sample_coords = torch.meshgrid( + torch.arange(0, latent_num_frames, 1, device=device), + torch.arange(0, latent_height, 1, device=device), + torch.arange(0, latent_width, 1, device=device), + indexing="ij", + ) + latent_sample_coords = torch.stack(latent_sample_coords, dim=0) + latent_coords = latent_sample_coords.unsqueeze(0).repeat(batch_size, 1, 1, 1, 1) + latent_coords = latent_coords.flatten(2) + pixel_coords = latent_coords * torch.tensor(rope_interpolation_scale, device=latent_coords.device)[None, :, None] + if latent_idx is not None: + if latent_idx <= 0: + frame_idx = latent_idx * rope_interpolation_scale[0] + else: + frame_idx = 1 + (latent_idx - 1) * rope_interpolation_scale[0] + if frame_idx == 0: + pixel_coords[:, 0] = (pixel_coords[:, 0] + 1 - rope_interpolation_scale[0]).clamp(min=0) + pixel_coords[:, 0] += frame_idx + return pixel_coords + + +# Copied from diffusers.pipelines.ltx.pipeline_ltx.rescale_noise_cfg +def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): + r""" + Rescales `noise_cfg` tensor based on `guidance_rescale` to improve image quality and fix overexposure. Based on + Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are + Flawed](https://huggingface.co/papers/2305.08891). + + Args: + noise_cfg (`torch.Tensor`): + The predicted noise tensor for the guided diffusion process. + noise_pred_text (`torch.Tensor`): + The predicted noise tensor for the text-guided diffusion process. + guidance_rescale (`float`, *optional*, defaults to 0.0): + A rescale factor applied to the noise predictions. + + Returns: + noise_cfg (`torch.Tensor`): The rescaled noise prediction tensor. + """ + std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True) + std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) + # rescale the results from guidance (fixes overexposure) + noise_pred_rescaled = noise_cfg * (std_text / std_cfg) + # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images + noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg + return noise_cfg + + +def adain_normalize_latents( + curr_latents: torch.Tensor, ref_latents: torch.Tensor | None, factor: float +) -> torch.Tensor: + """ + Optional AdaIN normalization: channel-wise mean/variance matching of curr_latents to ref_latents, controlled by + factor. + + Args: + curr_latents: Tensor [B, C, T, H, W]. Current window latents. + ref_latents: + Tensor | None [B, C, T_ref, H, W]. Reference latents (e.g., first window) used to compute target stats. + factor: float in [0, 1]. 0 keeps current stats; 1 matches reference stats. + + Returns: + Tensor with per-channel mean/std blended towards the reference. + """ + if ref_latents is None or factor is None or factor <= 0: + return curr_latents + + eps = torch.tensor(1e-6, device=curr_latents.device, dtype=curr_latents.dtype) + + # Compute per-channel means/stds for current and reference over (T, H, W) + mu_curr = curr_latents.mean(dim=(2, 3, 4), keepdim=True) + sigma_curr = curr_latents.std(dim=(2, 3, 4), keepdim=True) + + mu_ref = ref_latents.mean(dim=(2, 3, 4), keepdim=True).to(device=curr_latents.device, dtype=curr_latents.dtype) + sigma_ref = ref_latents.std(dim=(2, 3, 4), keepdim=True).to(device=curr_latents.device, dtype=curr_latents.dtype) + + # Blend target statistics + mu_blend = (1.0 - float(factor)) * mu_curr + float(factor) * mu_ref + sigma_blend = (1.0 - float(factor)) * sigma_curr + float(factor) * sigma_ref + sigma_blend = torch.clamp(sigma_blend, min=float(eps)) + + # Apply AdaIN + curr_norm = (curr_latents - mu_curr) / (sigma_curr + eps) + return curr_norm * sigma_blend + mu_blend + + +def split_into_temporal_windows( + latent_len: int, temporal_tile_size: int, temporal_overlap: int, compression: int +) -> list[tuple[int, int]]: + """ + Split latent frames into sliding windows. + + Args: + latent_len: int. Number of latent frames (T_lat). + temporal_tile_size: int. Window size in latent frames (> 0). + temporal_overlap: int. Overlap between windows in latent frames (>= 0). + compression: int. VAE temporal compression ratio (unused here; kept for parity). + + Returns: + list[tuple[int, int]]: inclusive-exclusive (start, end) indices per window. + """ + if temporal_tile_size <= 0: + raise ValueError("temporal_tile_size must be > 0") + stride = max(temporal_tile_size - temporal_overlap, 1) + windows = [] + start = 0 + while start < latent_len: + end = min(start + temporal_tile_size, latent_len) + windows.append((start, end)) + if end == latent_len: + break + start = start + stride + return windows + + +def linear_overlap_fuse(prev: torch.Tensor, new: torch.Tensor, overlap: int) -> torch.Tensor: + """ + Temporal linear crossfade between two latent clips over the overlap region. + + Args: + prev: Tensor [B, C, F, H, W]. Previous output segment. + new: Tensor [B, C, F, H, W]. New segment to be appended. + overlap: int. Number of frames to crossfade (overlap <= 1 concatenates without blend). + + Returns: + Tensor [B, C, F_prev + F_new - overlap, H, W] after crossfade at the seam. + """ + if overlap <= 1: + return torch.cat([prev, new], dim=2) + alpha = torch.linspace(1, 0, overlap + 2, device=prev.device, dtype=prev.dtype)[1:-1] + shape = [1] * prev.ndim + shape[2] = alpha.size(0) + alpha = alpha.reshape(shape) + blended = alpha * prev[:, :, -overlap:] + (1 - alpha) * new[:, :, :overlap] + return torch.cat([prev[:, :, :-overlap], blended, new[:, :, overlap:]], dim=2) + + +def inject_prev_tail_latents( + window_latents: torch.Tensor, + prev_tail_latents: torch.Tensor | None, + window_cond_mask_5d: torch.Tensor, + overlap_lat: int, + strength: float | None, + prev_overlap_len: int, +) -> tuple[torch.Tensor, torch.Tensor, int]: + """ + Inject the tail latents from the previous window at the beginning of the current window (first k frames), where k = + min(overlap_lat, T_curr, T_prev_tail). + + Args: + window_latents: Tensor [B, C, T, H, W]. Current window latents. + prev_tail_latents: Tensor | None [B, C, T_prev, H, W]. Tail segment from the previous window. + window_cond_mask_5d: Tensor [B, 1, T, H, W]. Per-token conditioning mask (1 = free, 0 = hard condition). + overlap_lat: int. Number of latent frames to inject from the previous tail. + strength: float | None in [0, 1]. Blend strength; 1.0 replaces, 0.0 keeps original. + prev_overlap_len: int. Accumulated overlap length so far (used for trimming later). + + Returns: + tuple[Tensor, Tensor, int]: (updated_window_latents, updated_cond_mask, updated_prev_overlap_len) + """ + if prev_tail_latents is None or overlap_lat <= 0 or strength is None or strength <= 0: + return window_latents, window_cond_mask_5d, prev_overlap_len + + # Expected shape: [B, C, T, H, W] + T = int(window_latents.shape[2]) + k = min(int(overlap_lat), T, int(prev_tail_latents.shape[2])) + if k <= 0: + return window_latents, window_cond_mask_5d, prev_overlap_len + + tail = prev_tail_latents[:, :, -k:] + mask = torch.full( + (window_cond_mask_5d.shape[0], 1, tail.shape[2], window_cond_mask_5d.shape[3], window_cond_mask_5d.shape[4]), + 1.0 - strength, + dtype=window_cond_mask_5d.dtype, + device=window_cond_mask_5d.device, + ) + + window_latents = torch.cat([window_latents, tail], dim=2) + window_cond_mask_5d = torch.cat([window_cond_mask_5d, mask], dim=2) + return window_latents, window_cond_mask_5d, prev_overlap_len + k + + +def build_video_coords_for_window( + latents: torch.Tensor, + overlap_len: int, + guiding_len: int, + negative_len: int, + rope_interpolation_scale: torch.Tensor, + frame_rate: int, +) -> torch.Tensor: + """ + Build video_coords: [B, 3, S] with order [t, y, x]. + + Args: + latents: Tensor [B, C, T, H, W]. Current window latents (before any trimming). + overlap_len: int. Number of frames from previous tail injected at the head. + guiding_len: int. Number of guidance frames appended at the head. + negative_len: int. Number of negative-index frames appended at the head (typically 1 or 0). + rope_interpolation_scale: tuple[int|float, int|float, int|float]. Scale for (t, y, x). + frame_rate: int. Used to convert time indices into seconds (t /= frame_rate). + + Returns: + Tensor [B, 3, T*H*W] of fractional pixel coordinates per latent patch. + """ + + b, c, f, h, w = latents.shape + pixel_coords = get_latent_coords(f, h, w, b, latents.device, rope_interpolation_scale, 0) + replace_corrds = [] + if overlap_len > 0: + replace_corrds.append(get_latent_coords(overlap_len, h, w, b, latents.device, rope_interpolation_scale, 0)) + if guiding_len > 0: + replace_corrds.append( + get_latent_coords(guiding_len, h, w, b, latents.device, rope_interpolation_scale, overlap_len) + ) + if negative_len > 0: + replace_corrds.append(get_latent_coords(negative_len, h, w, b, latents.device, rope_interpolation_scale, -1)) + if len(replace_corrds) > 0: + replace_corrds = torch.cat(replace_corrds, axis=2) + pixel_coords[:, :, -replace_corrds.shape[2] :] = replace_corrds + fractional_coords = pixel_coords.to(torch.float32) + fractional_coords[:, 0] = fractional_coords[:, 0] * (1.0 / frame_rate) + return fractional_coords + + +def parse_prompt_segments(prompt: str | list[str], prompt_segments: list[dict[str, Any]] | None) -> list[str]: + """ + Return a list of positive prompts per window index. + + Args: + prompt: str | list[str]. If str contains '|', parts are split by bars and trimmed. + prompt_segments: + list[dict], optional. Each dict with {"start_window", "end_window", "text"} overrides prompts per window. + + Returns: + list[str] containing the positive prompt for each window index. + """ + if prompt is None: + return [] + if prompt_segments: + max_w = 0 + for seg in prompt_segments: + max_w = max(max_w, int(seg.get("end_window", 0))) + texts = [""] * (max_w + 1) + for seg in prompt_segments: + s = int(seg.get("start_window", 0)) + e = int(seg.get("end_window", s)) + txt = seg.get("text", "") + for w in range(s, e + 1): + texts[w] = txt + # fill empty by last non-empty + last = "" + for i in range(len(texts)): + if texts[i] == "": + texts[i] = last + else: + last = texts[i] + return texts + + # bar-split mode + if isinstance(prompt, str): + parts = [p.strip() for p in prompt.split("|")] + else: + parts = prompt + parts = [p for p in parts if p is not None] + return parts + + +def batch_normalize(latents, reference, factor): + """ + Batch AdaIN-like normalization for latents in dict format (ComfyUI-compatible). + + Args: + latents: dict containing "samples" shaped [B, C, F, H, W] + reference: dict containing "samples" used to compute target stats + factor: float in [0, 1]; 0 = no change, 1 = full match to reference + Returns: + tuple[dict]: a single-element tuple with the updated latents dict. + """ + latents_copy = copy.deepcopy(latents) + t = latents_copy["samples"] # B x C x F x H x W + + for i in range(t.size(0)): # batch + for c in range(t.size(1)): # channel + r_sd, r_mean = torch.std_mean(reference["samples"][i, c], dim=None) # index by original dim order + i_sd, i_mean = torch.std_mean(t[i, c], dim=None) + + t[i, c] = ((t[i, c] - i_mean) / i_sd) * r_sd + r_mean + + latents_copy["samples"] = torch.lerp(latents["samples"], t, factor) + return (latents_copy,) + + +class LTXI2VLongMultiPromptPipeline(DiffusionPipeline, FromSingleFileMixin, LTXVideoLoraLoaderMixin): + r""" + Long-duration I2V (image-to-video) multi-prompt pipeline with ComfyUI parity. + + Key features: + - Temporal sliding-window sampling only (no spatial H/W sharding); autoregressive fusion across windows. + - Multi-prompt segmentation per window with smooth transitions at window heads. + - First-frame hard conditioning via per-token mask for I2V. + - VRAM control via temporal windowing and VAE tiled decoding. + + Reference: https://github.com/Lightricks/LTX-Video + + Args: + transformer ([`LTXVideoTransformer3DModel`]): + Conditional Transformer architecture to denoise the encoded video latents. + scheduler ([`FlowMatchEulerDiscreteScheduler`] or [`LTXEulerAncestralRFScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + vae ([`AutoencoderKLLTXVideo`]): + Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. + text_encoder ([`T5EncoderModel`]): + [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically + the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant. + tokenizer (`T5TokenizerFast`): + Tokenizer of class + [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast). + """ + + model_cpu_offload_seq = "text_encoder->transformer->vae" + _optional_components = [] + _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] + + def __init__( + self, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKLLTXVideo, + text_encoder: T5EncoderModel, + tokenizer: T5TokenizerFast, + transformer: LTXVideoTransformer3DModel, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + scheduler=scheduler, + ) + if not isinstance(scheduler, LTXEulerAncestralRFScheduler): + logger.warning( + "For ComfyUI parity, `LTXI2VLongMultiPromptPipeline` is typically run with " + "`LTXEulerAncestralRFScheduler`. Got %s.", + scheduler.__class__.__name__, + ) + + self.vae_spatial_compression_ratio = ( + self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 32 + ) + self.vae_temporal_compression_ratio = ( + self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 8 + ) + self.transformer_spatial_patch_size = ( + self.transformer.config.patch_size if getattr(self, "transformer", None) is not None else 1 + ) + self.transformer_temporal_patch_size = ( + self.transformer.config.patch_size_t if getattr(self, "transformer", None) is not None else 1 + ) + + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio) + self.tokenizer_max_length = ( + self.tokenizer.model_max_length if getattr(self, "tokenizer", None) is not None else 128 + ) + + self.default_height = 512 + self.default_width = 704 + self.default_frames = 121 + self._current_tile_T = None + + @property + # Copied from diffusers.pipelines.ltx.pipeline_ltx.LTXPipeline.guidance_scale + def guidance_scale(self): + return self._guidance_scale + + @property + # Copied from diffusers.pipelines.ltx.pipeline_ltx.LTXPipeline.guidance_rescale + def guidance_rescale(self): + return self._guidance_rescale + + @property + # Copied from diffusers.pipelines.ltx.pipeline_ltx.LTXPipeline.do_classifier_free_guidance + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + + @property + # Copied from diffusers.pipelines.ltx.pipeline_ltx.LTXPipeline.num_timesteps + def num_timesteps(self): + return self._num_timesteps + + @property + # Copied from diffusers.pipelines.ltx.pipeline_ltx.LTXPipeline.current_timestep + def current_timestep(self): + return self._current_timestep + + @property + # Copied from diffusers.pipelines.ltx.pipeline_ltx.LTXPipeline.attention_kwargs + def attention_kwargs(self): + return self._attention_kwargs + + @property + # Copied from diffusers.pipelines.ltx.pipeline_ltx.LTXPipeline.interrupt + def interrupt(self): + return self._interrupt + + # Copied from diffusers.pipelines.ltx.pipeline_ltx.LTXPipeline._get_t5_prompt_embeds + def _get_t5_prompt_embeds( + self, + prompt: str | list[str] = None, + num_videos_per_prompt: int = 1, + max_sequence_length: int = 128, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + prompt_attention_mask = text_inputs.attention_mask + prompt_attention_mask = prompt_attention_mask.bool().to(device) + + untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): + removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1]) + logger.warning( + "The following part of your input was truncated because `max_sequence_length` is set to " + f" {max_sequence_length} tokens: {removed_text}" + ) + + prompt_embeds = self.text_encoder(text_input_ids.to(device))[0] + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + + # duplicate text embeddings for each generation per prompt, using mps friendly method + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) + + prompt_attention_mask = prompt_attention_mask.view(batch_size, -1) + prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1) + + return prompt_embeds, prompt_attention_mask + + # Copied from diffusers.pipelines.ltx.pipeline_ltx.LTXPipeline.encode_prompt + def encode_prompt( + self, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, + do_classifier_free_guidance: bool = True, + num_videos_per_prompt: int = 1, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + max_sequence_length: int = 128, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `list[str]`, *optional*): + prompt to be encoded + negative_prompt (`str` or `list[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): + Whether to use classifier free guidance or not. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos that should be generated per prompt. torch device to place the resulting embeddings on + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + device: (`torch.device`, *optional*): + torch device + dtype: (`torch.dtype`, *optional*): + torch dtype + """ + device = device or self._execution_device + + prompt = [prompt] if isinstance(prompt, str) else prompt + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds, prompt_attention_mask = self._get_t5_prompt_embeds( + prompt=prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + + if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + + if prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + + negative_prompt_embeds, negative_prompt_attention_mask = self._get_t5_prompt_embeds( + prompt=negative_prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + + return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask + + @staticmethod + # Copied from diffusers.pipelines.ltx.pipeline_ltx.LTXPipeline._pack_latents + def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor: + # Unpacked latents of shape are [B, C, F, H, W] are patched into tokens of shape [B, C, F // p_t, p_t, H // p, p, W // p, p]. + # The patch dimensions are then permuted and collapsed into the channel dimension of shape: + # [B, F // p_t * H // p * W // p, C * p_t * p * p] (an ndim=3 tensor). + # dim=0 is the batch size, dim=1 is the effective video sequence length, dim=2 is the effective number of input features + batch_size, num_channels, num_frames, height, width = latents.shape + post_patch_num_frames = num_frames // patch_size_t + post_patch_height = height // patch_size + post_patch_width = width // patch_size + latents = latents.reshape( + batch_size, + -1, + post_patch_num_frames, + patch_size_t, + post_patch_height, + patch_size, + post_patch_width, + patch_size, + ) + latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3) + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx.pipeline_ltx.LTXPipeline._unpack_latents + def _unpack_latents( + latents: torch.Tensor, num_frames: int, height: int, width: int, patch_size: int = 1, patch_size_t: int = 1 + ) -> torch.Tensor: + # Packed latents of shape [B, S, D] (S is the effective video sequence length, D is the effective feature dimensions) + # are unpacked and reshaped into a video tensor of shape [B, C, F, H, W]. This is the inverse operation of + # what happens in the `_pack_latents` method. + batch_size = latents.size(0) + latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size) + latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3) + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx.pipeline_ltx.LTXPipeline._normalize_latents + def _normalize_latents( + latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 + ) -> torch.Tensor: + # Normalize latents across the channel dimension [B, C, F, H, W] + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents = (latents - latents_mean) * scaling_factor / latents_std + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx.pipeline_ltx.LTXPipeline._denormalize_latents + def _denormalize_latents( + latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 + ) -> torch.Tensor: + # Denormalize latents across the channel dimension [B, C, F, H, W] + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents = latents * latents_std / scaling_factor + latents_mean + return latents + + def prepare_latents( + self, + batch_size: int, + num_channels_latents: int, + height: int, + width: int, + num_frames: int, + device: torch.device, + generator: torch.Generator | None, + dtype: torch.dtype = torch.float32, + latents: torch.Tensor | None = None, + cond_latents: torch.Tensor | None = None, + cond_strength: float = 0.0, + negative_index_latents: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None, int, int, int]: + """ + Prepare base latents and optionally inject first-frame conditioning latents. + + Returns: + latents, negative_index_latents, latent_num_frames, latent_height, latent_width + """ + if latents is None: + latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 + latent_height = height // self.vae_spatial_compression_ratio + latent_width = width // self.vae_spatial_compression_ratio + latents = torch.zeros( + (batch_size, num_channels_latents, latent_num_frames, latent_height, latent_width), + device=device, + dtype=dtype, + ) + else: + latent_num_frames = latents.shape[2] + latent_height = latents.shape[3] + latent_width = latents.shape[4] + latents = latents.to(device=device, dtype=dtype) + + if cond_latents is not None and cond_strength > 0: + if negative_index_latents is None: + negative_index_latents = cond_latents + latents[:, :, :1, :, :] = cond_latents + + return latents, negative_index_latents, latent_num_frames, latent_height, latent_width + + # TODO: refactor this out + @torch.no_grad() + def vae_decode_tiled( + self, + latents: torch.Tensor, + decode_timestep: float | None = None, + decode_noise_scale: float | None = None, + horizontal_tiles: int = 4, + vertical_tiles: int = 4, + overlap: int = 3, + last_frame_fix: bool = True, + generator: torch.Generator | None = None, + output_type: str = "pt", + auto_denormalize: bool = True, + compute_dtype: torch.dtype = torch.float32, + enable_vae_tiling: bool = False, + ) -> torch.Tensor | np.ndarray | list[PIL.Image.Image]: + """ + VAE-based spatial tiled decoding (ComfyUI parity) implemented in Diffusers style. + - Linearly feather and blend overlapping tiles to avoid seams. + - Optional last_frame_fix: duplicate the last latent frame before decoding, then drop time_scale_factor frames + at the end. + - Supports timestep_conditioning and decode_noise_scale injection. + - By default, "normalized latents" (the denoising output) are de-normalized internally (auto_denormalize=True). + - Tile fusion is computed in compute_dtype (float32 by default) to reduce blur and color shifts. + + Args: + latents: [B, C_latent, F_latent, H_latent, W_latent] + decode_timestep: Optional decode timestep (effective only if VAE supports timestep_conditioning) + decode_noise_scale: + Optional decode noise interpolation (effective only if VAE supports timestep_conditioning) + horizontal_tiles, vertical_tiles: Number of tiles horizontally/vertically (>= 1) + overlap: Overlap in latent space (in latent pixels, >= 0) + last_frame_fix: Whether to enable the "repeat last frame" fix + generator: Random generator (used for decode_noise_scale noise) + output_type: "latent" | "pt" | "np" | "pil" + - "latent": return latents unchanged (useful for downstream processing) + - "pt": return tensor in VAE output space + - "np"/"pil": post-processed outputs via VideoProcessor.postprocess_video + auto_denormalize: If True, apply LTX de-normalization to `latents` internally (recommended) + compute_dtype: Precision used during tile fusion (float32 default; significantly reduces seam blur) + enable_vae_tiling: If True, delegate tiling to VAE's built-in `tiled_decode` (sets `vae.use_tiling`). + + Returns: + - If output_type="latent": returns input `latents` unchanged + - If output_type="pt": returns [B, C, F, H, W] (values roughly in [-1, 1]) + - If output_type="np"/"pil": returns post-processed outputs via postprocess_video + """ + if output_type == "latent": + return latents + if horizontal_tiles < 1 or vertical_tiles < 1: + raise ValueError("horizontal_tiles and vertical_tiles must be >= 1") + overlap = max(int(overlap), 0) + + # Device and precision + device = self._execution_device + latents = latents.to(device=device, dtype=compute_dtype) + + # De-normalize to VAE space (avoid color artifacts) + if auto_denormalize: + latents = self._denormalize_latents( + latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor + ) + # dtype required for VAE forward pass + latents = latents.to(dtype=self.vae.dtype) + + # Temporal/spatial upscaling ratios (parity with ComfyUI's downscale_index_formula) + tsf = int(self.vae_temporal_compression_ratio) + sf = int(self.vae_spatial_compression_ratio) + + # Optional: last_frame_fix (repeat last latent frame) + if last_frame_fix: + latents = torch.cat([latents, latents[:, :, -1:].contiguous()], dim=2) + + b, c_lat, f_lat, h_lat, w_lat = latents.shape + f_out = 1 + (f_lat - 1) * tsf + h_out = h_lat * sf + w_out = w_lat * sf + + # timestep_conditioning + decode-time noise injection (aligned with pipeline) + if getattr(self.vae.config, "timestep_conditioning", False): + dt = float(decode_timestep) if decode_timestep is not None else 0.0 + vt = torch.tensor([dt], device=device, dtype=latents.dtype) + if decode_noise_scale is not None: + dns = torch.tensor([float(decode_noise_scale)], device=device, dtype=latents.dtype)[ + :, None, None, None, None + ] + noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype) + latents = (1 - dns) * latents + dns * noise + else: + vt = None + + if enable_vae_tiling and hasattr(self.vae, "enable_tiling"): + self.vae.enable_tiling() + decoded = self.vae.decode(latents, vt, return_dict=False)[0] + if last_frame_fix: + decoded = decoded[:, :, :-tsf, :, :] + if output_type in ("np", "pil"): + return self.video_processor.postprocess_video(decoded, output_type=output_type) + return decoded + + # Compute base tile sizes (in latent space) + base_tile_h = (h_lat + (vertical_tiles - 1) * overlap) // vertical_tiles + base_tile_w = (w_lat + (horizontal_tiles - 1) * overlap) // horizontal_tiles + + output: torch.Tensor | None = None # [B, C_img, F, H, W], fused using compute_dtype + weights: torch.Tensor | None = None # [B, 1, F, H, W], fused using compute_dtype + + # Iterate tiles in latent space (no temporal tiling) + for v in range(vertical_tiles): + for h in range(horizontal_tiles): + h_start = h * (base_tile_w - overlap) + v_start = v * (base_tile_h - overlap) + + h_end = min(h_start + base_tile_w, w_lat) if h < horizontal_tiles - 1 else w_lat + v_end = min(v_start + base_tile_h, h_lat) if v < vertical_tiles - 1 else h_lat + + # Slice latent tile and decode + tile_latents = latents[:, :, :, v_start:v_end, h_start:h_end] + decoded_tile = self.vae.decode(tile_latents, vt, return_dict=False)[0] # [B, C, F, Ht, Wt] + # Cast to high precision to reduce blending blur + decoded_tile = decoded_tile.to(dtype=compute_dtype) + + # Initialize output buffers (compute_dtype) + if output is None: + output = torch.zeros( + (b, decoded_tile.shape[1], f_out, h_out, w_out), + device=decoded_tile.device, + dtype=compute_dtype, + ) + weights = torch.zeros( + (b, 1, f_out, h_out, w_out), + device=decoded_tile.device, + dtype=compute_dtype, + ) + + # Tile placement in output pixel space + out_h_start = v_start * sf + out_h_end = v_end * sf + out_w_start = h_start * sf + out_w_end = h_end * sf + + tile_out_h = out_h_end - out_h_start + tile_out_w = out_w_end - out_w_start + + # Linear feathering weights [B, 1, F, Ht, Wt] (compute_dtype) + tile_weights = torch.ones( + (b, 1, decoded_tile.shape[2], tile_out_h, tile_out_w), + device=decoded_tile.device, + dtype=compute_dtype, + ) + + overlap_out_h = overlap * sf + overlap_out_w = overlap * sf + + # Horizontal feathering: left/right overlaps + if overlap_out_w > 0: + if h > 0: + h_blend = torch.linspace( + 0, 1, steps=overlap_out_w, device=decoded_tile.device, dtype=compute_dtype + ) + tile_weights[:, :, :, :, :overlap_out_w] *= h_blend.view(1, 1, 1, 1, -1) + if h < horizontal_tiles - 1: + h_blend = torch.linspace( + 1, 0, steps=overlap_out_w, device=decoded_tile.device, dtype=compute_dtype + ) + tile_weights[:, :, :, :, -overlap_out_w:] *= h_blend.view(1, 1, 1, 1, -1) + + # Vertical feathering: top/bottom overlaps + if overlap_out_h > 0: + if v > 0: + v_blend = torch.linspace( + 0, 1, steps=overlap_out_h, device=decoded_tile.device, dtype=compute_dtype + ) + tile_weights[:, :, :, :overlap_out_h, :] *= v_blend.view(1, 1, 1, -1, 1) + if v < vertical_tiles - 1: + v_blend = torch.linspace( + 1, 0, steps=overlap_out_h, device=decoded_tile.device, dtype=compute_dtype + ) + tile_weights[:, :, :, -overlap_out_h:, :] *= v_blend.view(1, 1, 1, -1, 1) + + # Accumulate blended tile + output[:, :, :, out_h_start:out_h_end, out_w_start:out_w_end] += decoded_tile * tile_weights + weights[:, :, :, out_h_start:out_h_end, out_w_start:out_w_end] += tile_weights + + # Normalize, then clamp to [-1, 1] in compute_dtype to avoid color artifacts + output = output / (weights + 1e-8) + output = output.clamp(-1.0, 1.0) + output = output.to(dtype=self.vae.dtype) + + # Optional: drop the last tsf frames after last_frame_fix + if last_frame_fix: + output = output[:, :, :-tsf, :, :] + + if output_type in ("np", "pil"): + return self.video_processor.postprocess_video(output, output_type=output_type) + return output + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: str | list[str] = None, + negative_prompt: str | list[str] | None = None, + prompt_segments: list[dict[str, Any]] | None = None, + height: int = 512, + width: int = 704, + num_frames: int = 161, + frame_rate: float = 25, + guidance_scale: float = 1.0, + guidance_rescale: float = 0.0, + num_inference_steps: int | None = 8, + sigmas: list[float, torch.Tensor] | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + seed: int | None = 0, + cond_image: "PIL.Image.Image" | torch.Tensor | None = None, + cond_strength: float = 0.5, + latents: torch.Tensor | None = None, + temporal_tile_size: int = 80, + temporal_overlap: int = 24, + temporal_overlap_cond_strength: float = 0.5, + adain_factor: float = 0.25, + guidance_latents: torch.Tensor | None = None, + guiding_strength: float = 1.0, + negative_index_latents: torch.Tensor | None = None, + negative_index_strength: float = 1.0, + skip_steps_sigma_threshold: float | None = 1, + decode_timestep: float | None = 0.05, + decode_noise_scale: float | None = 0.025, + decode_horizontal_tiles: int = 4, + decode_vertical_tiles: int = 4, + decode_overlap: int = 3, + output_type: str | None = "latent", # "latent" | "pt" | "np" | "pil" + return_dict: bool = True, + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + max_sequence_length: int = 128, + ): + r""" + Generate an image-to-video sequence via temporal sliding windows and multi-prompt scheduling. + + Args: + prompt (`str` or `list[str]`, *optional*): + Positive text prompt(s) per window. If a single string contains '|', parts are split by bars. + negative_prompt (`str` or `list[str]`, *optional*): + Negative prompt(s) to suppress undesired content. + prompt_segments (`list[dict]`, *optional*): + Segment mapping with {"start_window", "end_window", "text"} to override prompts per window. + height (`int`, defaults to `512`): + Output image height in pixels; must be divisible by 32. + width (`int`, defaults to `704`): + Output image width in pixels; must be divisible by 32. + num_frames (`int`, defaults to `161`): + Number of output frames (in decoded pixel space). + frame_rate (`float`, defaults to `25`): + Frames-per-second; used to normalize temporal coordinates in `video_coords`. + guidance_scale (`float`, defaults to `1.0`): + CFG scale; values > 1 enable classifier-free guidance. + guidance_rescale (`float`, defaults to `0.0`): + Optional rescale to mitigate overexposure under CFG (see `rescale_noise_cfg`). + num_inference_steps (`int`, *optional*, defaults to `8`): + Denoising steps per window. Ignored if `sigmas` is provided. + sigmas (`list[float]` or `torch.Tensor`, *optional*): + Explicit sigma schedule per window; if set, overrides `num_inference_steps`. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + Controls stochasticity; list accepted but first element is used (batch=1). + seed (`int`, *optional*, defaults to `0`): + If provided, seeds the shared generator for global latents and derives a window-local generator with + `seed + w_start` per temporal window. + cond_image (`PIL.Image.Image` or `torch.Tensor`, *optional*): + Conditioning image; fixes frame 0 via per-token mask when `cond_strength > 0`. + cond_strength (`float`, defaults to `0.5`): + Strength of first-frame hard conditioning (smaller cond_mask ⇒ stronger preservation). + latents (`torch.Tensor`, *optional*): + Initial latents [B, C_lat, F_lat, H_lat, W_lat]; if None, sampled with `randn_tensor`. + temporal_tile_size (`int`, defaults to `80`): + Temporal window size (in decoded frames); internally scaled by VAE temporal compression. + temporal_overlap (`int`, defaults to `24`): + Overlap between consecutive windows (in decoded frames); internally scaled by compression. + temporal_overlap_cond_strength (`float`, defaults to `0.5`): + Strength for injecting previous window tail latents at new window head. + adain_factor (`float`, defaults to `0.25`): + AdaIN normalization strength for cross-window consistency (0 disables). + guidance_latents (`torch.Tensor`, *optional*): + Reference latents injected at window head; length trimmed by overlap for subsequent windows. + guiding_strength (`float`, defaults to `1.0`): + Injection strength for `guidance_latents`. + negative_index_latents (`torch.Tensor`, *optional*): + A single-frame latent appended at window head for "negative index" semantics. + negative_index_strength (`float`, defaults to `1.0`): + Injection strength for `negative_index_latents`. + skip_steps_sigma_threshold (`float`, *optional*, defaults to `1`): + Skip steps whose sigma exceeds this threshold. + decode_timestep (`float`, *optional*, defaults to `0.05`): + Decode-time timestep (if VAE supports timestep_conditioning). + decode_noise_scale (`float`, *optional*, defaults to `0.025`): + Decode-time noise mix scale (if VAE supports timestep_conditioning). + decode_horizontal_tiles (`int`, defaults to `4`): + Number of horizontal tiles during VAE decoding. + decode_vertical_tiles (`int`, defaults to `4`): + Number of vertical tiles during VAE decoding. + decode_overlap (`int`, defaults to `3`): + Overlap (in latent pixels) between tiles during VAE decoding. + output_type (`str`, *optional*, defaults to `"latent"`): + The output format of the generated video. Choose between "latent", "pt", "np", or "pil". If "latent", + returns latents without decoding. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.ltx.LTXPipelineOutput`] instead of a plain tuple. + attention_kwargs (`dict`, *optional*): + Extra attention parameters forwarded to the transformer. + callback_on_step_end (`PipelineCallback` or `MultiPipelineCallbacks`, *optional*): + Per-step callback hook. + callback_on_step_end_tensor_inputs (`list[str]`, defaults to `["latents"]`): + Keys from locals() to pass into the callback. + max_sequence_length (`int`, defaults to `128`): + Tokenizer max length for prompt encoding. + + Examples: + + Returns: + [`~pipelines.ltx.LTXPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`~pipelines.ltx.LTXPipelineOutput`] is returned, otherwise a `tuple` is + returned where the first element is a list with the generated frames. The output format depends on + `output_type`: + - "latent"/"pt": `torch.Tensor` [B, C, F, H, W]; "latent" is in normalized latent space, "pt" is VAE + output space. + - "np": `np.ndarray` post-processed. + - "pil": `list[PIL.Image.Image]` list of PIL images. + + Shapes: + Latent sizes (when auto-generated): + - F_lat = (num_frames - 1) // vae_temporal_compression_ratio + 1 + - H_lat = height // vae_spatial_compression_ratio + - W_lat = width // vae_spatial_compression_ratio + + Notes: + - Seeding: when `seed` is provided, each temporal window uses a local generator seeded with `seed + + w_start`, while the shared generator is seeded once for global latents if no generator is passed; + otherwise the passed-in generator is reused. + - CFG: unified `noise_pred = uncond + w * (text - uncond)` with optional `guidance_rescale`. + - Memory: denoising performs full-frame predictions (no spatial tiling); decoding can be tiled to avoid + OOM. + """ + + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + # 0. Input validation: height/width must be divisible by 32 + if height % 32 != 0 or width % 32 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.") + + self._guidance_scale = guidance_scale + self._guidance_rescale = guidance_rescale + self._attention_kwargs = attention_kwargs + self._interrupt = False + self._current_timestep = None + + # 1. Device & generator + device = self._execution_device + # Normalize generator input: accept list but use the first (batch_size=1) + if isinstance(generator, list): + generator = generator[0] + if seed is not None and generator is None: + generator = torch.Generator(device=device).manual_seed(seed) + + # 2. Optional i2v first frame conditioning: encode cond_image and inject at frame 0 via prepare_latents + cond_latents = None + if cond_image is not None and cond_strength > 0: + img = self.video_processor.preprocess(cond_image, height=height, width=width) + img = img.to(device=device, dtype=self.vae.dtype) + enc = self.vae.encode(img.unsqueeze(2)) # [B, C, 1, h, w] + cond_latents = enc.latent_dist.mode() if hasattr(enc, "latent_dist") else enc.latents + cond_latents = cond_latents.to(torch.float32) + cond_latents = self._normalize_latents( + cond_latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor + ) + + # 3. Global initial latents [B,C,F,H,W], optionally seeded/conditioned + latents, negative_index_latents, latent_num_frames, latent_height, latent_width = self.prepare_latents( + batch_size=1, + num_channels_latents=self.transformer.config.in_channels, + height=height, + width=width, + num_frames=num_frames, + device=device, + generator=generator, + dtype=torch.float32, + latents=latents, + cond_latents=cond_latents, + cond_strength=cond_strength, + negative_index_latents=negative_index_latents, + ) + if guidance_latents is not None: + guidance_latents = guidance_latents.to(device=device, dtype=torch.float32) + if latents.shape[2] != guidance_latents.shape[2]: + raise ValueError("The number of frames in `latents` and `guidance_latents` must be the same") + + # 4. Sliding windows in latent frames + tile_size_lat = max(1, temporal_tile_size // self.vae_temporal_compression_ratio) + overlap_lat = max(0, temporal_overlap // self.vae_temporal_compression_ratio) + windows = split_into_temporal_windows( + latent_num_frames, tile_size_lat, overlap_lat, self.vae_temporal_compression_ratio + ) + + # 5. Multi-prompt segments parsing + segment_texts = parse_prompt_segments(prompt, prompt_segments) + + out_latents = None + first_window_latents = None + + # 6. Process each temporal window + for w_idx, (w_start, w_end) in enumerate(windows): + if self.interrupt: + break + + # 6.1 Encode prompt embeddings per window segment + seg_index = min(w_idx, len(segment_texts) - 1) if segment_texts else 0 + pos_text = segment_texts[seg_index] if segment_texts else (prompt if isinstance(prompt, str) else "") + ( + prompt_embeds, + prompt_attention_mask, + negative_prompt_embeds, + negative_prompt_attention_mask, + ) = self.encode_prompt( + prompt=[pos_text], + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + num_videos_per_prompt=1, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + max_sequence_length=max_sequence_length, + device=device, + dtype=None, + ) + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) + + # 6.2 Window-level timesteps reset: fresh sampling for each temporal window + if sigmas is not None: + s = torch.tensor(sigmas, dtype=torch.float32) if not isinstance(sigmas, torch.Tensor) else sigmas + self.scheduler.set_timesteps(sigmas=s, device=device) + self._num_timesteps = len(sigmas) + else: + self.scheduler.set_timesteps(num_inference_steps=num_inference_steps, device=device) + self._num_timesteps = num_inference_steps + + # 6.3 Extract window latents [B,C,T,H,W] + window_latents = latents[:, :, w_start:w_end] + window_guidance_latents = guidance_latents[:, :, w_start:w_end] if guidance_latents is not None else None + window_T = window_latents.shape[2] + + # 6.4 Build per-window cond mask and inject previous tails / reference + window_cond_mask_5d = torch.ones( + (1, 1, window_T, latent_height, latent_width), device=device, dtype=torch.float32 + ) + self._current_tile_T = window_T + prev_overlap_len = 0 + # Inter-window tail latent injection (Extend) + if w_idx > 0 and overlap_lat > 0 and out_latents is not None: + k = min(overlap_lat, out_latents.shape[2]) + prev_tail = out_latents[:, :, -k:] + window_latents, window_cond_mask_5d, prev_overlap_len = inject_prev_tail_latents( + window_latents, + prev_tail, + window_cond_mask_5d, + overlap_lat, + temporal_overlap_cond_strength, + prev_overlap_len, + ) + # Reference/negative-index latent injection (append 1 frame at window head; controlled by negative_index_strength) + if window_guidance_latents is not None: + guiding_len = ( + window_guidance_latents.shape[2] if w_idx == 0 else window_guidance_latents.shape[2] - overlap_lat + ) + window_latents, window_cond_mask_5d, prev_overlap_len = inject_prev_tail_latents( + window_latents, + window_guidance_latents[:, :, -guiding_len:], + window_cond_mask_5d, + guiding_len, + guiding_strength, + prev_overlap_len, + ) + else: + guiding_len = 0 + window_latents, window_cond_mask_5d, prev_overlap_len = inject_prev_tail_latents( + window_latents, + negative_index_latents, + window_cond_mask_5d, + 1, + negative_index_strength, + prev_overlap_len, + ) + if w_idx == 0 and cond_image is not None and cond_strength > 0: + # First-frame I2V: smaller mask means stronger preservation of the original latent + window_cond_mask_5d[:, :, 0] = 1.0 - cond_strength + + # Update effective window latent sizes (consider injections on T/H/W) + w_B, w_C, w_T_eff, w_H_eff, w_W_eff = window_latents.shape + p = self.transformer_spatial_patch_size + pt = self.transformer_temporal_patch_size + + # 6.5 Pack full-window latents/masks once + # Seeding policy: derive a window-local generator to decouple RNG across windows + if seed is not None: + tile_seed = int(seed) + int(w_start) + local_gen = torch.Generator(device=device).manual_seed(tile_seed) + else: + local_gen = generator + # randn*mask + (1-mask)*latents implements hard-condition initialization + init_rand = randn_tensor(window_latents.shape, generator=local_gen, device=device, dtype=torch.float32) + mixed_latents = init_rand * window_cond_mask_5d + (1 - window_cond_mask_5d) * window_latents + window_latents_packed = self._pack_latents( + window_latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size + ) + latents_packed = self._pack_latents( + mixed_latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size + ) + cond_mask_tokens = self._pack_latents( + window_cond_mask_5d, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size + ) + if self.do_classifier_free_guidance: + cond_mask = torch.cat([cond_mask_tokens, cond_mask_tokens], dim=0) + else: + cond_mask = cond_mask_tokens + + # 6.6 Denoising loop per full window (no spatial tiling) + sigmas_current = self.scheduler.sigmas.to(device=latents_packed.device) + if sigmas_current.shape[0] >= 2: + for i, t in enumerate(self.progress_bar(self.scheduler.timesteps[:-1])): + if self.interrupt: + break + # Skip semantics: if sigma exceeds threshold, skip this step (do not call scheduler.step) + sigma_val = float(sigmas_current[i].item()) + if skip_steps_sigma_threshold is not None and float(skip_steps_sigma_threshold) > 0.0: + if sigma_val > float(skip_steps_sigma_threshold): + continue + + self._current_timestep = t + + # Model input (stack 2 copies under CFG) + latent_model_input = ( + torch.cat([latents_packed] * 2) if self.do_classifier_free_guidance else latents_packed + ) + # Broadcast timesteps, combine with per-token cond mask (I2V at window head) + timestep = t.expand(latent_model_input.shape[0]) + if cond_mask is not None: + # Broadcast timestep to per-token mask under CFG: [B] -> [B, S, 1] + timestep = timestep[:, None, None] * cond_mask + + # Micro-conditions: only provide video_coords (num_frames/height/width set to 1) + rope_interpolation_scale = ( + self.vae_temporal_compression_ratio, + self.vae_spatial_compression_ratio, + self.vae_spatial_compression_ratio, + ) + # Inpainting pre-blend (ComfyUI parity: KSamplerX0Inpaint:400) + if cond_mask_tokens is not None: + latents_packed = latents_packed * cond_mask_tokens + window_latents_packed * ( + 1.0 - cond_mask_tokens + ) + + # Negative-index/overlap lengths (for segmenting time coordinates; RoPE-compatible) + k_negative_count = ( + 1 if (negative_index_latents is not None and float(negative_index_strength) > 0.0) else 0 + ) + k_overlap_count = overlap_lat if (w_idx > 0 and overlap_lat > 0) else 0 + video_coords = build_video_coords_for_window( + latents=window_latents, + overlap_len=int(k_overlap_count), + guiding_len=int(guiding_len), + negative_len=int(k_negative_count), + rope_interpolation_scale=rope_interpolation_scale, + frame_rate=frame_rate, + ) + with self.transformer.cache_context("cond_uncond"): + noise_pred = self.transformer( + hidden_states=latent_model_input.to(dtype=self.transformer.dtype), + encoder_hidden_states=prompt_embeds, + timestep=timestep, + encoder_attention_mask=prompt_attention_mask, + num_frames=1, + height=1, + width=1, + rope_interpolation_scale=rope_interpolation_scale, + video_coords=video_coords, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] + + # Unified CFG + if self.do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) + if self.guidance_rescale > 0: + noise_pred = rescale_noise_cfg( + noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale + ) + + # Use global timestep for scheduling, but apply suppressive blending with hard-condition tokens (e.g., first frame) after step to avoid brightness/flicker due to time misalignment + latents_packed = self.scheduler.step( + noise_pred, t, latents_packed, generator=local_gen, return_dict=False + )[0] + # Inpainting post-blend (ComfyUI parity: restore hard-conditioned regions after update) + if cond_mask_tokens is not None: + latents_packed = latents_packed * cond_mask_tokens + window_latents_packed * ( + 1.0 - cond_mask_tokens + ) + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents_packed = callback_outputs.pop("latents", latents_packed) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + if XLA_AVAILABLE: + xm.mark_step() + else: + # Not enough sigmas to perform a valid step; skip this window safely. + pass + + # 6.7 Unpack back to [B,C,T,H,W] once + window_out = self._unpack_latents( + latents_packed, + w_T_eff, + w_H_eff, + w_W_eff, + p, + pt, + ) + if prev_overlap_len > 0: + window_out = window_out[:, :, :-prev_overlap_len] + + # 6.8 Overlap handling and fusion + if out_latents is None: + # First window: keep all latent frames and cache as AdaIN reference + out_latents = window_out + first_window_latents = out_latents + else: + window_out = window_out[:, :, 1:] # Drop the first frame of the new window + if adain_factor > 0 and first_window_latents is not None: + window_out = adain_normalize_latents(window_out, first_window_latents, adain_factor) + overlap_len = max(overlap_lat - 1, 1) + prev_tail_chunk = out_latents[:, :, -window_out.shape[2] :] + fused = linear_overlap_fuse(prev_tail_chunk, window_out, overlap_len) + out_latents = torch.cat([out_latents[:, :, : -window_out.shape[2]], fused], dim=2) + + # 7. Decode or return latent + if output_type == "latent": + video = out_latents + else: + # Decode via tiling to avoid OOM from full-frame decoding; latents are already de-normalized, so keep auto_denormalize disabled + video = self.vae_decode_tiled( + out_latents, + decode_timestep=decode_timestep, + decode_noise_scale=decode_noise_scale, + horizontal_tiles=int(decode_horizontal_tiles), + vertical_tiles=int(decode_vertical_tiles), + overlap=int(decode_overlap), + generator=generator, + output_type=output_type, # Keep type consistent; postprocess is applied afterwards + ) + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (video,) + + return LTXPipelineOutput(frames=video) diff --git a/src/diffusers/pipelines/ltx/pipeline_ltx_image2video.py b/src/diffusers/pipelines/ltx/pipeline_ltx_image2video.py index f30f8a3dc8f6..497f505c4dd8 100644 --- a/src/diffusers/pipelines/ltx/pipeline_ltx_image2video.py +++ b/src/diffusers/pipelines/ltx/pipeline_ltx_image2video.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -88,10 +88,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -106,15 +106,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -147,7 +147,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -256,11 +256,11 @@ def __init__( def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 128, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -305,25 +305,25 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.mochi.pipeline_mochi.MochiPipeline.encode_prompt with 256->128 def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, max_sequence_length: int = 128, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -501,16 +501,16 @@ def _denormalize_latents( def prepare_latents( self, - image: Optional[torch.Tensor] = None, + image: torch.Tensor | None = None, batch_size: int = 1, num_channels_latents: int = 128, height: int = 512, width: int = 704, num_frames: int = 161, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[torch.Generator] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: height = height // self.vae_spatial_compression_ratio width = width // self.vae_spatial_compression_ratio @@ -598,30 +598,30 @@ def interrupt(self): def __call__( self, image: PipelineImageInput = None, - prompt: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] | None = None, height: int = 512, width: int = 704, num_frames: int = 161, frame_rate: int = 25, num_inference_steps: int = 50, - timesteps: List[int] = None, + timesteps: list[int] = None, guidance_scale: float = 3, guidance_rescale: float = 0.0, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - decode_timestep: Union[float, List[float]] = 0.0, - decode_noise_scale: Optional[Union[float, List[float]]] = None, - output_type: Optional[str] = "pil", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + decode_timestep: float | list[float] = 0.0, + decode_noise_scale: float | list[float] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 128, ): r""" @@ -630,7 +630,7 @@ def __call__( Args: image (`PipelineImageInput`): The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. height (`int`, defaults to `512`): @@ -642,7 +642,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. @@ -660,7 +660,7 @@ def __call__( using zero terminal SNR. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of videos to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -695,7 +695,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -798,10 +798,14 @@ def __call__( self.scheduler.config.get("base_shift", 0.5), self.scheduler.config.get("max_shift", 1.15), ) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, - device, + timestep_device, timesteps, sigmas=sigmas, mu=mu, diff --git a/src/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py b/src/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py index 9acff105e56d..17d4e1d8fc57 100644 --- a/src/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py +++ b/src/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Optional, Union - import torch from ...image_processor import PipelineImageInput @@ -31,7 +29,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -65,12 +63,12 @@ def __init__( def prepare_latents( self, - video: Optional[torch.Tensor] = None, + video: torch.Tensor | None = None, batch_size: int = 1, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[torch.Generator] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: if latents is not None: return latents.to(device=device, dtype=dtype) @@ -243,16 +241,16 @@ def check_inputs(self, video, height, width, latents, tone_map_compression_ratio @torch.no_grad() def __call__( self, - video: Optional[List[PipelineImageInput]] = None, + video: list[PipelineImageInput] | None = None, height: int = 512, width: int = 704, - latents: Optional[torch.Tensor] = None, - decode_timestep: Union[float, List[float]] = 0.0, - decode_noise_scale: Optional[Union[float, List[float]]] = None, + latents: torch.Tensor | None = None, + decode_timestep: float | list[float] = 0.0, + decode_noise_scale: float | list[float] | None = None, adain_factor: float = 0.0, tone_map_compression_ratio: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + output_type: str | None = "pil", return_dict: bool = True, ): self.check_inputs( diff --git a/src/diffusers/pipelines/ltx/pipeline_output.py b/src/diffusers/pipelines/ltx/pipeline_output.py index 36ec3ea884a2..f5cb34aa508d 100644 --- a/src/diffusers/pipelines/ltx/pipeline_output.py +++ b/src/diffusers/pipelines/ltx/pipeline_output.py @@ -11,8 +11,8 @@ class LTXPipelineOutput(BaseOutput): Output class for LTX pipelines. Args: - frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): - List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): + list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape `(batch_size, num_frames, channels, height, width)`. """ diff --git a/src/diffusers/pipelines/ltx2/__init__.py b/src/diffusers/pipelines/ltx2/__init__.py new file mode 100644 index 000000000000..d6a408d5c546 --- /dev/null +++ b/src/diffusers/pipelines/ltx2/__init__.py @@ -0,0 +1,60 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["connectors"] = ["LTX2TextConnectors"] + _import_structure["latent_upsampler"] = ["LTX2LatentUpsamplerModel"] + _import_structure["pipeline_ltx2"] = ["LTX2Pipeline"] + _import_structure["pipeline_ltx2_condition"] = ["LTX2ConditionPipeline"] + _import_structure["pipeline_ltx2_image2video"] = ["LTX2ImageToVideoPipeline"] + _import_structure["pipeline_ltx2_latent_upsample"] = ["LTX2LatentUpsamplePipeline"] + _import_structure["vocoder"] = ["LTX2Vocoder"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * + else: + from .connectors import LTX2TextConnectors + from .latent_upsampler import LTX2LatentUpsamplerModel + from .pipeline_ltx2 import LTX2Pipeline + from .pipeline_ltx2_condition import LTX2ConditionPipeline + from .pipeline_ltx2_image2video import LTX2ImageToVideoPipeline + from .pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline + from .vocoder import LTX2Vocoder + +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/ltx2/connectors.py b/src/diffusers/pipelines/ltx2/connectors.py new file mode 100644 index 000000000000..4b2a81a9dc2c --- /dev/null +++ b/src/diffusers/pipelines/ltx2/connectors.py @@ -0,0 +1,324 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import PeftAdapterMixin +from ...models.attention import FeedForward +from ...models.modeling_utils import ModelMixin +from ...models.transformers.transformer_ltx2 import LTX2Attention, LTX2AudioVideoAttnProcessor + + +class LTX2RotaryPosEmbed1d(nn.Module): + """ + 1D rotary positional embeddings (RoPE) for the LTX 2.0 text encoder connectors. + """ + + def __init__( + self, + dim: int, + base_seq_len: int = 4096, + theta: float = 10000.0, + double_precision: bool = True, + rope_type: str = "interleaved", + num_attention_heads: int = 32, + ): + super().__init__() + if rope_type not in ["interleaved", "split"]: + raise ValueError(f"{rope_type=} not supported. Choose between 'interleaved' and 'split'.") + + self.dim = dim + self.base_seq_len = base_seq_len + self.theta = theta + self.double_precision = double_precision + self.rope_type = rope_type + self.num_attention_heads = num_attention_heads + + def forward( + self, + batch_size: int, + pos: int, + device: str | torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + # 1. Get 1D position ids + grid_1d = torch.arange(pos, dtype=torch.float32, device=device) + # Get fractional indices relative to self.base_seq_len + grid_1d = grid_1d / self.base_seq_len + grid = grid_1d.unsqueeze(0).repeat(batch_size, 1) # [batch_size, seq_len] + + # 2. Calculate 1D RoPE frequencies + num_rope_elems = 2 # 1 (because 1D) * 2 (for cos, sin) = 2 + freqs_dtype = torch.float64 if self.double_precision else torch.float32 + pow_indices = torch.pow( + self.theta, + torch.linspace(start=0.0, end=1.0, steps=self.dim // num_rope_elems, dtype=freqs_dtype, device=device), + ) + freqs = (pow_indices * torch.pi / 2.0).to(dtype=torch.float32) + + # 3. Matrix-vector outer product between pos ids of shape (batch_size, seq_len) and freqs vector of shape + # (self.dim // 2,). + freqs = (grid.unsqueeze(-1) * 2 - 1) * freqs # [B, seq_len, self.dim // 2] + + # 4. Get real, interleaved (cos, sin) frequencies, padded to self.dim + if self.rope_type == "interleaved": + cos_freqs = freqs.cos().repeat_interleave(2, dim=-1) + sin_freqs = freqs.sin().repeat_interleave(2, dim=-1) + + if self.dim % num_rope_elems != 0: + cos_padding = torch.ones_like(cos_freqs[:, :, : self.dim % num_rope_elems]) + sin_padding = torch.zeros_like(sin_freqs[:, :, : self.dim % num_rope_elems]) + cos_freqs = torch.cat([cos_padding, cos_freqs], dim=-1) + sin_freqs = torch.cat([sin_padding, sin_freqs], dim=-1) + + elif self.rope_type == "split": + expected_freqs = self.dim // 2 + current_freqs = freqs.shape[-1] + pad_size = expected_freqs - current_freqs + cos_freq = freqs.cos() + sin_freq = freqs.sin() + + if pad_size != 0: + cos_padding = torch.ones_like(cos_freq[:, :, :pad_size]) + sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size]) + + cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1) + sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1) + + # Reshape freqs to be compatible with multi-head attention + b = cos_freq.shape[0] + t = cos_freq.shape[1] + + cos_freq = cos_freq.reshape(b, t, self.num_attention_heads, -1) + sin_freq = sin_freq.reshape(b, t, self.num_attention_heads, -1) + + cos_freqs = torch.swapaxes(cos_freq, 1, 2) # (B,H,T,D//2) + sin_freqs = torch.swapaxes(sin_freq, 1, 2) # (B,H,T,D//2) + + return cos_freqs, sin_freqs + + +class LTX2TransformerBlock1d(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + activation_fn: str = "gelu-approximate", + eps: float = 1e-6, + rope_type: str = "interleaved", + ): + super().__init__() + + self.norm1 = torch.nn.RMSNorm(dim, eps=eps, elementwise_affine=False) + self.attn1 = LTX2Attention( + query_dim=dim, + heads=num_attention_heads, + kv_heads=num_attention_heads, + dim_head=attention_head_dim, + processor=LTX2AudioVideoAttnProcessor(), + rope_type=rope_type, + ) + + self.norm2 = torch.nn.RMSNorm(dim, eps=eps, elementwise_affine=False) + self.ff = FeedForward(dim, activation_fn=activation_fn) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + rotary_emb: torch.Tensor | None = None, + ) -> torch.Tensor: + norm_hidden_states = self.norm1(hidden_states) + attn_hidden_states = self.attn1(norm_hidden_states, attention_mask=attention_mask, query_rotary_emb=rotary_emb) + hidden_states = hidden_states + attn_hidden_states + + norm_hidden_states = self.norm2(hidden_states) + ff_hidden_states = self.ff(norm_hidden_states) + hidden_states = hidden_states + ff_hidden_states + + return hidden_states + + +class LTX2ConnectorTransformer1d(nn.Module): + """ + A 1D sequence transformer for modalities such as text. + + In LTX 2.0, this is used to process the text encoder hidden states for each of the video and audio streams. + """ + + _supports_gradient_checkpointing = True + + def __init__( + self, + num_attention_heads: int = 30, + attention_head_dim: int = 128, + num_layers: int = 2, + num_learnable_registers: int | None = 128, + rope_base_seq_len: int = 4096, + rope_theta: float = 10000.0, + rope_double_precision: bool = True, + eps: float = 1e-6, + causal_temporal_positioning: bool = False, + rope_type: str = "interleaved", + ): + super().__init__() + self.num_attention_heads = num_attention_heads + self.inner_dim = num_attention_heads * attention_head_dim + self.causal_temporal_positioning = causal_temporal_positioning + + self.num_learnable_registers = num_learnable_registers + self.learnable_registers = None + if num_learnable_registers is not None: + init_registers = torch.rand(num_learnable_registers, self.inner_dim) * 2.0 - 1.0 + self.learnable_registers = torch.nn.Parameter(init_registers) + + self.rope = LTX2RotaryPosEmbed1d( + self.inner_dim, + base_seq_len=rope_base_seq_len, + theta=rope_theta, + double_precision=rope_double_precision, + rope_type=rope_type, + num_attention_heads=num_attention_heads, + ) + + self.transformer_blocks = torch.nn.ModuleList( + [ + LTX2TransformerBlock1d( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + rope_type=rope_type, + ) + for _ in range(num_layers) + ] + ) + + self.norm_out = torch.nn.RMSNorm(self.inner_dim, eps=eps, elementwise_affine=False) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + attn_mask_binarize_threshold: float = -9000.0, + ) -> tuple[torch.Tensor, torch.Tensor]: + # hidden_states shape: [batch_size, seq_len, hidden_dim] + # attention_mask shape: [batch_size, seq_len] or [batch_size, 1, 1, seq_len] + batch_size, seq_len, _ = hidden_states.shape + + # 1. Replace padding with learned registers, if using + if self.learnable_registers is not None: + if seq_len % self.num_learnable_registers != 0: + raise ValueError( + f"The `hidden_states` sequence length {hidden_states.shape[1]} should be divisible by the number" + f" of learnable registers {self.num_learnable_registers}" + ) + + num_register_repeats = seq_len // self.num_learnable_registers + registers = torch.tile(self.learnable_registers, (num_register_repeats, 1)) # [seq_len, inner_dim] + + binary_attn_mask = (attention_mask >= attn_mask_binarize_threshold).int() + if binary_attn_mask.ndim == 4: + binary_attn_mask = binary_attn_mask.squeeze(1).squeeze(1) # [B, 1, 1, L] --> [B, L] + + hidden_states_non_padded = [hidden_states[i, binary_attn_mask[i].bool(), :] for i in range(batch_size)] + valid_seq_lens = [x.shape[0] for x in hidden_states_non_padded] + pad_lengths = [seq_len - valid_seq_len for valid_seq_len in valid_seq_lens] + padded_hidden_states = [ + F.pad(x, pad=(0, 0, 0, p), value=0) for x, p in zip(hidden_states_non_padded, pad_lengths) + ] + padded_hidden_states = torch.cat([x.unsqueeze(0) for x in padded_hidden_states], dim=0) # [B, L, D] + + flipped_mask = torch.flip(binary_attn_mask, dims=[1]).unsqueeze(-1) # [B, L, 1] + hidden_states = flipped_mask * padded_hidden_states + (1 - flipped_mask) * registers + + # Overwrite attention_mask with an all-zeros mask if using registers. + attention_mask = torch.zeros_like(attention_mask) + + # 2. Calculate 1D RoPE positional embeddings + rotary_emb = self.rope(batch_size, seq_len, device=hidden_states.device) + + # 3. Run 1D transformer blocks + for block in self.transformer_blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func(block, hidden_states, attention_mask, rotary_emb) + else: + hidden_states = block(hidden_states, attention_mask=attention_mask, rotary_emb=rotary_emb) + + hidden_states = self.norm_out(hidden_states) + + return hidden_states, attention_mask + + +class LTX2TextConnectors(ModelMixin, PeftAdapterMixin, ConfigMixin): + """ + Text connector stack used by LTX 2.0 to process the packed text encoder hidden states for both the video and audio + streams. + """ + + @register_to_config + def __init__( + self, + caption_channels: int, + text_proj_in_factor: int, + video_connector_num_attention_heads: int, + video_connector_attention_head_dim: int, + video_connector_num_layers: int, + video_connector_num_learnable_registers: int | None, + audio_connector_num_attention_heads: int, + audio_connector_attention_head_dim: int, + audio_connector_num_layers: int, + audio_connector_num_learnable_registers: int | None, + connector_rope_base_seq_len: int, + rope_theta: float, + rope_double_precision: bool, + causal_temporal_positioning: bool, + rope_type: str = "interleaved", + ): + super().__init__() + self.text_proj_in = nn.Linear(caption_channels * text_proj_in_factor, caption_channels, bias=False) + self.video_connector = LTX2ConnectorTransformer1d( + num_attention_heads=video_connector_num_attention_heads, + attention_head_dim=video_connector_attention_head_dim, + num_layers=video_connector_num_layers, + num_learnable_registers=video_connector_num_learnable_registers, + rope_base_seq_len=connector_rope_base_seq_len, + rope_theta=rope_theta, + rope_double_precision=rope_double_precision, + causal_temporal_positioning=causal_temporal_positioning, + rope_type=rope_type, + ) + self.audio_connector = LTX2ConnectorTransformer1d( + num_attention_heads=audio_connector_num_attention_heads, + attention_head_dim=audio_connector_attention_head_dim, + num_layers=audio_connector_num_layers, + num_learnable_registers=audio_connector_num_learnable_registers, + rope_base_seq_len=connector_rope_base_seq_len, + rope_theta=rope_theta, + rope_double_precision=rope_double_precision, + causal_temporal_positioning=causal_temporal_positioning, + rope_type=rope_type, + ) + + def forward( + self, text_encoder_hidden_states: torch.Tensor, attention_mask: torch.Tensor, additive_mask: bool = False + ): + # Convert to additive attention mask, if necessary + if not additive_mask: + text_dtype = text_encoder_hidden_states.dtype + attention_mask = (attention_mask - 1).reshape(attention_mask.shape[0], 1, -1, attention_mask.shape[-1]) + attention_mask = attention_mask.to(text_dtype) * torch.finfo(text_dtype).max + + text_encoder_hidden_states = self.text_proj_in(text_encoder_hidden_states) + + video_text_embedding, new_attn_mask = self.video_connector(text_encoder_hidden_states, attention_mask) + + attn_mask = (new_attn_mask < 1e-6).to(torch.int64) + attn_mask = attn_mask.reshape(video_text_embedding.shape[0], video_text_embedding.shape[1], 1) + video_text_embedding = video_text_embedding * attn_mask + new_attn_mask = attn_mask.squeeze(-1) + + audio_text_embedding, _ = self.audio_connector(text_encoder_hidden_states, attention_mask) + + return video_text_embedding, audio_text_embedding, new_attn_mask diff --git a/src/diffusers/pipelines/ltx2/export_utils.py b/src/diffusers/pipelines/ltx2/export_utils.py new file mode 100644 index 000000000000..f0287506b8db --- /dev/null +++ b/src/diffusers/pipelines/ltx2/export_utils.py @@ -0,0 +1,191 @@ +# Copyright 2025 The Lightricks team and The HuggingFace Team. +# All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Iterator +from fractions import Fraction +from itertools import chain + +import numpy as np +import PIL.Image +import torch +from tqdm import tqdm + +from ...utils import get_logger, is_av_available + + +logger = get_logger(__name__) # pylint: disable=invalid-name + + +_CAN_USE_AV = is_av_available() +if _CAN_USE_AV: + import av +else: + raise ImportError( + "PyAV is required to use LTX 2.0 video export utilities. You can install it with `pip install av`" + ) + + +def _prepare_audio_stream(container: av.container.Container, audio_sample_rate: int) -> av.audio.AudioStream: + """ + Prepare the audio stream for writing. + """ + audio_stream = container.add_stream("aac", rate=audio_sample_rate) + audio_stream.codec_context.sample_rate = audio_sample_rate + audio_stream.codec_context.layout = "stereo" + audio_stream.codec_context.time_base = Fraction(1, audio_sample_rate) + return audio_stream + + +def _resample_audio( + container: av.container.Container, audio_stream: av.audio.AudioStream, frame_in: av.AudioFrame +) -> None: + cc = audio_stream.codec_context + + # Use the encoder's format/layout/rate as the *target* + target_format = cc.format or "fltp" # AAC → usually fltp + target_layout = cc.layout or "stereo" + target_rate = cc.sample_rate or frame_in.sample_rate + + audio_resampler = av.audio.resampler.AudioResampler( + format=target_format, + layout=target_layout, + rate=target_rate, + ) + + audio_next_pts = 0 + for rframe in audio_resampler.resample(frame_in): + if rframe.pts is None: + rframe.pts = audio_next_pts + audio_next_pts += rframe.samples + rframe.sample_rate = frame_in.sample_rate + container.mux(audio_stream.encode(rframe)) + + # flush audio encoder + for packet in audio_stream.encode(): + container.mux(packet) + + +def _write_audio( + container: av.container.Container, + audio_stream: av.audio.AudioStream, + samples: torch.Tensor, + audio_sample_rate: int, +) -> None: + if samples.ndim == 1: + samples = samples[:, None] + + if samples.shape[1] != 2 and samples.shape[0] == 2: + samples = samples.T + + if samples.shape[1] != 2: + raise ValueError(f"Expected samples with 2 channels; got shape {samples.shape}.") + + # Convert to int16 packed for ingestion; resampler converts to encoder fmt. + if samples.dtype != torch.int16: + samples = torch.clip(samples, -1.0, 1.0) + samples = (samples * 32767.0).to(torch.int16) + + frame_in = av.AudioFrame.from_ndarray( + samples.contiguous().reshape(1, -1).cpu().numpy(), + format="s16", + layout="stereo", + ) + frame_in.sample_rate = audio_sample_rate + + _resample_audio(container, audio_stream, frame_in) + + +def encode_video( + video: list[PIL.Image.Image] | np.ndarray | torch.Tensor | Iterator[torch.Tensor], + fps: int, + audio: torch.Tensor, + audio_sample_rate: int, + output_path: str, + video_chunks_number: int = 1, +) -> None: + """ + Encodes a video with audio using the PyAV library. Based on code from the original LTX-2 repo: + https://github.com/Lightricks/LTX-2/blob/4f410820b198e05074a1e92de793e3b59e9ab5a0/packages/ltx-pipelines/src/ltx_pipelines/utils/media_io.py#L182 + + Args: + video (`List[PIL.Image.Image]` or `np.ndarray` or `torch.Tensor`): + A video tensor of shape [frames, height, width, channels] with integer pixel values in [0, 255]. If the + input is a `np.ndarray`, it is expected to be a float array with values in [0, 1] (which is what pipelines + usually return with `output_type="np"`). + fps (`int`) + The frames per second (FPS) of the encoded video. + audio (`torch.Tensor`, *optional*): + An audio waveform of shape [audio_channels, samples]. + audio_sample_rate: (`int`, *optional*): + The sampling rate of the audio waveform. For LTX 2, this is typically 24000 (24 kHz). + output_path (`str`): + The path to save the encoded video to. + video_chunks_number (`int`, *optional*, defaults to `1`): + The number of chunks to split the video into for encoding. Each chunk will be encoded separately. The + number of chunks to use often depends on the tiling config for the video VAE. + """ + if isinstance(video, list) and isinstance(video[0], PIL.Image.Image): + # Pipeline output_type="pil"; assumes each image is in "RGB" mode + video_frames = [np.array(frame) for frame in video] + video = np.stack(video_frames, axis=0) + video = torch.from_numpy(video) + elif isinstance(video, np.ndarray): + # Pipeline output_type="np" + is_denormalized = np.logical_and(np.zeros_like(video) <= video, video <= np.ones_like(video)) + if np.all(is_denormalized): + video = (video * 255).round().astype("uint8") + else: + logger.warning( + "Supplied `numpy.ndarray` does not have values in [0, 1]. The values will be assumed to be pixel " + "values in [0, ..., 255] and will be used as is." + ) + video = torch.from_numpy(video) + + if isinstance(video, torch.Tensor): + # Split into video_chunks_number along the frame dimension + video = torch.tensor_split(video, video_chunks_number, dim=0) + video = iter(video) + + first_chunk = next(video) + + _, height, width, _ = first_chunk.shape + + container = av.open(output_path, mode="w") + stream = container.add_stream("libx264", rate=int(fps)) + stream.width = width + stream.height = height + stream.pix_fmt = "yuv420p" + + if audio is not None: + if audio_sample_rate is None: + raise ValueError("audio_sample_rate is required when audio is provided") + + audio_stream = _prepare_audio_stream(container, audio_sample_rate) + + for video_chunk in tqdm(chain([first_chunk], video), total=video_chunks_number, desc="Encoding video chunks"): + video_chunk_cpu = video_chunk.to("cpu").numpy() + for frame_array in video_chunk_cpu: + frame = av.VideoFrame.from_ndarray(frame_array, format="rgb24") + for packet in stream.encode(frame): + container.mux(packet) + + # Flush encoder + for packet in stream.encode(): + container.mux(packet) + + if audio is not None: + _write_audio(container, audio_stream, audio, audio_sample_rate) + + container.close() diff --git a/src/diffusers/pipelines/ltx2/latent_upsampler.py b/src/diffusers/pipelines/ltx2/latent_upsampler.py new file mode 100644 index 000000000000..f6c589a70ab6 --- /dev/null +++ b/src/diffusers/pipelines/ltx2/latent_upsampler.py @@ -0,0 +1,284 @@ +# Copyright 2025 Lightricks and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math + +import torch +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...models.modeling_utils import ModelMixin + + +RATIONAL_RESAMPLER_SCALE_MAPPING = { + 0.75: (3, 4), + 1.5: (3, 2), + 2.0: (2, 1), + 4.0: (4, 1), +} + + +# Copied from diffusers.pipelines.ltx.modeling_latent_upsampler.ResBlock +class ResBlock(torch.nn.Module): + def __init__(self, channels: int, mid_channels: int | None = None, dims: int = 3): + super().__init__() + if mid_channels is None: + mid_channels = channels + + Conv = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d + + self.conv1 = Conv(channels, mid_channels, kernel_size=3, padding=1) + self.norm1 = torch.nn.GroupNorm(32, mid_channels) + self.conv2 = Conv(mid_channels, channels, kernel_size=3, padding=1) + self.norm2 = torch.nn.GroupNorm(32, channels) + self.activation = torch.nn.SiLU() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + residual = hidden_states + hidden_states = self.conv1(hidden_states) + hidden_states = self.norm1(hidden_states) + hidden_states = self.activation(hidden_states) + hidden_states = self.conv2(hidden_states) + hidden_states = self.norm2(hidden_states) + hidden_states = self.activation(hidden_states + residual) + return hidden_states + + +# Copied from diffusers.pipelines.ltx.modeling_latent_upsampler.PixelShuffleND +class PixelShuffleND(torch.nn.Module): + def __init__(self, dims, upscale_factors=(2, 2, 2)): + super().__init__() + + self.dims = dims + self.upscale_factors = upscale_factors + + if dims not in [1, 2, 3]: + raise ValueError("dims must be 1, 2, or 3") + + def forward(self, x): + if self.dims == 3: + # spatiotemporal: b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3) + return ( + x.unflatten(1, (-1, *self.upscale_factors[:3])) + .permute(0, 1, 5, 2, 6, 3, 7, 4) + .flatten(6, 7) + .flatten(4, 5) + .flatten(2, 3) + ) + elif self.dims == 2: + # spatial: b (c p1 p2) h w -> b c (h p1) (w p2) + return ( + x.unflatten(1, (-1, *self.upscale_factors[:2])).permute(0, 1, 4, 2, 5, 3).flatten(4, 5).flatten(2, 3) + ) + elif self.dims == 1: + # temporal: b (c p1) f h w -> b c (f p1) h w + return x.unflatten(1, (-1, *self.upscale_factors[:1])).permute(0, 1, 3, 2, 4, 5).flatten(2, 3) + + +class BlurDownsample(torch.nn.Module): + """ + Anti-aliased spatial downsampling by integer stride using a fixed separable binomial kernel. Applies only on H,W. + Works for dims=2 or dims=3 (per-frame). + """ + + def __init__(self, dims: int, stride: int, kernel_size: int = 5) -> None: + super().__init__() + + if dims not in (2, 3): + raise ValueError(f"`dims` must be either 2 or 3 but is {dims}") + if kernel_size < 3 or kernel_size % 2 != 1: + raise ValueError(f"`kernel_size` must be an odd number >= 3 but is {kernel_size}") + + self.dims = dims + self.stride = stride + self.kernel_size = kernel_size + + # 5x5 separable binomial kernel using binomial coefficients [1, 4, 6, 4, 1] from + # the 4th row of Pascal's triangle. This kernel is used for anti-aliasing and + # provides a smooth approximation of a Gaussian filter (often called a "binomial filter"). + # The 2D kernel is constructed as the outer product and normalized. + k = torch.tensor([math.comb(kernel_size - 1, k) for k in range(kernel_size)]) + k2d = k[:, None] @ k[None, :] + k2d = (k2d / k2d.sum()).float() # shape (kernel_size, kernel_size) + self.register_buffer("kernel", k2d[None, None, :, :]) # (1, 1, kernel_size, kernel_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.stride == 1: + return x + + if self.dims == 2: + c = x.shape[1] + weight = self.kernel.expand(c, 1, self.kernel_size, self.kernel_size) # depthwise + x = F.conv2d(x, weight=weight, bias=None, stride=self.stride, padding=self.kernel_size // 2, groups=c) + else: + # dims == 3: apply per-frame on H,W + b, c, f, _, _ = x.shape + x = x.transpose(1, 2).flatten(0, 1) # [B, C, F, H, W] --> [B * F, C, H, W] + + weight = self.kernel.expand(c, 1, self.kernel_size, self.kernel_size) # depthwise + x = F.conv2d(x, weight=weight, bias=None, stride=self.stride, padding=self.kernel_size // 2, groups=c) + + h2, w2 = x.shape[-2:] + x = x.unflatten(0, (b, f)).reshape(b, -1, f, h2, w2) # [B * F, C, H, W] --> [B, C, F, H, W] + return x + + +class SpatialRationalResampler(torch.nn.Module): + """ + Scales by the spatial size of the input by a rational number `scale`. For example, `scale = 0.75` will downsample + by a factor of 3 / 4, while `scale = 1.5` will upsample by a factor of 3 / 2. This works by first upsampling the + input by the (integer) numerator of `scale`, and then performing a blur + stride anti-aliased downsample by the + (integer) denominator. + """ + + def __init__(self, mid_channels: int = 1024, scale: float = 2.0): + super().__init__() + self.scale = float(scale) + num_denom = RATIONAL_RESAMPLER_SCALE_MAPPING.get(scale, None) + if num_denom is None: + raise ValueError( + f"The supplied `scale` {scale} is not supported; supported scales are {list(RATIONAL_RESAMPLER_SCALE_MAPPING.keys())}" + ) + self.num, self.den = num_denom + + self.conv = torch.nn.Conv2d(mid_channels, (self.num**2) * mid_channels, kernel_size=3, padding=1) + self.pixel_shuffle = PixelShuffleND(2, upscale_factors=(self.num, self.num)) + self.blur_down = BlurDownsample(dims=2, stride=self.den) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Expected x shape: [B * F, C, H, W] + # b, _, f, h, w = x.shape + # x = x.transpose(1, 2).flatten(0, 1) # [B, C, F, H, W] --> [B * F, C, H, W] + x = self.conv(x) + x = self.pixel_shuffle(x) + x = self.blur_down(x) + # x = x.unflatten(0, (b, f)).reshape(b, -1, f, h, w) # [B * F, C, H, W] --> [B, C, F, H, W] + return x + + +class LTX2LatentUpsamplerModel(ModelMixin, ConfigMixin): + """ + Model to spatially upsample VAE latents. + + Args: + in_channels (`int`, defaults to `128`): + Number of channels in the input latent + mid_channels (`int`, defaults to `512`): + Number of channels in the middle layers + num_blocks_per_stage (`int`, defaults to `4`): + Number of ResBlocks to use in each stage (pre/post upsampling) + dims (`int`, defaults to `3`): + Number of dimensions for convolutions (2 or 3) + spatial_upsample (`bool`, defaults to `True`): + Whether to spatially upsample the latent + temporal_upsample (`bool`, defaults to `False`): + Whether to temporally upsample the latent + """ + + @register_to_config + def __init__( + self, + in_channels: int = 128, + mid_channels: int = 1024, + num_blocks_per_stage: int = 4, + dims: int = 3, + spatial_upsample: bool = True, + temporal_upsample: bool = False, + rational_spatial_scale: float | None = 2.0, + ): + super().__init__() + + self.in_channels = in_channels + self.mid_channels = mid_channels + self.num_blocks_per_stage = num_blocks_per_stage + self.dims = dims + self.spatial_upsample = spatial_upsample + self.temporal_upsample = temporal_upsample + + ConvNd = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d + + self.initial_conv = ConvNd(in_channels, mid_channels, kernel_size=3, padding=1) + self.initial_norm = torch.nn.GroupNorm(32, mid_channels) + self.initial_activation = torch.nn.SiLU() + + self.res_blocks = torch.nn.ModuleList([ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]) + + if spatial_upsample and temporal_upsample: + self.upsampler = torch.nn.Sequential( + torch.nn.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1), + PixelShuffleND(3), + ) + elif spatial_upsample: + if rational_spatial_scale is not None: + self.upsampler = SpatialRationalResampler(mid_channels=mid_channels, scale=rational_spatial_scale) + else: + self.upsampler = torch.nn.Sequential( + torch.nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1), + PixelShuffleND(2), + ) + elif temporal_upsample: + self.upsampler = torch.nn.Sequential( + torch.nn.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1), + PixelShuffleND(1), + ) + else: + raise ValueError("Either spatial_upsample or temporal_upsample must be True") + + self.post_upsample_res_blocks = torch.nn.ModuleList( + [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)] + ) + + self.final_conv = ConvNd(mid_channels, in_channels, kernel_size=3, padding=1) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = hidden_states.shape + + if self.dims == 2: + hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) + hidden_states = self.initial_conv(hidden_states) + hidden_states = self.initial_norm(hidden_states) + hidden_states = self.initial_activation(hidden_states) + + for block in self.res_blocks: + hidden_states = block(hidden_states) + + hidden_states = self.upsampler(hidden_states) + + for block in self.post_upsample_res_blocks: + hidden_states = block(hidden_states) + + hidden_states = self.final_conv(hidden_states) + hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) + else: + hidden_states = self.initial_conv(hidden_states) + hidden_states = self.initial_norm(hidden_states) + hidden_states = self.initial_activation(hidden_states) + + for block in self.res_blocks: + hidden_states = block(hidden_states) + + if self.temporal_upsample: + hidden_states = self.upsampler(hidden_states) + hidden_states = hidden_states[:, :, 1:, :, :] + else: + hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) + hidden_states = self.upsampler(hidden_states) + hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) + + for block in self.post_upsample_res_blocks: + hidden_states = block(hidden_states) + + hidden_states = self.final_conv(hidden_states) + + return hidden_states diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py new file mode 100644 index 000000000000..037840360137 --- /dev/null +++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py @@ -0,0 +1,1226 @@ +# Copyright 2025 Lightricks and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import inspect +from typing import Any, Callable + +import numpy as np +import torch +from transformers import Gemma3ForConditionalGeneration, GemmaTokenizer, GemmaTokenizerFast + +from ...callbacks import MultiPipelineCallbacks, PipelineCallback +from ...loaders import FromSingleFileMixin, LTX2LoraLoaderMixin +from ...models.autoencoders import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video +from ...models.transformers import LTX2VideoTransformer3DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import is_torch_xla_available, logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor +from ..pipeline_utils import DiffusionPipeline +from .connectors import LTX2TextConnectors +from .pipeline_output import LTX2PipelineOutput +from .vocoder import LTX2Vocoder + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import LTX2Pipeline + >>> from diffusers.pipelines.ltx2.export_utils import encode_video + + >>> pipe = LTX2Pipeline.from_pretrained("Lightricks/LTX-2", torch_dtype=torch.bfloat16) + >>> pipe.enable_model_cpu_offload() + + >>> prompt = "A woman with long brown hair and light skin smiles at another woman with long blonde hair. The woman with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek. The camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be real-life footage" + >>> negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted" + + >>> frame_rate = 24.0 + >>> video, audio = pipe( + ... prompt=prompt, + ... negative_prompt=negative_prompt, + ... width=768, + ... height=512, + ... num_frames=121, + ... frame_rate=frame_rate, + ... num_inference_steps=40, + ... guidance_scale=4.0, + ... output_type="np", + ... return_dict=False, + ... ) + + >>> encode_video( + ... video[0], + ... fps=frame_rate, + ... audio=audio[0].float().cpu(), + ... audio_sample_rate=pipe.vocoder.config.output_sampling_rate, # should be 24000 + ... output_path="video.mp4", + ... ) + ``` +""" + + +# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`list[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`list[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg +def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): + r""" + Rescales `noise_cfg` tensor based on `guidance_rescale` to improve image quality and fix overexposure. Based on + Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are + Flawed](https://huggingface.co/papers/2305.08891). + + Args: + noise_cfg (`torch.Tensor`): + The predicted noise tensor for the guided diffusion process. + noise_pred_text (`torch.Tensor`): + The predicted noise tensor for the text-guided diffusion process. + guidance_rescale (`float`, *optional*, defaults to 0.0): + A rescale factor applied to the noise predictions. + + Returns: + noise_cfg (`torch.Tensor`): The rescaled noise prediction tensor. + """ + std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True) + std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) + # rescale the results from guidance (fixes overexposure) + noise_pred_rescaled = noise_cfg * (std_text / std_cfg) + # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images + noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg + return noise_cfg + + +class LTX2Pipeline(DiffusionPipeline, FromSingleFileMixin, LTX2LoraLoaderMixin): + r""" + Pipeline for text-to-video generation. + + Reference: https://github.com/Lightricks/LTX-Video + + Args: + transformer ([`LTXVideoTransformer3DModel`]): + Conditional Transformer architecture to denoise the encoded video latents. + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + vae ([`AutoencoderKLLTXVideo`]): + Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. + text_encoder ([`T5EncoderModel`]): + [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically + the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant. + tokenizer (`CLIPTokenizer`): + Tokenizer of class + [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer). + tokenizer (`T5TokenizerFast`): + Second Tokenizer of class + [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast). + connectors ([`LTX2TextConnectors`]): + Text connector stack used to adapt text encoder hidden states for the video and audio branches. + """ + + model_cpu_offload_seq = "text_encoder->connectors->transformer->vae->audio_vae->vocoder" + _optional_components = [] + _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] + + def __init__( + self, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKLLTX2Video, + audio_vae: AutoencoderKLLTX2Audio, + text_encoder: Gemma3ForConditionalGeneration, + tokenizer: GemmaTokenizer | GemmaTokenizerFast, + connectors: LTX2TextConnectors, + transformer: LTX2VideoTransformer3DModel, + vocoder: LTX2Vocoder, + ): + super().__init__() + + self.register_modules( + vae=vae, + audio_vae=audio_vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + connectors=connectors, + transformer=transformer, + vocoder=vocoder, + scheduler=scheduler, + ) + + self.vae_spatial_compression_ratio = ( + self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 32 + ) + self.vae_temporal_compression_ratio = ( + self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 8 + ) + # TODO: check whether the MEL compression ratio logic here is corrct + self.audio_vae_mel_compression_ratio = ( + self.audio_vae.mel_compression_ratio if getattr(self, "audio_vae", None) is not None else 4 + ) + self.audio_vae_temporal_compression_ratio = ( + self.audio_vae.temporal_compression_ratio if getattr(self, "audio_vae", None) is not None else 4 + ) + self.transformer_spatial_patch_size = ( + self.transformer.config.patch_size if getattr(self, "transformer", None) is not None else 1 + ) + self.transformer_temporal_patch_size = ( + self.transformer.config.patch_size_t if getattr(self, "transformer") is not None else 1 + ) + + self.audio_sampling_rate = ( + self.audio_vae.config.sample_rate if getattr(self, "audio_vae", None) is not None else 16000 + ) + self.audio_hop_length = ( + self.audio_vae.config.mel_hop_length if getattr(self, "audio_vae", None) is not None else 160 + ) + + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio) + self.tokenizer_max_length = ( + self.tokenizer.model_max_length if getattr(self, "tokenizer", None) is not None else 1024 + ) + + @staticmethod + def _pack_text_embeds( + text_hidden_states: torch.Tensor, + sequence_lengths: torch.Tensor, + device: str | torch.device, + padding_side: str = "left", + scale_factor: int = 8, + eps: float = 1e-6, + ) -> torch.Tensor: + """ + Packs and normalizes text encoder hidden states, respecting padding. Normalization is performed per-batch and + per-layer in a masked fashion (only over non-padded positions). + + Args: + text_hidden_states (`torch.Tensor` of shape `(batch_size, seq_len, hidden_dim, num_layers)`): + Per-layer hidden_states from a text encoder (e.g. `Gemma3ForConditionalGeneration`). + sequence_lengths (`torch.Tensor of shape `(batch_size,)`): + The number of valid (non-padded) tokens for each batch instance. + device: (`str` or `torch.device`, *optional*): + torch device to place the resulting embeddings on + padding_side: (`str`, *optional*, defaults to `"left"`): + Whether the text tokenizer performs padding on the `"left"` or `"right"`. + scale_factor (`int`, *optional*, defaults to `8`): + Scaling factor to multiply the normalized hidden states by. + eps (`float`, *optional*, defaults to `1e-6`): + A small positive value for numerical stability when performing normalization. + + Returns: + `torch.Tensor` of shape `(batch_size, seq_len, hidden_dim * num_layers)`: + Normed and flattened text encoder hidden states. + """ + batch_size, seq_len, hidden_dim, num_layers = text_hidden_states.shape + original_dtype = text_hidden_states.dtype + + # Create padding mask + token_indices = torch.arange(seq_len, device=device).unsqueeze(0) + if padding_side == "right": + # For right padding, valid tokens are from 0 to sequence_length-1 + mask = token_indices < sequence_lengths[:, None] # [batch_size, seq_len] + elif padding_side == "left": + # For left padding, valid tokens are from (T - sequence_length) to T-1 + start_indices = seq_len - sequence_lengths[:, None] # [batch_size, 1] + mask = token_indices >= start_indices # [B, T] + else: + raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}") + mask = mask[:, :, None, None] # [batch_size, seq_len] --> [batch_size, seq_len, 1, 1] + + # Compute masked mean over non-padding positions of shape (batch_size, 1, 1, seq_len) + masked_text_hidden_states = text_hidden_states.masked_fill(~mask, 0.0) + num_valid_positions = (sequence_lengths * hidden_dim).view(batch_size, 1, 1, 1) + masked_mean = masked_text_hidden_states.sum(dim=(1, 2), keepdim=True) / (num_valid_positions + eps) + + # Compute min/max over non-padding positions of shape (batch_size, 1, 1 seq_len) + x_min = text_hidden_states.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True) + x_max = text_hidden_states.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True) + + # Normalization + normalized_hidden_states = (text_hidden_states - masked_mean) / (x_max - x_min + eps) + normalized_hidden_states = normalized_hidden_states * scale_factor + + # Pack the hidden states to a 3D tensor (batch_size, seq_len, hidden_dim * num_layers) + normalized_hidden_states = normalized_hidden_states.flatten(2) + mask_flat = mask.squeeze(-1).expand(-1, -1, hidden_dim * num_layers) + normalized_hidden_states = normalized_hidden_states.masked_fill(~mask_flat, 0.0) + normalized_hidden_states = normalized_hidden_states.to(dtype=original_dtype) + return normalized_hidden_states + + def _get_gemma_prompt_embeds( + self, + prompt: str | list[str], + num_videos_per_prompt: int = 1, + max_sequence_length: int = 1024, + scale_factor: int = 8, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `list[str]`, *optional*): + prompt to be encoded + device: (`str` or `torch.device`): + torch device to place the resulting embeddings on + dtype: (`torch.dtype`): + torch dtype to cast the prompt embeds to + max_sequence_length (`int`, defaults to 1024): Maximum sequence length to use for the prompt. + """ + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + if getattr(self, "tokenizer", None) is not None: + # Gemma expects left padding for chat-style prompts + self.tokenizer.padding_side = "left" + if self.tokenizer.pad_token is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + + prompt = [p.strip() for p in prompt] + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + prompt_attention_mask = text_inputs.attention_mask + text_input_ids = text_input_ids.to(device) + prompt_attention_mask = prompt_attention_mask.to(device) + + text_encoder_outputs = self.text_encoder( + input_ids=text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True + ) + text_encoder_hidden_states = text_encoder_outputs.hidden_states + text_encoder_hidden_states = torch.stack(text_encoder_hidden_states, dim=-1) + sequence_lengths = prompt_attention_mask.sum(dim=-1) + + prompt_embeds = self._pack_text_embeds( + text_encoder_hidden_states, + sequence_lengths, + device=device, + padding_side=self.tokenizer.padding_side, + scale_factor=scale_factor, + ) + prompt_embeds = prompt_embeds.to(dtype=dtype) + + # duplicate text embeddings for each generation per prompt, using mps friendly method + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) + + prompt_attention_mask = prompt_attention_mask.view(batch_size, -1) + prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1) + + return prompt_embeds, prompt_attention_mask + + def encode_prompt( + self, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, + do_classifier_free_guidance: bool = True, + num_videos_per_prompt: int = 1, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + max_sequence_length: int = 1024, + scale_factor: int = 8, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `list[str]`, *optional*): + prompt to be encoded + negative_prompt (`str` or `list[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): + Whether to use classifier free guidance or not. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos that should be generated per prompt. torch device to place the resulting embeddings on + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + device: (`torch.device`, *optional*): + torch device + dtype: (`torch.dtype`, *optional*): + torch dtype + """ + device = device or self._execution_device + + prompt = [prompt] if isinstance(prompt, str) else prompt + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds( + prompt=prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + scale_factor=scale_factor, + device=device, + dtype=dtype, + ) + + if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + + if prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + + negative_prompt_embeds, negative_prompt_attention_mask = self._get_gemma_prompt_embeds( + prompt=negative_prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + scale_factor=scale_factor, + device=device, + dtype=dtype, + ) + + return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask + + def check_inputs( + self, + prompt, + height, + width, + callback_on_step_end_tensor_inputs=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + ): + if height % 32 != 0 or width % 32 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.") + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if prompt_embeds is not None and prompt_attention_mask is None: + raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.") + + if negative_prompt_embeds is not None and negative_prompt_attention_mask is None: + raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.") + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}." + ) + if prompt_attention_mask.shape != negative_prompt_attention_mask.shape: + raise ValueError( + "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but" + f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`" + f" {negative_prompt_attention_mask.shape}." + ) + + @staticmethod + def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor: + # Unpacked latents of shape are [B, C, F, H, W] are patched into tokens of shape [B, C, F // p_t, p_t, H // p, p, W // p, p]. + # The patch dimensions are then permuted and collapsed into the channel dimension of shape: + # [B, F // p_t * H // p * W // p, C * p_t * p * p] (an ndim=3 tensor). + # dim=0 is the batch size, dim=1 is the effective video sequence length, dim=2 is the effective number of input features + batch_size, num_channels, num_frames, height, width = latents.shape + post_patch_num_frames = num_frames // patch_size_t + post_patch_height = height // patch_size + post_patch_width = width // patch_size + latents = latents.reshape( + batch_size, + -1, + post_patch_num_frames, + patch_size_t, + post_patch_height, + patch_size, + post_patch_width, + patch_size, + ) + latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3) + return latents + + @staticmethod + def _unpack_latents( + latents: torch.Tensor, num_frames: int, height: int, width: int, patch_size: int = 1, patch_size_t: int = 1 + ) -> torch.Tensor: + # Packed latents of shape [B, S, D] (S is the effective video sequence length, D is the effective feature dimensions) + # are unpacked and reshaped into a video tensor of shape [B, C, F, H, W]. This is the inverse operation of + # what happens in the `_pack_latents` method. + batch_size = latents.size(0) + latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size) + latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3) + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2_image2video.LTX2ImageToVideoPipeline._normalize_latents + def _normalize_latents( + latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 + ) -> torch.Tensor: + # Normalize latents across the channel dimension [B, C, F, H, W] + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents = (latents - latents_mean) * scaling_factor / latents_std + return latents + + @staticmethod + def _denormalize_latents( + latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 + ) -> torch.Tensor: + # Denormalize latents across the channel dimension [B, C, F, H, W] + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents = latents * latents_std / scaling_factor + latents_mean + return latents + + @staticmethod + def _normalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor): + latents_mean = latents_mean.to(latents.device, latents.dtype) + latents_std = latents_std.to(latents.device, latents.dtype) + return (latents - latents_mean) / latents_std + + @staticmethod + def _denormalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor): + latents_mean = latents_mean.to(latents.device, latents.dtype) + latents_std = latents_std.to(latents.device, latents.dtype) + return (latents * latents_std) + latents_mean + + @staticmethod + def _create_noised_state( + latents: torch.Tensor, noise_scale: float | torch.Tensor, generator: torch.Generator | None = None + ): + noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype) + noised_latents = noise_scale * noise + (1 - noise_scale) * latents + return noised_latents + + @staticmethod + def _pack_audio_latents( + latents: torch.Tensor, patch_size: int | None = None, patch_size_t: int | None = None + ) -> torch.Tensor: + # Audio latents shape: [B, C, L, M], where L is the latent audio length and M is the number of mel bins + if patch_size is not None and patch_size_t is not None: + # Packs the latents into a patch sequence of shape [B, L // p_t * M // p, C * p_t * p] (a ndim=3 tnesor). + # dim=1 is the effective audio sequence length and dim=2 is the effective audio input feature size. + batch_size, num_channels, latent_length, latent_mel_bins = latents.shape + post_patch_latent_length = latent_length / patch_size_t + post_patch_mel_bins = latent_mel_bins / patch_size + latents = latents.reshape( + batch_size, -1, post_patch_latent_length, patch_size_t, post_patch_mel_bins, patch_size + ) + latents = latents.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2) + else: + # Packs the latents into a patch sequence of shape [B, L, C * M]. This implicitly assumes a (mel) + # patch_size of M (all mel bins constitutes a single patch) and a patch_size_t of 1. + latents = latents.transpose(1, 2).flatten(2, 3) # [B, C, L, M] --> [B, L, C * M] + return latents + + @staticmethod + def _unpack_audio_latents( + latents: torch.Tensor, + latent_length: int, + num_mel_bins: int, + patch_size: int | None = None, + patch_size_t: int | None = None, + ) -> torch.Tensor: + # Unpacks an audio patch sequence of shape [B, S, D] into a latent spectrogram tensor of shape [B, C, L, M], + # where L is the latent audio length and M is the number of mel bins. + if patch_size is not None and patch_size_t is not None: + batch_size = latents.size(0) + latents = latents.reshape(batch_size, latent_length, num_mel_bins, -1, patch_size_t, patch_size) + latents = latents.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) + else: + # Assume [B, S, D] = [B, L, C * M], which implies that patch_size = M and patch_size_t = 1. + latents = latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2) + return latents + + def prepare_latents( + self, + batch_size: int = 1, + num_channels_latents: int = 128, + height: int = 512, + width: int = 768, + num_frames: int = 121, + noise_scale: float = 0.0, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + if latents is not None: + if latents.ndim == 5: + latents = self._normalize_latents( + latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor + ) + # latents are of shape [B, C, F, H, W], need to be packed + latents = self._pack_latents( + latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size + ) + if latents.ndim != 3: + raise ValueError( + f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is [batch_size, num_seq, num_features]." + ) + latents = self._create_noised_state(latents, noise_scale, generator) + return latents.to(device=device, dtype=dtype) + + height = height // self.vae_spatial_compression_ratio + width = width // self.vae_spatial_compression_ratio + num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 + + shape = (batch_size, num_channels_latents, num_frames, height, width) + + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = self._pack_latents( + latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size + ) + return latents + + def prepare_audio_latents( + self, + batch_size: int = 1, + num_channels_latents: int = 8, + audio_latent_length: int = 1, # 1 is just a dummy value + num_mel_bins: int = 64, + noise_scale: float = 0.0, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + if latents is not None: + if latents.ndim == 4: + # latents are of shape [B, C, L, M], need to be packed + latents = self._pack_audio_latents(latents) + if latents.ndim != 3: + raise ValueError( + f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is [batch_size, num_seq, num_features]." + ) + latents = self._normalize_audio_latents(latents, self.audio_vae.latents_mean, self.audio_vae.latents_std) + latents = self._create_noised_state(latents, noise_scale, generator) + return latents.to(device=device, dtype=dtype) + + # TODO: confirm whether this logic is correct + latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio + + shape = (batch_size, num_channels_latents, audio_latent_length, latent_mel_bins) + + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = self._pack_audio_latents(latents) + return latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def guidance_rescale(self): + return self._guidance_rescale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: str | list[str] = None, + negative_prompt: str | list[str] | None = None, + height: int = 512, + width: int = 768, + num_frames: int = 121, + frame_rate: float = 24.0, + num_inference_steps: int = 40, + sigmas: list[float] | None = None, + timesteps: list[int] = None, + guidance_scale: float = 4.0, + guidance_rescale: float = 0.0, + noise_scale: float = 0.0, + num_videos_per_prompt: int = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + audio_latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + decode_timestep: float | list[float] = 0.0, + decode_noise_scale: float | list[float] | None = None, + output_type: str = "pil", + return_dict: bool = True, + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + max_sequence_length: int = 1024, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `list[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + height (`int`, *optional*, defaults to `512`): + The height in pixels of the generated image. This is set to 480 by default for the best results. + width (`int`, *optional*, defaults to `768`): + The width in pixels of the generated image. This is set to 848 by default for the best results. + num_frames (`int`, *optional*, defaults to `121`): + The number of video frames to generate + frame_rate (`float`, *optional*, defaults to `24.0`): + The frames per second (FPS) of the generated video. + num_inference_steps (`int`, *optional*, defaults to 40): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + timesteps (`list[int]`, *optional*): + Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument + in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is + passed will be used. Must be in descending order. + guidance_scale (`float`, *optional*, defaults to `4.0`): + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. + of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting + `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to + the text `prompt`, usually at the expense of lower image quality. + guidance_rescale (`float`, *optional*, defaults to 0.0): + Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are + Flawed](https://huggingface.co/papers/2305.08891) `guidance_scale` is defined as `φ` in equation 16. of + [Common Diffusion Noise Schedules and Sample Steps are + Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when + using zero terminal SNR. + noise_scale (`float`, *optional*, defaults to `0.0`): + The interpolation factor between random noise and denoised latents at each timestep. Applying noise to + the `latents` and `audio_latents` before continue denoising. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of videos to generate per prompt. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for video + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + audio_latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for audio + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + prompt_attention_mask (`torch.Tensor`, *optional*): + Pre-generated attention mask for text embeddings. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not + provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. + negative_prompt_attention_mask (`torch.FloatTensor`, *optional*): + Pre-generated attention mask for negative text embeddings. + decode_timestep (`float`, defaults to `0.0`): + The timestep at which generated video is decoded. + decode_noise_scale (`float`, defaults to `None`): + The interpolation factor between random noise and denoised latents at the decode timestep. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.ltx.LTX2PipelineOutput`] instead of a plain tuple. + attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*, defaults to `["latents"]`): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int`, *optional*, defaults to `1024`): + Maximum sequence length to use with the `prompt`. + + Examples: + + Returns: + [`~pipelines.ltx.LTX2PipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`~pipelines.ltx.LTX2PipelineOutput`] is returned, otherwise a `tuple` is + returned where the first element is a list with the generated images. + """ + + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt=prompt, + height=height, + width=width, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + ) + + self._guidance_scale = guidance_scale + self._guidance_rescale = guidance_rescale + self._attention_kwargs = attention_kwargs + self._interrupt = False + self._current_timestep = None + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + + # 3. Prepare text embeddings + ( + prompt_embeds, + prompt_attention_mask, + negative_prompt_embeds, + negative_prompt_attention_mask, + ) = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + num_videos_per_prompt=num_videos_per_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + max_sequence_length=max_sequence_length, + device=device, + ) + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) + + additive_attention_mask = (1 - prompt_attention_mask.to(prompt_embeds.dtype)) * -1000000.0 + connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors( + prompt_embeds, additive_attention_mask, additive_mask=True + ) + + # 4. Prepare latent variables + latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 + latent_height = height // self.vae_spatial_compression_ratio + latent_width = width // self.vae_spatial_compression_ratio + if latents is not None: + if latents.ndim == 5: + logger.info( + "Got latents of shape [batch_size, latent_dim, latent_frames, latent_height, latent_width], `latent_num_frames`, `latent_height`, `latent_width` will be inferred." + ) + _, _, latent_num_frames, latent_height, latent_width = latents.shape # [B, C, F, H, W] + elif latents.ndim == 3: + logger.warning( + f"You have supplied packed `latents` of shape {latents.shape}, so the latent dims cannot be" + f" inferred. Make sure the supplied `height`, `width`, and `num_frames` are correct." + ) + else: + raise ValueError( + f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is either [batch_size, seq_len, num_features] or [batch_size, latent_dim, latent_frames, latent_height, latent_width]." + ) + video_sequence_length = latent_num_frames * latent_height * latent_width + + num_channels_latents = self.transformer.config.in_channels + latents = self.prepare_latents( + batch_size * num_videos_per_prompt, + num_channels_latents, + height, + width, + num_frames, + noise_scale, + torch.float32, + device, + generator, + latents, + ) + + duration_s = num_frames / frame_rate + audio_latents_per_second = ( + self.audio_sampling_rate / self.audio_hop_length / float(self.audio_vae_temporal_compression_ratio) + ) + audio_num_frames = round(duration_s * audio_latents_per_second) + if audio_latents is not None: + if audio_latents.ndim == 4: + logger.info( + "Got audio_latents of shape [batch_size, num_channels, audio_length, mel_bins], `audio_num_frames` will be inferred." + ) + _, _, audio_num_frames, _ = audio_latents.shape # [B, C, L, M] + elif audio_latents.ndim == 3: + logger.warning( + f"You have supplied packed `audio_latents` of shape {audio_latents.shape}, so the latent dims" + f" cannot be inferred. Make sure the supplied `num_frames` and `frame_rate` are correct." + ) + else: + raise ValueError( + f"Provided `audio_latents` tensor has shape {audio_latents.shape}, but the expected shape is either [batch_size, seq_len, num_features] or [batch_size, num_channels, audio_length, mel_bins]." + ) + + num_mel_bins = self.audio_vae.config.mel_bins if getattr(self, "audio_vae", None) is not None else 64 + latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio + num_channels_latents_audio = ( + self.audio_vae.config.latent_channels if getattr(self, "audio_vae", None) is not None else 8 + ) + audio_latents = self.prepare_audio_latents( + batch_size * num_videos_per_prompt, + num_channels_latents=num_channels_latents_audio, + audio_latent_length=audio_num_frames, + num_mel_bins=num_mel_bins, + noise_scale=noise_scale, + dtype=torch.float32, + device=device, + generator=generator, + latents=audio_latents, + ) + + # 5. Prepare timesteps + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas + mu = calculate_shift( + video_sequence_length, + self.scheduler.config.get("base_image_seq_len", 1024), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.95), + self.scheduler.config.get("max_shift", 2.05), + ) + # For now, duplicate the scheduler for use with the audio latents + audio_scheduler = copy.deepcopy(self.scheduler) + _, _ = retrieve_timesteps( + audio_scheduler, + num_inference_steps, + device, + timesteps, + sigmas=sigmas, + mu=mu, + ) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + timesteps, + sigmas=sigmas, + mu=mu, + ) + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + # 6. Prepare micro-conditions + rope_interpolation_scale = ( + self.vae_temporal_compression_ratio / frame_rate, + self.vae_spatial_compression_ratio, + self.vae_spatial_compression_ratio, + ) + # Pre-compute video and audio positional ids as they will be the same at each step of the denoising loop + video_coords = self.transformer.rope.prepare_video_coords( + latents.shape[0], latent_num_frames, latent_height, latent_width, latents.device, fps=frame_rate + ) + audio_coords = self.transformer.audio_rope.prepare_audio_coords( + audio_latents.shape[0], audio_num_frames, audio_latents.device + ) + # Duplicate the positional ids as well if using CFG + if self.do_classifier_free_guidance: + video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1)) # Repeat twice in batch dim + audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1)) + + # 7. Denoising loop + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + self._current_timestep = t + + latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents + latent_model_input = latent_model_input.to(prompt_embeds.dtype) + audio_latent_model_input = ( + torch.cat([audio_latents] * 2) if self.do_classifier_free_guidance else audio_latents + ) + audio_latent_model_input = audio_latent_model_input.to(prompt_embeds.dtype) + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latent_model_input.shape[0]) + + with self.transformer.cache_context("cond_uncond"): + noise_pred_video, noise_pred_audio = self.transformer( + hidden_states=latent_model_input, + audio_hidden_states=audio_latent_model_input, + encoder_hidden_states=connector_prompt_embeds, + audio_encoder_hidden_states=connector_audio_prompt_embeds, + timestep=timestep, + encoder_attention_mask=connector_attention_mask, + audio_encoder_attention_mask=connector_attention_mask, + num_frames=latent_num_frames, + height=latent_height, + width=latent_width, + fps=frame_rate, + audio_num_frames=audio_num_frames, + video_coords=video_coords, + audio_coords=audio_coords, + # rope_interpolation_scale=rope_interpolation_scale, + attention_kwargs=attention_kwargs, + return_dict=False, + ) + noise_pred_video = noise_pred_video.float() + noise_pred_audio = noise_pred_audio.float() + + if self.do_classifier_free_guidance: + noise_pred_video_uncond, noise_pred_video_text = noise_pred_video.chunk(2) + noise_pred_video = noise_pred_video_uncond + self.guidance_scale * ( + noise_pred_video_text - noise_pred_video_uncond + ) + + noise_pred_audio_uncond, noise_pred_audio_text = noise_pred_audio.chunk(2) + noise_pred_audio = noise_pred_audio_uncond + self.guidance_scale * ( + noise_pred_audio_text - noise_pred_audio_uncond + ) + + if self.guidance_rescale > 0: + # Based on 3.4. in https://huggingface.co/papers/2305.08891 + noise_pred_video = rescale_noise_cfg( + noise_pred_video, noise_pred_video_text, guidance_rescale=self.guidance_rescale + ) + noise_pred_audio = rescale_noise_cfg( + noise_pred_audio, noise_pred_audio_text, guidance_rescale=self.guidance_rescale + ) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0] + # NOTE: for now duplicate scheduler for audio latents in case self.scheduler sets internal state in + # the step method (such as _step_index) + audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0] + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + latents = self._unpack_latents( + latents, + latent_num_frames, + latent_height, + latent_width, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) + latents = self._denormalize_latents( + latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor + ) + + audio_latents = self._denormalize_audio_latents( + audio_latents, self.audio_vae.latents_mean, self.audio_vae.latents_std + ) + audio_latents = self._unpack_audio_latents(audio_latents, audio_num_frames, num_mel_bins=latent_mel_bins) + + if output_type == "latent": + video = latents + audio = audio_latents + else: + latents = latents.to(prompt_embeds.dtype) + + if not self.vae.config.timestep_conditioning: + timestep = None + else: + noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype) + if not isinstance(decode_timestep, list): + decode_timestep = [decode_timestep] * batch_size + if decode_noise_scale is None: + decode_noise_scale = decode_timestep + elif not isinstance(decode_noise_scale, list): + decode_noise_scale = [decode_noise_scale] * batch_size + + timestep = torch.tensor(decode_timestep, device=device, dtype=latents.dtype) + decode_noise_scale = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[ + :, None, None, None, None + ] + latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise + + latents = latents.to(self.vae.dtype) + video = self.vae.decode(latents, timestep, return_dict=False)[0] + video = self.video_processor.postprocess_video(video, output_type=output_type) + + audio_latents = audio_latents.to(self.audio_vae.dtype) + generated_mel_spectrograms = self.audio_vae.decode(audio_latents, return_dict=False)[0] + audio = self.vocoder(generated_mel_spectrograms) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (video, audio) + + return LTX2PipelineOutput(frames=video, audio=audio) diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py new file mode 100644 index 000000000000..4c451330f439 --- /dev/null +++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py @@ -0,0 +1,1474 @@ +# Copyright 2025 Lightricks and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import inspect +from dataclasses import dataclass +from typing import Any, Callable + +import numpy as np +import PIL.Image +import torch +from transformers import Gemma3ForConditionalGeneration, GemmaTokenizer, GemmaTokenizerFast + +from ...callbacks import MultiPipelineCallbacks, PipelineCallback +from ...loaders import FromSingleFileMixin, LTX2LoraLoaderMixin +from ...models.autoencoders import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video +from ...models.transformers import LTX2VideoTransformer3DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import is_torch_xla_available, logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor +from ..pipeline_utils import DiffusionPipeline +from .connectors import LTX2TextConnectors +from .pipeline_output import LTX2PipelineOutput +from .vocoder import LTX2Vocoder + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import LTX2ConditionPipeline + >>> from diffusers.pipelines.ltx2.export_utils import encode_video + >>> from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition + >>> from diffusers.utils import load_image + + >>> pipe = LTX2ConditionPipeline.from_pretrained("Lightricks/LTX-2", torch_dtype=torch.bfloat16) + >>> pipe.enable_model_cpu_offload() + + >>> first_image = load_image( + ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_first_frame.png" + ... ) + >>> last_image = load_image( + ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_last_frame.png" + ... ) + >>> first_cond = LTX2VideoCondition(frames=first_image, index=0, strength=1.0) + >>> last_cond = LTX2VideoCondition(frames=last_image, index=-1, strength=1.0) + >>> conditions = [first_cond, last_cond] + >>> prompt = "CG animation style, a small blue bird takes off from the ground, flapping its wings." + >>> negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted, static" + + >>> frame_rate = 24.0 + >>> video = pipe( + ... conditions=conditions, + ... prompt=prompt, + ... negative_prompt=negative_prompt, + ... width=768, + ... height=512, + ... num_frames=121, + ... frame_rate=frame_rate, + ... num_inference_steps=40, + ... guidance_scale=4.0, + ... output_type="np", + ... return_dict=False, + ... ) + >>> video = (video * 255).round().astype("uint8") + >>> video = torch.from_numpy(video) + + >>> encode_video( + ... video[0], + ... fps=frame_rate, + ... audio=audio[0].float().cpu(), + ... audio_sample_rate=pipe.vocoder.config.output_sampling_rate, # should be 24000 + ... output_path="video.mp4", + ... ) + ``` +""" + + +@dataclass +class LTX2VideoCondition: + """ + Defines a single frame-conditioning item for LTX-2 Video - a single frame or a sequence of frames. + + Attributes: + frames (`PIL.Image.Image` or `List[PIL.Image.Image]` or `np.ndarray` or `torch.Tensor`): + The image (or video) to condition the video on. Accepts any type that can be handled by + VideoProcessor.preprocess_video. + index (`int`, defaults to `0`): + The index at which the image or video will conditionally affect the video generation. + strength (`float`, defaults to `1.0`): + The strength of the conditioning effect. A value of `1.0` means the conditioning effect is fully applied. + """ + + frames: PIL.Image.Image | list[PIL.Image.Image] | np.ndarray | torch.Tensor + index: int = 0 + strength: float = 1.0 + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents +def retrieve_latents( + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" +): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`list[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`list[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg +def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): + r""" + Rescales `noise_cfg` tensor based on `guidance_rescale` to improve image quality and fix overexposure. Based on + Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are + Flawed](https://huggingface.co/papers/2305.08891). + + Args: + noise_cfg (`torch.Tensor`): + The predicted noise tensor for the guided diffusion process. + noise_pred_text (`torch.Tensor`): + The predicted noise tensor for the text-guided diffusion process. + guidance_rescale (`float`, *optional*, defaults to 0.0): + A rescale factor applied to the noise predictions. + + Returns: + noise_cfg (`torch.Tensor`): The rescaled noise prediction tensor. + """ + std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True) + std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) + # rescale the results from guidance (fixes overexposure) + noise_pred_rescaled = noise_cfg * (std_text / std_cfg) + # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images + noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg + return noise_cfg + + +class LTX2ConditionPipeline(DiffusionPipeline, FromSingleFileMixin, LTX2LoraLoaderMixin): + r""" + Pipeline for video generation which allows image conditions to be inserted at arbitary parts of the video. + + Reference: https://github.com/Lightricks/LTX-Video + + TODO + """ + + model_cpu_offload_seq = "text_encoder->connectors->transformer->vae->audio_vae->vocoder" + _optional_components = [] + _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] + + def __init__( + self, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKLLTX2Video, + audio_vae: AutoencoderKLLTX2Audio, + text_encoder: Gemma3ForConditionalGeneration, + tokenizer: GemmaTokenizer | GemmaTokenizerFast, + connectors: LTX2TextConnectors, + transformer: LTX2VideoTransformer3DModel, + vocoder: LTX2Vocoder, + ): + super().__init__() + + self.register_modules( + vae=vae, + audio_vae=audio_vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + connectors=connectors, + transformer=transformer, + vocoder=vocoder, + scheduler=scheduler, + ) + + self.vae_spatial_compression_ratio = ( + self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 32 + ) + self.vae_temporal_compression_ratio = ( + self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 8 + ) + self.audio_vae_mel_compression_ratio = ( + self.audio_vae.mel_compression_ratio if getattr(self, "audio_vae", None) is not None else 4 + ) + self.audio_vae_temporal_compression_ratio = ( + self.audio_vae.temporal_compression_ratio if getattr(self, "audio_vae", None) is not None else 4 + ) + self.transformer_spatial_patch_size = ( + self.transformer.config.patch_size if getattr(self, "transformer", None) is not None else 1 + ) + self.transformer_temporal_patch_size = ( + self.transformer.config.patch_size_t if getattr(self, "transformer") is not None else 1 + ) + + self.audio_sampling_rate = ( + self.audio_vae.config.sample_rate if getattr(self, "audio_vae", None) is not None else 16000 + ) + self.audio_hop_length = ( + self.audio_vae.config.mel_hop_length if getattr(self, "audio_vae", None) is not None else 160 + ) + + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio, resample="bilinear") + self.tokenizer_max_length = ( + self.tokenizer.model_max_length if getattr(self, "tokenizer", None) is not None else 1024 + ) + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._pack_text_embeds + def _pack_text_embeds( + text_hidden_states: torch.Tensor, + sequence_lengths: torch.Tensor, + device: str | torch.device, + padding_side: str = "left", + scale_factor: int = 8, + eps: float = 1e-6, + ) -> torch.Tensor: + """ + Packs and normalizes text encoder hidden states, respecting padding. Normalization is performed per-batch and + per-layer in a masked fashion (only over non-padded positions). + + Args: + text_hidden_states (`torch.Tensor` of shape `(batch_size, seq_len, hidden_dim, num_layers)`): + Per-layer hidden_states from a text encoder (e.g. `Gemma3ForConditionalGeneration`). + sequence_lengths (`torch.Tensor of shape `(batch_size,)`): + The number of valid (non-padded) tokens for each batch instance. + device: (`str` or `torch.device`, *optional*): + torch device to place the resulting embeddings on + padding_side: (`str`, *optional*, defaults to `"left"`): + Whether the text tokenizer performs padding on the `"left"` or `"right"`. + scale_factor (`int`, *optional*, defaults to `8`): + Scaling factor to multiply the normalized hidden states by. + eps (`float`, *optional*, defaults to `1e-6`): + A small positive value for numerical stability when performing normalization. + + Returns: + `torch.Tensor` of shape `(batch_size, seq_len, hidden_dim * num_layers)`: + Normed and flattened text encoder hidden states. + """ + batch_size, seq_len, hidden_dim, num_layers = text_hidden_states.shape + original_dtype = text_hidden_states.dtype + + # Create padding mask + token_indices = torch.arange(seq_len, device=device).unsqueeze(0) + if padding_side == "right": + # For right padding, valid tokens are from 0 to sequence_length-1 + mask = token_indices < sequence_lengths[:, None] # [batch_size, seq_len] + elif padding_side == "left": + # For left padding, valid tokens are from (T - sequence_length) to T-1 + start_indices = seq_len - sequence_lengths[:, None] # [batch_size, 1] + mask = token_indices >= start_indices # [B, T] + else: + raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}") + mask = mask[:, :, None, None] # [batch_size, seq_len] --> [batch_size, seq_len, 1, 1] + + # Compute masked mean over non-padding positions of shape (batch_size, 1, 1, seq_len) + masked_text_hidden_states = text_hidden_states.masked_fill(~mask, 0.0) + num_valid_positions = (sequence_lengths * hidden_dim).view(batch_size, 1, 1, 1) + masked_mean = masked_text_hidden_states.sum(dim=(1, 2), keepdim=True) / (num_valid_positions + eps) + + # Compute min/max over non-padding positions of shape (batch_size, 1, 1 seq_len) + x_min = text_hidden_states.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True) + x_max = text_hidden_states.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True) + + # Normalization + normalized_hidden_states = (text_hidden_states - masked_mean) / (x_max - x_min + eps) + normalized_hidden_states = normalized_hidden_states * scale_factor + + # Pack the hidden states to a 3D tensor (batch_size, seq_len, hidden_dim * num_layers) + normalized_hidden_states = normalized_hidden_states.flatten(2) + mask_flat = mask.squeeze(-1).expand(-1, -1, hidden_dim * num_layers) + normalized_hidden_states = normalized_hidden_states.masked_fill(~mask_flat, 0.0) + normalized_hidden_states = normalized_hidden_states.to(dtype=original_dtype) + return normalized_hidden_states + + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._get_gemma_prompt_embeds + def _get_gemma_prompt_embeds( + self, + prompt: str | list[str], + num_videos_per_prompt: int = 1, + max_sequence_length: int = 1024, + scale_factor: int = 8, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `list[str]`, *optional*): + prompt to be encoded + device: (`str` or `torch.device`): + torch device to place the resulting embeddings on + dtype: (`torch.dtype`): + torch dtype to cast the prompt embeds to + max_sequence_length (`int`, defaults to 1024): Maximum sequence length to use for the prompt. + """ + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + if getattr(self, "tokenizer", None) is not None: + # Gemma expects left padding for chat-style prompts + self.tokenizer.padding_side = "left" + if self.tokenizer.pad_token is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + + prompt = [p.strip() for p in prompt] + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + prompt_attention_mask = text_inputs.attention_mask + text_input_ids = text_input_ids.to(device) + prompt_attention_mask = prompt_attention_mask.to(device) + + text_encoder_outputs = self.text_encoder( + input_ids=text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True + ) + text_encoder_hidden_states = text_encoder_outputs.hidden_states + text_encoder_hidden_states = torch.stack(text_encoder_hidden_states, dim=-1) + sequence_lengths = prompt_attention_mask.sum(dim=-1) + + prompt_embeds = self._pack_text_embeds( + text_encoder_hidden_states, + sequence_lengths, + device=device, + padding_side=self.tokenizer.padding_side, + scale_factor=scale_factor, + ) + prompt_embeds = prompt_embeds.to(dtype=dtype) + + # duplicate text embeddings for each generation per prompt, using mps friendly method + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) + + prompt_attention_mask = prompt_attention_mask.view(batch_size, -1) + prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1) + + return prompt_embeds, prompt_attention_mask + + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline.encode_prompt + def encode_prompt( + self, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, + do_classifier_free_guidance: bool = True, + num_videos_per_prompt: int = 1, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + max_sequence_length: int = 1024, + scale_factor: int = 8, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `list[str]`, *optional*): + prompt to be encoded + negative_prompt (`str` or `list[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): + Whether to use classifier free guidance or not. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos that should be generated per prompt. torch device to place the resulting embeddings on + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + device: (`torch.device`, *optional*): + torch device + dtype: (`torch.dtype`, *optional*): + torch dtype + """ + device = device or self._execution_device + + prompt = [prompt] if isinstance(prompt, str) else prompt + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds( + prompt=prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + scale_factor=scale_factor, + device=device, + dtype=dtype, + ) + + if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + + if prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + + negative_prompt_embeds, negative_prompt_attention_mask = self._get_gemma_prompt_embeds( + prompt=negative_prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + scale_factor=scale_factor, + device=device, + dtype=dtype, + ) + + return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask + + def check_inputs( + self, + prompt, + height, + width, + callback_on_step_end_tensor_inputs=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + latents=None, + audio_latents=None, + ): + if height % 32 != 0 or width % 32 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.") + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if prompt_embeds is not None and prompt_attention_mask is None: + raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.") + + if negative_prompt_embeds is not None and negative_prompt_attention_mask is None: + raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.") + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}." + ) + if prompt_attention_mask.shape != negative_prompt_attention_mask.shape: + raise ValueError( + "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but" + f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`" + f" {negative_prompt_attention_mask.shape}." + ) + + if latents is not None and latents.ndim != 5: + raise ValueError( + f"Only unpacked (5D) video latents of shape `[batch_size, latent_channels, latent_frames," + f" latent_height, latent_width] are supported, but got {latents.ndim} dims. If you have packed (3D)" + f" latents, please unpack them (e.g. using the `_unpack_latents` method)." + ) + if audio_latents is not None and audio_latents.ndim != 4: + raise ValueError( + f"Only unpacked (4D) audio latents of shape `[batch_size, num_channels, audio_length, mel_bins] are" + f" supported, but got {latents.ndim} dims. If you have packed (3D) latents, please unpack them (e.g." + f" using the `_unpack_audio_latents` method)." + ) + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._pack_latents + def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor: + # Unpacked latents of shape are [B, C, F, H, W] are patched into tokens of shape [B, C, F // p_t, p_t, H // p, p, W // p, p]. + # The patch dimensions are then permuted and collapsed into the channel dimension of shape: + # [B, F // p_t * H // p * W // p, C * p_t * p * p] (an ndim=3 tensor). + # dim=0 is the batch size, dim=1 is the effective video sequence length, dim=2 is the effective number of input features + batch_size, num_channels, num_frames, height, width = latents.shape + post_patch_num_frames = num_frames // patch_size_t + post_patch_height = height // patch_size + post_patch_width = width // patch_size + latents = latents.reshape( + batch_size, + -1, + post_patch_num_frames, + patch_size_t, + post_patch_height, + patch_size, + post_patch_width, + patch_size, + ) + latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3) + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._unpack_latents + def _unpack_latents( + latents: torch.Tensor, num_frames: int, height: int, width: int, patch_size: int = 1, patch_size_t: int = 1 + ) -> torch.Tensor: + # Packed latents of shape [B, S, D] (S is the effective video sequence length, D is the effective feature dimensions) + # are unpacked and reshaped into a video tensor of shape [B, C, F, H, W]. This is the inverse operation of + # what happens in the `_pack_latents` method. + batch_size = latents.size(0) + latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size) + latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3) + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2_image2video.LTX2ImageToVideoPipeline._normalize_latents + def _normalize_latents( + latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 + ) -> torch.Tensor: + # Normalize latents across the channel dimension [B, C, F, H, W] + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents = (latents - latents_mean) * scaling_factor / latents_std + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._denormalize_latents + def _denormalize_latents( + latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 + ) -> torch.Tensor: + # Denormalize latents across the channel dimension [B, C, F, H, W] + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents = latents * latents_std / scaling_factor + latents_mean + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._create_noised_state + def _create_noised_state( + latents: torch.Tensor, noise_scale: float | torch.Tensor, generator: torch.Generator | None = None + ): + noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype) + noised_latents = noise_scale * noise + (1 - noise_scale) * latents + return noised_latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._pack_audio_latents + def _pack_audio_latents( + latents: torch.Tensor, patch_size: int | None = None, patch_size_t: int | None = None + ) -> torch.Tensor: + # Audio latents shape: [B, C, L, M], where L is the latent audio length and M is the number of mel bins + if patch_size is not None and patch_size_t is not None: + # Packs the latents into a patch sequence of shape [B, L // p_t * M // p, C * p_t * p] (a ndim=3 tnesor). + # dim=1 is the effective audio sequence length and dim=2 is the effective audio input feature size. + batch_size, num_channels, latent_length, latent_mel_bins = latents.shape + post_patch_latent_length = latent_length / patch_size_t + post_patch_mel_bins = latent_mel_bins / patch_size + latents = latents.reshape( + batch_size, -1, post_patch_latent_length, patch_size_t, post_patch_mel_bins, patch_size + ) + latents = latents.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2) + else: + # Packs the latents into a patch sequence of shape [B, L, C * M]. This implicitly assumes a (mel) + # patch_size of M (all mel bins constitutes a single patch) and a patch_size_t of 1. + latents = latents.transpose(1, 2).flatten(2, 3) # [B, C, L, M] --> [B, L, C * M] + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._unpack_audio_latents + def _unpack_audio_latents( + latents: torch.Tensor, + latent_length: int, + num_mel_bins: int, + patch_size: int | None = None, + patch_size_t: int | None = None, + ) -> torch.Tensor: + # Unpacks an audio patch sequence of shape [B, S, D] into a latent spectrogram tensor of shape [B, C, L, M], + # where L is the latent audio length and M is the number of mel bins. + if patch_size is not None and patch_size_t is not None: + batch_size = latents.size(0) + latents = latents.reshape(batch_size, latent_length, num_mel_bins, -1, patch_size_t, patch_size) + latents = latents.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) + else: + # Assume [B, S, D] = [B, L, C * M], which implies that patch_size = M and patch_size_t = 1. + latents = latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2) + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._normalize_audio_latents + def _normalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor): + latents_mean = latents_mean.to(latents.device, latents.dtype) + latents_std = latents_std.to(latents.device, latents.dtype) + return (latents - latents_mean) / latents_std + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._denormalize_audio_latents + def _denormalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor): + latents_mean = latents_mean.to(latents.device, latents.dtype) + latents_std = latents_std.to(latents.device, latents.dtype) + return (latents * latents_std) + latents_mean + + # Copied from diffusers.pipelines.ltx.pipeline_ltx_condition.LTXConditionPipeline.trim_conditioning_sequence + def trim_conditioning_sequence(self, start_frame: int, sequence_num_frames: int, target_num_frames: int) -> int: + """ + Trim a conditioning sequence to the allowed number of frames. + + Args: + start_frame (int): The target frame number of the first frame in the sequence. + sequence_num_frames (int): The number of frames in the sequence. + target_num_frames (int): The target number of frames in the generated video. + Returns: + int: updated sequence length + """ + scale_factor = self.vae_temporal_compression_ratio + num_frames = min(sequence_num_frames, target_num_frames - start_frame) + # Trim down to a multiple of temporal_scale_factor frames plus 1 + num_frames = (num_frames - 1) // scale_factor * scale_factor + 1 + return num_frames + + def preprocess_conditions( + self, + conditions: LTX2VideoCondition | list[LTX2VideoCondition] | None = None, + height: int = 512, + width: int = 768, + num_frames: int = 121, + device: torch.device | None = None, + ) -> tuple[list[torch.Tensor], list[float], list[int]]: + """ + Preprocesses the condition images/videos to torch tensors. + + Args: + conditions (`LTX2VideoCondition` or `List[LTX2VideoCondition]`, *optional*, defaults to `None`): + A list of image/video condition instances. + height (`int`, *optional*, defaults to `512`): + The desired height in pixels. + width (`int`, *optional*, defaults to `768`): + The desired width in pixels. + num_frames (`int`, *optional*, defaults to `121`): + The desired number of frames in the generated video. + device (`torch.device`, *optional*, defaults to `None`): + The device on which to put the preprocessed image/video tensors. + + Returns: + `Tuple[List[torch.Tensor], List[float], List[int]]`: + Returns a 3-tuple of lists of length `len(conditions)` as follows: + 1. The first list is a list of preprocessed video tensors of shape [batch_size=1, num_channels, + num_frames, height, width]. + 2. The second list is a list of conditioning strengths. + 3. The third list is a list of indices in latent space to insert the corresponding condition. + """ + conditioning_frames, conditioning_strengths, conditioning_indices = [], [], [] + + if conditions is None: + conditions = [] + if isinstance(conditions, LTX2VideoCondition): + conditions = [conditions] + + frame_scale_factor = self.vae_temporal_compression_ratio + latent_num_frames = (num_frames - 1) // frame_scale_factor + 1 + for i, condition in enumerate(conditions): + if isinstance(condition.frames, PIL.Image.Image): + # Single image, convert to List[PIL.Image.Image] + video_like_cond = [condition.frames] + elif isinstance(condition.frames, np.ndarray) and condition.frames.ndim == 3: + # Image-like ndarray of shape (H, W, C), insert frame dim in first axis + video_like_cond = np.expand_dims(condition.frames, axis=0) + elif isinstance(condition.frames, torch.Tensor) and condition.frames.ndim == 3: + # Image-like tensor of shape (C, H, W), insert frame dim in first dim + video_like_cond = condition.frames.unsqueeze(0) + else: + # Treat all other as videos. Note that this means 4D ndarrays and tensors will be treated as videos of + # shape (F, H, W, C) and (F, C, H, W), respectively. + video_like_cond = condition.frames + condition_pixels = self.video_processor.preprocess_video( + video_like_cond, height, width, resize_mode="crop" + ) + + # Interpret the index as a latent index, following the original LTX-2 code. + latent_start_idx = condition.index + # Support negative latent indices (e.g. -1 for the last latent index) + if latent_start_idx < 0: + # latent_start_idx will be positive because latent_num_frames is positive + latent_start_idx = latent_start_idx % latent_num_frames + if latent_start_idx >= latent_num_frames: + logger.warning( + f"The starting latent index {latent_start_idx} of condition {i} is too big for the specified number" + f" of latent frames {latent_num_frames}. This condition will be skipped." + ) + continue + + cond_num_frames = condition_pixels.size(2) + start_idx = max((latent_start_idx - 1) * frame_scale_factor + 1, 0) + truncated_cond_frames = self.trim_conditioning_sequence(start_idx, cond_num_frames, num_frames) + condition_pixels = condition_pixels[:, :, :truncated_cond_frames] + + conditioning_frames.append(condition_pixels.to(dtype=self.vae.dtype, device=device)) + conditioning_strengths.append(condition.strength) + conditioning_indices.append(latent_start_idx) + + return conditioning_frames, conditioning_strengths, conditioning_indices + + def apply_visual_conditioning( + self, + latents: torch.Tensor, + conditioning_mask: torch.Tensor, + condition_latents: list[torch.Tensor], + condition_strengths: list[float], + condition_indices: list[int], + latent_height: int, + latent_width: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Applies visual conditioning frames to an initial latent. + + Args: + latents (`torch.Tensor`): + Initial packed (patchified) latents of shape [batch_size, patch_seq_len, hidden_dim]. + conditioning_mask (`torch.Tensor`, *optional*): + Initial packed (patchified) conditioning mask of shape [batch_size, patch_seq_len, 1] with values in + [0, 1] where 0 means that the denoising model output will be fully used and 1 means that the condition + will be fully used (with intermediate values specifying a blend of the denoised and latent values). + + Returns: + `Tuple[torch.Tensor, torch.Tensor, torch.Tensor]`: + Returns a 3-tuple of tensors where: + 1. The first element is the packed video latents (with unchanged shape [batch_size, patch_seq_len, + hidden_dim]) with the conditions applied + 2. The second element is the packed conditioning mask with conditioning strengths applied + 3. The third element holds the clean conditioning latents. + """ + # Latents-like tensor which holds the clean conditioning latents + clean_latents = torch.zeros_like(latents) + for cond, strength, latent_idx in zip(condition_latents, condition_strengths, condition_indices): + num_cond_tokens = cond.size(1) + start_token_idx = latent_idx * latent_height * latent_width + end_token_idx = start_token_idx + num_cond_tokens + + # Overwrite the portion of latents starting with start_token_idx with the condition + latents[:, start_token_idx:end_token_idx] = cond + conditioning_mask[:, start_token_idx:end_token_idx] = strength + clean_latents[:, start_token_idx:end_token_idx] = cond + + return latents, conditioning_mask, clean_latents + + def prepare_latents( + self, + conditions: LTX2VideoCondition | list[LTX2VideoCondition] | None = None, + batch_size: int = 1, + num_channels_latents: int = 128, + height: int = 512, + width: int = 768, + num_frames: int = 121, + noise_scale: float = 1.0, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + latent_height = height // self.vae_spatial_compression_ratio + latent_width = width // self.vae_spatial_compression_ratio + latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 + + shape = (batch_size, num_channels_latents, latent_num_frames, latent_height, latent_width) + mask_shape = (batch_size, 1, latent_num_frames, latent_height, latent_width) + + if latents is not None: + # Latents are expected to be unpacked (5D) with shape [B, F, C, H, W] + latents = self._normalize_latents( + latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor + ) + else: + # NOTE: we set the initial latents to zeros rather a sample from the standard Gaussian prior because we + # will sample from the prior later once we have calculated the conditioning mask + latents = torch.zeros(shape, device=device, dtype=dtype) + + conditioning_mask = latents.new_zeros(mask_shape) + latents = self._pack_latents( + latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size + ) + conditioning_mask = self._pack_latents( + conditioning_mask, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size + ) # [B, seq_len, 1] + + if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape[:2]: + raise ValueError( + f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is {conditioning_mask.shape[:2] + (num_channels_latents,)}." + ) + + if isinstance(generator, list): + logger.warning( + f"{self.__class__.__name__} does not support using a list of generators. The first generator in the" + f" list will be used for all (pseudo-)random operations." + ) + generator = generator[0] + + condition_frames, condition_strengths, condition_indices = self.preprocess_conditions( + conditions, height, width, num_frames, device=device + ) + condition_latents = [] + for condition_tensor in condition_frames: + condition_latent = retrieve_latents( + self.vae.encode(condition_tensor), generator=generator, sample_mode="argmax" + ) + condition_latent = self._normalize_latents( + condition_latent, self.vae.latents_mean, self.vae.latents_std + ).to(device=device, dtype=dtype) + condition_latent = self._pack_latents( + condition_latent, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size + ) + condition_latents.append(condition_latent) + + # NOTE: following the I2V pipeline, we return a conditioning mask. The original LTX 2 code uses a denoising + # mask, which is the inverse of the conditioning mask (`denoise_mask = 1 - conditioning_mask`) + latents, conditioning_mask, clean_latents = self.apply_visual_conditioning( + latents, + conditioning_mask, + condition_latents, + condition_strengths, + condition_indices, + latent_height=latent_height, + latent_width=latent_width, + ) + + # Sample from the standard Gaussian prior (or an intermediate Gaussian distribution if noise_scale < 1.0). + noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype) + scaled_mask = (1.0 - conditioning_mask) * noise_scale + # Add noise to the `latents` so that it is at the noise level specified by `noise_scale`. + latents = noise * scaled_mask + latents * (1 - scaled_mask) + + return latents, conditioning_mask, clean_latents + + def prepare_audio_latents( + self, + batch_size: int = 1, + num_channels_latents: int = 8, + audio_latent_length: int = 1, # 1 is just a dummy value + num_mel_bins: int = 64, + noise_scale: float = 0.0, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + if latents is not None: + # latents expected to be unpacked (4D) with shape [B, C, L, M] + latents = self._pack_audio_latents(latents) + latents = self._normalize_audio_latents(latents, self.audio_vae.latents_mean, self.audio_vae.latents_std) + latents = self._create_noised_state(latents, noise_scale, generator) + return latents.to(device=device, dtype=dtype) + + latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio + + shape = (batch_size, num_channels_latents, audio_latent_length, latent_mel_bins) + + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = self._pack_audio_latents(latents) + return latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def guidance_rescale(self): + return self._guidance_rescale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + conditions: LTX2VideoCondition | list[LTX2VideoCondition] | None = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] | None = None, + height: int = 512, + width: int = 768, + num_frames: int = 121, + frame_rate: float = 24.0, + num_inference_steps: int = 40, + sigmas: list[float] | None = None, + timesteps: list[float] | None = None, + guidance_scale: float = 4.0, + guidance_rescale: float = 0.0, + noise_scale: float | None = None, + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + audio_latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + decode_timestep: float | list[float] = 0.0, + decode_noise_scale: float | list[float] | None = None, + output_type: str = "pil", + return_dict: bool = True, + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + max_sequence_length: int = 1024, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + conditions (`List[LTXVideoCondition], *optional*`): + The list of frame-conditioning items for the video generation. + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + height (`int`, *optional*, defaults to `512`): + The height in pixels of the generated image. This is set to 480 by default for the best results. + width (`int`, *optional*, defaults to `768`): + The width in pixels of the generated image. This is set to 848 by default for the best results. + num_frames (`int`, *optional*, defaults to `121`): + The number of video frames to generate + frame_rate (`float`, *optional*, defaults to `24.0`): + The frames per second (FPS) of the generated video. + num_inference_steps (`int`, *optional*, defaults to 40): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + timesteps (`List[int]`, *optional*): + Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument + in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is + passed will be used. Must be in descending order. + guidance_scale (`float`, *optional*, defaults to `4.0`): + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. + of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting + `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to + the text `prompt`, usually at the expense of lower image quality. + guidance_rescale (`float`, *optional*, defaults to 0.0): + Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are + Flawed](https://huggingface.co/papers/2305.08891) `guidance_scale` is defined as `φ` in equation 16. of + [Common Diffusion Noise Schedules and Sample Steps are + Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when + using zero terminal SNR. + noise_scale (`float`, *optional*, defaults to `None`): + The interpolation factor between random noise and denoised latents at each timestep. Applying noise to + the `latents` and `audio_latents` before continue denoising. If not set, will be inferred from the + sigma schedule. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of videos to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for video + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + audio_latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for audio + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + prompt_attention_mask (`torch.Tensor`, *optional*): + Pre-generated attention mask for text embeddings. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not + provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. + negative_prompt_attention_mask (`torch.FloatTensor`, *optional*): + Pre-generated attention mask for negative text embeddings. + decode_timestep (`float`, defaults to `0.0`): + The timestep at which generated video is decoded. + decode_noise_scale (`float`, defaults to `None`): + The interpolation factor between random noise and denoised latents at the decode timestep. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.ltx.LTX2PipelineOutput`] instead of a plain tuple. + attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int`, *optional*, defaults to `1024`): + Maximum sequence length to use with the `prompt`. + + Examples: + + Returns: + [`~pipelines.ltx.LTX2PipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`~pipelines.ltx.LTX2PipelineOutput`] is returned, otherwise a `tuple` is + returned where the first element is a list with the generated images. + """ + + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt=prompt, + height=height, + width=width, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + latents=latents, + audio_latents=audio_latents, + ) + + self._guidance_scale = guidance_scale + self._guidance_rescale = guidance_rescale + self._attention_kwargs = attention_kwargs + self._interrupt = False + self._current_timestep = None + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if conditions is not None and not isinstance(conditions, list): + conditions = [conditions] + + # Infer noise scale: first (largest) sigma value if using custom sigmas, else 1.0 + if noise_scale is None: + noise_scale = sigmas[0] if sigmas is not None else 1.0 + + device = self._execution_device + + # 3. Prepare text embeddings + ( + prompt_embeds, + prompt_attention_mask, + negative_prompt_embeds, + negative_prompt_attention_mask, + ) = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + num_videos_per_prompt=num_videos_per_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + max_sequence_length=max_sequence_length, + device=device, + ) + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) + + additive_attention_mask = (1 - prompt_attention_mask.to(prompt_embeds.dtype)) * -1000000.0 + connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors( + prompt_embeds, additive_attention_mask, additive_mask=True + ) + + # 4. Prepare latent variables + latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 + latent_height = height // self.vae_spatial_compression_ratio + latent_width = width // self.vae_spatial_compression_ratio + if latents is not None: + logger.info( + "Got latents of shape [batch_size, latent_dim, latent_frames, latent_height, latent_width], `latent_num_frames`, `latent_height`, `latent_width` will be inferred." + ) + _, _, latent_num_frames, latent_height, latent_width = latents.shape # [B, C, F, H, W] + video_sequence_length = latent_num_frames * latent_height * latent_width + + num_channels_latents = self.transformer.config.in_channels + latents, conditioning_mask, clean_latents = self.prepare_latents( + conditions, + batch_size * num_videos_per_prompt, + num_channels_latents, + height, + width, + num_frames, + noise_scale, + torch.float32, + device, + generator, + latents, + ) + if self.do_classifier_free_guidance: + conditioning_mask = torch.cat([conditioning_mask, conditioning_mask]) + + duration_s = num_frames / frame_rate + audio_latents_per_second = ( + self.audio_sampling_rate / self.audio_hop_length / float(self.audio_vae_temporal_compression_ratio) + ) + audio_num_frames = round(duration_s * audio_latents_per_second) + if audio_latents is not None: + logger.info( + "Got audio_latents of shape [batch_size, num_channels, audio_num_frames, mel_bins], `audio_num_frames` will be inferred." + ) + _, _, audio_num_frames, _ = audio_latents.shape # [B, C, L, M] + + num_mel_bins = self.audio_vae.config.mel_bins if getattr(self, "audio_vae", None) is not None else 64 + latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio + num_channels_latents_audio = ( + self.audio_vae.config.latent_channels if getattr(self, "audio_vae", None) is not None else 8 + ) + audio_latents = self.prepare_audio_latents( + batch_size * num_videos_per_prompt, + num_channels_latents=num_channels_latents_audio, + audio_latent_length=audio_num_frames, + num_mel_bins=num_mel_bins, + noise_scale=noise_scale, + dtype=torch.float32, + device=device, + generator=generator, + latents=audio_latents, + ) + + # 5. Prepare timesteps + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas + mu = calculate_shift( + video_sequence_length, + self.scheduler.config.get("base_image_seq_len", 1024), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.95), + self.scheduler.config.get("max_shift", 2.05), + ) + + # For now, duplicate the scheduler for use with the audio latents + audio_scheduler = copy.deepcopy(self.scheduler) + _, _ = retrieve_timesteps( + audio_scheduler, + num_inference_steps, + device, + timesteps, + sigmas=sigmas, + mu=mu, + ) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + timesteps, + sigmas=sigmas, + mu=mu, + ) + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + # 6. Prepare micro-conditions + rope_interpolation_scale = ( + self.vae_temporal_compression_ratio / frame_rate, + self.vae_spatial_compression_ratio, + self.vae_spatial_compression_ratio, + ) + # Pre-compute video and audio positional ids as they will be the same at each step of the denoising loop + video_coords = self.transformer.rope.prepare_video_coords( + latents.shape[0], latent_num_frames, latent_height, latent_width, latents.device, fps=frame_rate + ) + audio_coords = self.transformer.audio_rope.prepare_audio_coords( + audio_latents.shape[0], audio_num_frames, audio_latents.device + ) + # Duplicate the positional ids as well if using CFG + if self.do_classifier_free_guidance: + video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1)) # Repeat twice in batch dim + audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1)) + + # 7. Denoising loop + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + self._current_timestep = t + + latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents + latent_model_input = latent_model_input.to(prompt_embeds.dtype) + audio_latent_model_input = ( + torch.cat([audio_latents] * 2) if self.do_classifier_free_guidance else audio_latents + ) + audio_latent_model_input = audio_latent_model_input.to(prompt_embeds.dtype) + + timestep = t.expand(latent_model_input.shape[0]) + video_timestep = timestep.unsqueeze(-1) * (1 - conditioning_mask.squeeze(-1)) + + with self.transformer.cache_context("cond_uncond"): + noise_pred_video, noise_pred_audio = self.transformer( + hidden_states=latent_model_input, + audio_hidden_states=audio_latent_model_input, + encoder_hidden_states=connector_prompt_embeds, + audio_encoder_hidden_states=connector_audio_prompt_embeds, + timestep=video_timestep, + audio_timestep=timestep, + encoder_attention_mask=connector_attention_mask, + audio_encoder_attention_mask=connector_attention_mask, + num_frames=latent_num_frames, + height=latent_height, + width=latent_width, + fps=frame_rate, + audio_num_frames=audio_num_frames, + video_coords=video_coords, + audio_coords=audio_coords, + # rope_interpolation_scale=rope_interpolation_scale, + attention_kwargs=attention_kwargs, + return_dict=False, + ) + noise_pred_video = noise_pred_video.float() + noise_pred_audio = noise_pred_audio.float() + + if self.do_classifier_free_guidance: + noise_pred_video_uncond, noise_pred_video_text = noise_pred_video.chunk(2) + noise_pred_video = noise_pred_video_uncond + self.guidance_scale * ( + noise_pred_video_text - noise_pred_video_uncond + ) + + noise_pred_audio_uncond, noise_pred_audio_text = noise_pred_audio.chunk(2) + noise_pred_audio = noise_pred_audio_uncond + self.guidance_scale * ( + noise_pred_audio_text - noise_pred_audio_uncond + ) + + if self.guidance_rescale > 0: + # Based on 3.4. in https://huggingface.co/papers/2305.08891 + noise_pred_video = rescale_noise_cfg( + noise_pred_video, noise_pred_video_text, guidance_rescale=self.guidance_rescale + ) + noise_pred_audio = rescale_noise_cfg( + noise_pred_audio, noise_pred_audio_text, guidance_rescale=self.guidance_rescale + ) + + # NOTE: use only the first chunk of conditioning mask in case it is duplicated for CFG + bsz = noise_pred_video.size(0) + sigma = self.scheduler.sigmas[i] + # Convert the noise_pred_video velocity model prediction into a sample (x0) prediction + denoised_sample = latents - noise_pred_video * sigma + # Apply the (packed) conditioning mask to the denoised (x0) sample and clean conditioning. The + # conditioning mask contains conditioning strengths from 0 (always use denoised sample) to 1 (always + # use conditions), with intermediate values specifying how strongly to follow the conditions. + denoised_sample_cond = ( + denoised_sample * (1 - conditioning_mask[:bsz]) + clean_latents.float() * conditioning_mask[:bsz] + ).to(noise_pred_video.dtype) + # Convert the denoised (x0) sample back to a velocity for the scheduler + denoised_latents_cond = ((latents - denoised_sample_cond) / sigma).to(noise_pred_video.dtype) + + # Compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(denoised_latents_cond, t, latents, return_dict=False)[0] + + # NOTE: for now duplicate scheduler for audio latents in case self.scheduler sets internal state in + # the step method (such as _step_index) + audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0] + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + latents = self._unpack_latents( + latents, + latent_num_frames, + latent_height, + latent_width, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) + latents = self._denormalize_latents( + latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor + ) + + audio_latents = self._denormalize_audio_latents( + audio_latents, self.audio_vae.latents_mean, self.audio_vae.latents_std + ) + audio_latents = self._unpack_audio_latents(audio_latents, audio_num_frames, num_mel_bins=latent_mel_bins) + + if output_type == "latent": + video = latents + audio = audio_latents + else: + latents = latents.to(prompt_embeds.dtype) + + if not self.vae.config.timestep_conditioning: + timestep = None + else: + noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype) + if not isinstance(decode_timestep, list): + decode_timestep = [decode_timestep] * batch_size + if decode_noise_scale is None: + decode_noise_scale = decode_timestep + elif not isinstance(decode_noise_scale, list): + decode_noise_scale = [decode_noise_scale] * batch_size + + timestep = torch.tensor(decode_timestep, device=device, dtype=latents.dtype) + decode_noise_scale = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[ + :, None, None, None, None + ] + latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise + + latents = latents.to(self.vae.dtype) + video = self.vae.decode(latents, timestep, return_dict=False)[0] + video = self.video_processor.postprocess_video(video, output_type=output_type) + + audio_latents = audio_latents.to(self.audio_vae.dtype) + generated_mel_spectrograms = self.audio_vae.decode(audio_latents, return_dict=False)[0] + audio = self.vocoder(generated_mel_spectrograms) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (video, audio) + + return LTX2PipelineOutput(frames=video, audio=audio) diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py new file mode 100644 index 000000000000..83ba2cd7c685 --- /dev/null +++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py @@ -0,0 +1,1317 @@ +# Copyright 2025 Lightricks and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import inspect +from typing import Any, Callable + +import numpy as np +import torch +from transformers import Gemma3ForConditionalGeneration, GemmaTokenizer, GemmaTokenizerFast + +from ...callbacks import MultiPipelineCallbacks, PipelineCallback +from ...image_processor import PipelineImageInput +from ...loaders import FromSingleFileMixin, LTX2LoraLoaderMixin +from ...models.autoencoders import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video +from ...models.transformers import LTX2VideoTransformer3DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import is_torch_xla_available, logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor +from ..pipeline_utils import DiffusionPipeline +from .connectors import LTX2TextConnectors +from .pipeline_output import LTX2PipelineOutput +from .vocoder import LTX2Vocoder + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import LTX2ImageToVideoPipeline + >>> from diffusers.pipelines.ltx2.export_utils import encode_video + >>> from diffusers.utils import load_image + + >>> pipe = LTX2ImageToVideoPipeline.from_pretrained("Lightricks/LTX-2", torch_dtype=torch.bfloat16) + >>> pipe.enable_model_cpu_offload() + + >>> image = load_image( + ... "https://huggingface.co/datasets/a-r-r-o-w/tiny-meme-dataset-captioned/resolve/main/images/8.png" + ... ) + >>> prompt = "A young girl stands calmly in the foreground, looking directly at the camera, as a house fire rages in the background." + >>> negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted" + + >>> frame_rate = 24.0 + >>> video, audio = pipe( + ... image=image, + ... prompt=prompt, + ... negative_prompt=negative_prompt, + ... width=768, + ... height=512, + ... num_frames=121, + ... frame_rate=frame_rate, + ... num_inference_steps=40, + ... guidance_scale=4.0, + ... output_type="np", + ... return_dict=False, + ... ) + + >>> encode_video( + ... video[0], + ... fps=frame_rate, + ... audio=audio[0].float().cpu(), + ... audio_sample_rate=pipe.vocoder.config.output_sampling_rate, # should be 24000 + ... output_path="video.mp4", + ... ) + ``` +""" + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents +def retrieve_latents( + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" +): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`list[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`list[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg +def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): + r""" + Rescales `noise_cfg` tensor based on `guidance_rescale` to improve image quality and fix overexposure. Based on + Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are + Flawed](https://huggingface.co/papers/2305.08891). + + Args: + noise_cfg (`torch.Tensor`): + The predicted noise tensor for the guided diffusion process. + noise_pred_text (`torch.Tensor`): + The predicted noise tensor for the text-guided diffusion process. + guidance_rescale (`float`, *optional*, defaults to 0.0): + A rescale factor applied to the noise predictions. + + Returns: + noise_cfg (`torch.Tensor`): The rescaled noise prediction tensor. + """ + std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True) + std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) + # rescale the results from guidance (fixes overexposure) + noise_pred_rescaled = noise_cfg * (std_text / std_cfg) + # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images + noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg + return noise_cfg + + +class LTX2ImageToVideoPipeline(DiffusionPipeline, FromSingleFileMixin, LTX2LoraLoaderMixin): + r""" + Pipeline for image-to-video generation. + + Reference: https://github.com/Lightricks/LTX-Video + + TODO + """ + + model_cpu_offload_seq = "text_encoder->connectors->transformer->vae->audio_vae->vocoder" + _optional_components = [] + _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] + + def __init__( + self, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKLLTX2Video, + audio_vae: AutoencoderKLLTX2Audio, + text_encoder: Gemma3ForConditionalGeneration, + tokenizer: GemmaTokenizer | GemmaTokenizerFast, + connectors: LTX2TextConnectors, + transformer: LTX2VideoTransformer3DModel, + vocoder: LTX2Vocoder, + ): + super().__init__() + + self.register_modules( + vae=vae, + audio_vae=audio_vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + connectors=connectors, + transformer=transformer, + vocoder=vocoder, + scheduler=scheduler, + ) + + self.vae_spatial_compression_ratio = ( + self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 32 + ) + self.vae_temporal_compression_ratio = ( + self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 8 + ) + # TODO: check whether the MEL compression ratio logic here is corrct + self.audio_vae_mel_compression_ratio = ( + self.audio_vae.mel_compression_ratio if getattr(self, "audio_vae", None) is not None else 4 + ) + self.audio_vae_temporal_compression_ratio = ( + self.audio_vae.temporal_compression_ratio if getattr(self, "audio_vae", None) is not None else 4 + ) + self.transformer_spatial_patch_size = ( + self.transformer.config.patch_size if getattr(self, "transformer", None) is not None else 1 + ) + self.transformer_temporal_patch_size = ( + self.transformer.config.patch_size_t if getattr(self, "transformer") is not None else 1 + ) + + self.audio_sampling_rate = ( + self.audio_vae.config.sample_rate if getattr(self, "audio_vae", None) is not None else 16000 + ) + self.audio_hop_length = ( + self.audio_vae.config.mel_hop_length if getattr(self, "audio_vae", None) is not None else 160 + ) + + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio, resample="bilinear") + self.tokenizer_max_length = ( + self.tokenizer.model_max_length if getattr(self, "tokenizer", None) is not None else 1024 + ) + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._pack_text_embeds + def _pack_text_embeds( + text_hidden_states: torch.Tensor, + sequence_lengths: torch.Tensor, + device: str | torch.device, + padding_side: str = "left", + scale_factor: int = 8, + eps: float = 1e-6, + ) -> torch.Tensor: + """ + Packs and normalizes text encoder hidden states, respecting padding. Normalization is performed per-batch and + per-layer in a masked fashion (only over non-padded positions). + + Args: + text_hidden_states (`torch.Tensor` of shape `(batch_size, seq_len, hidden_dim, num_layers)`): + Per-layer hidden_states from a text encoder (e.g. `Gemma3ForConditionalGeneration`). + sequence_lengths (`torch.Tensor of shape `(batch_size,)`): + The number of valid (non-padded) tokens for each batch instance. + device: (`str` or `torch.device`, *optional*): + torch device to place the resulting embeddings on + padding_side: (`str`, *optional*, defaults to `"left"`): + Whether the text tokenizer performs padding on the `"left"` or `"right"`. + scale_factor (`int`, *optional*, defaults to `8`): + Scaling factor to multiply the normalized hidden states by. + eps (`float`, *optional*, defaults to `1e-6`): + A small positive value for numerical stability when performing normalization. + + Returns: + `torch.Tensor` of shape `(batch_size, seq_len, hidden_dim * num_layers)`: + Normed and flattened text encoder hidden states. + """ + batch_size, seq_len, hidden_dim, num_layers = text_hidden_states.shape + original_dtype = text_hidden_states.dtype + + # Create padding mask + token_indices = torch.arange(seq_len, device=device).unsqueeze(0) + if padding_side == "right": + # For right padding, valid tokens are from 0 to sequence_length-1 + mask = token_indices < sequence_lengths[:, None] # [batch_size, seq_len] + elif padding_side == "left": + # For left padding, valid tokens are from (T - sequence_length) to T-1 + start_indices = seq_len - sequence_lengths[:, None] # [batch_size, 1] + mask = token_indices >= start_indices # [B, T] + else: + raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}") + mask = mask[:, :, None, None] # [batch_size, seq_len] --> [batch_size, seq_len, 1, 1] + + # Compute masked mean over non-padding positions of shape (batch_size, 1, 1, seq_len) + masked_text_hidden_states = text_hidden_states.masked_fill(~mask, 0.0) + num_valid_positions = (sequence_lengths * hidden_dim).view(batch_size, 1, 1, 1) + masked_mean = masked_text_hidden_states.sum(dim=(1, 2), keepdim=True) / (num_valid_positions + eps) + + # Compute min/max over non-padding positions of shape (batch_size, 1, 1 seq_len) + x_min = text_hidden_states.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True) + x_max = text_hidden_states.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True) + + # Normalization + normalized_hidden_states = (text_hidden_states - masked_mean) / (x_max - x_min + eps) + normalized_hidden_states = normalized_hidden_states * scale_factor + + # Pack the hidden states to a 3D tensor (batch_size, seq_len, hidden_dim * num_layers) + normalized_hidden_states = normalized_hidden_states.flatten(2) + mask_flat = mask.squeeze(-1).expand(-1, -1, hidden_dim * num_layers) + normalized_hidden_states = normalized_hidden_states.masked_fill(~mask_flat, 0.0) + normalized_hidden_states = normalized_hidden_states.to(dtype=original_dtype) + return normalized_hidden_states + + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._get_gemma_prompt_embeds + def _get_gemma_prompt_embeds( + self, + prompt: str | list[str], + num_videos_per_prompt: int = 1, + max_sequence_length: int = 1024, + scale_factor: int = 8, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `list[str]`, *optional*): + prompt to be encoded + device: (`str` or `torch.device`): + torch device to place the resulting embeddings on + dtype: (`torch.dtype`): + torch dtype to cast the prompt embeds to + max_sequence_length (`int`, defaults to 1024): Maximum sequence length to use for the prompt. + """ + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + if getattr(self, "tokenizer", None) is not None: + # Gemma expects left padding for chat-style prompts + self.tokenizer.padding_side = "left" + if self.tokenizer.pad_token is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + + prompt = [p.strip() for p in prompt] + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + prompt_attention_mask = text_inputs.attention_mask + text_input_ids = text_input_ids.to(device) + prompt_attention_mask = prompt_attention_mask.to(device) + + text_encoder_outputs = self.text_encoder( + input_ids=text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True + ) + text_encoder_hidden_states = text_encoder_outputs.hidden_states + text_encoder_hidden_states = torch.stack(text_encoder_hidden_states, dim=-1) + sequence_lengths = prompt_attention_mask.sum(dim=-1) + + prompt_embeds = self._pack_text_embeds( + text_encoder_hidden_states, + sequence_lengths, + device=device, + padding_side=self.tokenizer.padding_side, + scale_factor=scale_factor, + ) + prompt_embeds = prompt_embeds.to(dtype=dtype) + + # duplicate text embeddings for each generation per prompt, using mps friendly method + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) + + prompt_attention_mask = prompt_attention_mask.view(batch_size, -1) + prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1) + + return prompt_embeds, prompt_attention_mask + + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline.encode_prompt + def encode_prompt( + self, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, + do_classifier_free_guidance: bool = True, + num_videos_per_prompt: int = 1, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + max_sequence_length: int = 1024, + scale_factor: int = 8, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `list[str]`, *optional*): + prompt to be encoded + negative_prompt (`str` or `list[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): + Whether to use classifier free guidance or not. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos that should be generated per prompt. torch device to place the resulting embeddings on + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + device: (`torch.device`, *optional*): + torch device + dtype: (`torch.dtype`, *optional*): + torch dtype + """ + device = device or self._execution_device + + prompt = [prompt] if isinstance(prompt, str) else prompt + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds( + prompt=prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + scale_factor=scale_factor, + device=device, + dtype=dtype, + ) + + if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + + if prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + + negative_prompt_embeds, negative_prompt_attention_mask = self._get_gemma_prompt_embeds( + prompt=negative_prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + scale_factor=scale_factor, + device=device, + dtype=dtype, + ) + + return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask + + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline.check_inputs + def check_inputs( + self, + prompt, + height, + width, + callback_on_step_end_tensor_inputs=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + ): + if height % 32 != 0 or width % 32 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.") + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if prompt_embeds is not None and prompt_attention_mask is None: + raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.") + + if negative_prompt_embeds is not None and negative_prompt_attention_mask is None: + raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.") + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}." + ) + if prompt_attention_mask.shape != negative_prompt_attention_mask.shape: + raise ValueError( + "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but" + f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`" + f" {negative_prompt_attention_mask.shape}." + ) + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._pack_latents + def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor: + # Unpacked latents of shape are [B, C, F, H, W] are patched into tokens of shape [B, C, F // p_t, p_t, H // p, p, W // p, p]. + # The patch dimensions are then permuted and collapsed into the channel dimension of shape: + # [B, F // p_t * H // p * W // p, C * p_t * p * p] (an ndim=3 tensor). + # dim=0 is the batch size, dim=1 is the effective video sequence length, dim=2 is the effective number of input features + batch_size, num_channels, num_frames, height, width = latents.shape + post_patch_num_frames = num_frames // patch_size_t + post_patch_height = height // patch_size + post_patch_width = width // patch_size + latents = latents.reshape( + batch_size, + -1, + post_patch_num_frames, + patch_size_t, + post_patch_height, + patch_size, + post_patch_width, + patch_size, + ) + latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3) + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._unpack_latents + def _unpack_latents( + latents: torch.Tensor, num_frames: int, height: int, width: int, patch_size: int = 1, patch_size_t: int = 1 + ) -> torch.Tensor: + # Packed latents of shape [B, S, D] (S is the effective video sequence length, D is the effective feature dimensions) + # are unpacked and reshaped into a video tensor of shape [B, C, F, H, W]. This is the inverse operation of + # what happens in the `_pack_latents` method. + batch_size = latents.size(0) + latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size) + latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3) + return latents + + @staticmethod + def _normalize_latents( + latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 + ) -> torch.Tensor: + # Normalize latents across the channel dimension [B, C, F, H, W] + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents = (latents - latents_mean) * scaling_factor / latents_std + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._denormalize_latents + def _denormalize_latents( + latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 + ) -> torch.Tensor: + # Denormalize latents across the channel dimension [B, C, F, H, W] + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents = latents * latents_std / scaling_factor + latents_mean + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._create_noised_state + def _create_noised_state( + latents: torch.Tensor, noise_scale: float | torch.Tensor, generator: torch.Generator | None = None + ): + noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype) + noised_latents = noise_scale * noise + (1 - noise_scale) * latents + return noised_latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._pack_audio_latents + def _pack_audio_latents( + latents: torch.Tensor, patch_size: int | None = None, patch_size_t: int | None = None + ) -> torch.Tensor: + # Audio latents shape: [B, C, L, M], where L is the latent audio length and M is the number of mel bins + if patch_size is not None and patch_size_t is not None: + # Packs the latents into a patch sequence of shape [B, L // p_t * M // p, C * p_t * p] (a ndim=3 tnesor). + # dim=1 is the effective audio sequence length and dim=2 is the effective audio input feature size. + batch_size, num_channels, latent_length, latent_mel_bins = latents.shape + post_patch_latent_length = latent_length / patch_size_t + post_patch_mel_bins = latent_mel_bins / patch_size + latents = latents.reshape( + batch_size, -1, post_patch_latent_length, patch_size_t, post_patch_mel_bins, patch_size + ) + latents = latents.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2) + else: + # Packs the latents into a patch sequence of shape [B, L, C * M]. This implicitly assumes a (mel) + # patch_size of M (all mel bins constitutes a single patch) and a patch_size_t of 1. + latents = latents.transpose(1, 2).flatten(2, 3) # [B, C, L, M] --> [B, L, C * M] + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._unpack_audio_latents + def _unpack_audio_latents( + latents: torch.Tensor, + latent_length: int, + num_mel_bins: int, + patch_size: int | None = None, + patch_size_t: int | None = None, + ) -> torch.Tensor: + # Unpacks an audio patch sequence of shape [B, S, D] into a latent spectrogram tensor of shape [B, C, L, M], + # where L is the latent audio length and M is the number of mel bins. + if patch_size is not None and patch_size_t is not None: + batch_size = latents.size(0) + latents = latents.reshape(batch_size, latent_length, num_mel_bins, -1, patch_size_t, patch_size) + latents = latents.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) + else: + # Assume [B, S, D] = [B, L, C * M], which implies that patch_size = M and patch_size_t = 1. + latents = latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2) + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._normalize_audio_latents + def _normalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor): + latents_mean = latents_mean.to(latents.device, latents.dtype) + latents_std = latents_std.to(latents.device, latents.dtype) + return (latents - latents_mean) / latents_std + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._denormalize_audio_latents + def _denormalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor): + latents_mean = latents_mean.to(latents.device, latents.dtype) + latents_std = latents_std.to(latents.device, latents.dtype) + return (latents * latents_std) + latents_mean + + def prepare_latents( + self, + image: torch.Tensor | None = None, + batch_size: int = 1, + num_channels_latents: int = 128, + height: int = 512, + width: int = 704, + num_frames: int = 161, + noise_scale: float = 0.0, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + height = height // self.vae_spatial_compression_ratio + width = width // self.vae_spatial_compression_ratio + num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 + + shape = (batch_size, num_channels_latents, num_frames, height, width) + mask_shape = (batch_size, 1, num_frames, height, width) + + if latents is not None: + if latents.ndim == 5: + # conditioning_mask needs to the same shape as latents in two stages generation. + batch_size, _, num_frames, height, width = latents.shape + mask_shape = (batch_size, 1, num_frames, height, width) + conditioning_mask = latents.new_zeros(mask_shape) + conditioning_mask[:, :, 0] = 1.0 + + latents = self._normalize_latents( + latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor + ) + latents = self._create_noised_state(latents, noise_scale * (1 - conditioning_mask), generator) + # latents are of shape [B, C, F, H, W], need to be packed + latents = self._pack_latents( + latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size + ) + else: + conditioning_mask = latents.new_zeros(mask_shape) + conditioning_mask[:, :, 0] = 1.0 + conditioning_mask = self._pack_latents( + conditioning_mask, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size + ).squeeze(-1) + if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape: + raise ValueError( + f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is {conditioning_mask.shape + (num_channels_latents,)}." + ) + return latents.to(device=device, dtype=dtype), conditioning_mask + + if isinstance(generator, list): + if len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + init_latents = [ + retrieve_latents(self.vae.encode(image[i].unsqueeze(0).unsqueeze(2)), generator[i], "argmax") + for i in range(batch_size) + ] + else: + init_latents = [ + retrieve_latents(self.vae.encode(img.unsqueeze(0).unsqueeze(2)), generator, "argmax") for img in image + ] + + init_latents = torch.cat(init_latents, dim=0).to(dtype) + init_latents = self._normalize_latents(init_latents, self.vae.latents_mean, self.vae.latents_std) + init_latents = init_latents.repeat(1, 1, num_frames, 1, 1) + + # First condition is image latents and those should be kept clean. + conditioning_mask = torch.zeros(mask_shape, device=device, dtype=dtype) + conditioning_mask[:, :, 0] = 1.0 + + noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + # Interpolation. + latents = init_latents * conditioning_mask + noise * (1 - conditioning_mask) + + conditioning_mask = self._pack_latents( + conditioning_mask, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size + ).squeeze(-1) + latents = self._pack_latents( + latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size + ) + + return latents, conditioning_mask + + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline.prepare_audio_latents + def prepare_audio_latents( + self, + batch_size: int = 1, + num_channels_latents: int = 8, + audio_latent_length: int = 1, # 1 is just a dummy value + num_mel_bins: int = 64, + noise_scale: float = 0.0, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + if latents is not None: + if latents.ndim == 4: + # latents are of shape [B, C, L, M], need to be packed + latents = self._pack_audio_latents(latents) + if latents.ndim != 3: + raise ValueError( + f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is [batch_size, num_seq, num_features]." + ) + latents = self._normalize_audio_latents(latents, self.audio_vae.latents_mean, self.audio_vae.latents_std) + latents = self._create_noised_state(latents, noise_scale, generator) + return latents.to(device=device, dtype=dtype) + + # TODO: confirm whether this logic is correct + latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio + + shape = (batch_size, num_channels_latents, audio_latent_length, latent_mel_bins) + + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = self._pack_audio_latents(latents) + return latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def guidance_rescale(self): + return self._guidance_rescale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + image: PipelineImageInput = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] | None = None, + height: int = 512, + width: int = 768, + num_frames: int = 121, + frame_rate: float = 24.0, + num_inference_steps: int = 40, + sigmas: list[float] | None = None, + timesteps: list[int] | None = None, + guidance_scale: float = 4.0, + guidance_rescale: float = 0.0, + noise_scale: float = 0.0, + num_videos_per_prompt: int = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + audio_latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + decode_timestep: float | list[float] = 0.0, + decode_noise_scale: float | list[float] | None = None, + output_type: str = "pil", + return_dict: bool = True, + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + max_sequence_length: int = 1024, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + image (`PipelineImageInput`): + The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`. + prompt (`str` or `list[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + height (`int`, *optional*, defaults to `512`): + The height in pixels of the generated image. This is set to 480 by default for the best results. + width (`int`, *optional*, defaults to `768`): + The width in pixels of the generated image. This is set to 848 by default for the best results. + num_frames (`int`, *optional*, defaults to `121`): + The number of video frames to generate + frame_rate (`float`, *optional*, defaults to `24.0`): + The frames per second (FPS) of the generated video. + num_inference_steps (`int`, *optional*, defaults to 40): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + timesteps (`List[int]`, *optional*): + Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument + in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is + passed will be used. Must be in descending order. + guidance_scale (`float`, *optional*, defaults to `4.0`): + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. + of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting + `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to + the text `prompt`, usually at the expense of lower image quality. + guidance_rescale (`float`, *optional*, defaults to 0.0): + Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are + Flawed](https://huggingface.co/papers/2305.08891) `guidance_scale` is defined as `φ` in equation 16. of + [Common Diffusion Noise Schedules and Sample Steps are + Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when + using zero terminal SNR. + noise_scale (`float`, *optional*, defaults to `0.0`): + The interpolation factor between random noise and denoised latents at each timestep. Applying noise to + the `latents` and `audio_latents` before continue denoising. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of videos to generate per prompt. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for video + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + audio_latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for audio + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + prompt_attention_mask (`torch.Tensor`, *optional*): + Pre-generated attention mask for text embeddings. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not + provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. + negative_prompt_attention_mask (`torch.FloatTensor`, *optional*): + Pre-generated attention mask for negative text embeddings. + decode_timestep (`float`, defaults to `0.0`): + The timestep at which generated video is decoded. + decode_noise_scale (`float`, defaults to `None`): + The interpolation factor between random noise and denoised latents at the decode timestep. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.ltx.LTX2PipelineOutput`] instead of a plain tuple. + attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int`, *optional*, defaults to `1024`): + Maximum sequence length to use with the `prompt`. + + Examples: + + Returns: + [`~pipelines.ltx.LTX2PipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`~pipelines.ltx.LTX2PipelineOutput`] is returned, otherwise a `tuple` is + returned where the first element is a list with the generated images. + """ + + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt=prompt, + height=height, + width=width, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + ) + + self._guidance_scale = guidance_scale + self._guidance_rescale = guidance_rescale + self._attention_kwargs = attention_kwargs + self._interrupt = False + self._current_timestep = None + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + + # 3. Prepare text embeddings + ( + prompt_embeds, + prompt_attention_mask, + negative_prompt_embeds, + negative_prompt_attention_mask, + ) = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + num_videos_per_prompt=num_videos_per_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + max_sequence_length=max_sequence_length, + device=device, + ) + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) + + additive_attention_mask = (1 - prompt_attention_mask.to(prompt_embeds.dtype)) * -1000000.0 + connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors( + prompt_embeds, additive_attention_mask, additive_mask=True + ) + + # 4. Prepare latent variables + latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 + latent_height = height // self.vae_spatial_compression_ratio + latent_width = width // self.vae_spatial_compression_ratio + if latents is not None: + if latents.ndim == 5: + logger.info( + "Got latents of shape [batch_size, latent_dim, latent_frames, latent_height, latent_width], `latent_num_frames`, `latent_height`, `latent_width` will be inferred." + ) + _, _, latent_num_frames, latent_height, latent_width = latents.shape # [B, C, F, H, W] + elif latents.ndim == 3: + logger.warning( + f"You have supplied packed `latents` of shape {latents.shape}, so the latent dims cannot be" + f" inferred. Make sure the supplied `height`, `width`, and `num_frames` are correct." + ) + else: + raise ValueError( + f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is either [batch_size, seq_len, num_features] or [batch_size, latent_dim, latent_frames, latent_height, latent_width]." + ) + video_sequence_length = latent_num_frames * latent_height * latent_width + + if latents is None: + image = self.video_processor.preprocess(image, height=height, width=width) + image = image.to(device=device, dtype=prompt_embeds.dtype) + + num_channels_latents = self.transformer.config.in_channels + latents, conditioning_mask = self.prepare_latents( + image, + batch_size * num_videos_per_prompt, + num_channels_latents, + height, + width, + num_frames, + noise_scale, + torch.float32, + device, + generator, + latents, + ) + if self.do_classifier_free_guidance: + conditioning_mask = torch.cat([conditioning_mask, conditioning_mask]) + + duration_s = num_frames / frame_rate + audio_latents_per_second = ( + self.audio_sampling_rate / self.audio_hop_length / float(self.audio_vae_temporal_compression_ratio) + ) + audio_num_frames = round(duration_s * audio_latents_per_second) + if audio_latents is not None: + if audio_latents.ndim == 4: + logger.info( + "Got audio_latents of shape [batch_size, num_channels, audio_length, mel_bins], `audio_num_frames` will be inferred." + ) + _, _, audio_num_frames, _ = audio_latents.shape # [B, C, L, M] + elif audio_latents.ndim == 3: + logger.warning( + f"You have supplied packed `audio_latents` of shape {audio_latents.shape}, so the latent dims" + f" cannot be inferred. Make sure the supplied `num_frames` and `frame_rate` are correct." + ) + else: + raise ValueError( + f"Provided `audio_latents` tensor has shape {audio_latents.shape}, but the expected shape is either [batch_size, seq_len, num_features] or [batch_size, num_channels, audio_length, mel_bins]." + ) + + num_mel_bins = self.audio_vae.config.mel_bins if getattr(self, "audio_vae", None) is not None else 64 + latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio + num_channels_latents_audio = ( + self.audio_vae.config.latent_channels if getattr(self, "audio_vae", None) is not None else 8 + ) + audio_latents = self.prepare_audio_latents( + batch_size * num_videos_per_prompt, + num_channels_latents=num_channels_latents_audio, + audio_latent_length=audio_num_frames, + num_mel_bins=num_mel_bins, + noise_scale=noise_scale, + dtype=torch.float32, + device=device, + generator=generator, + latents=audio_latents, + ) + + # 5. Prepare timesteps + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas + mu = calculate_shift( + video_sequence_length, + self.scheduler.config.get("base_image_seq_len", 1024), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.95), + self.scheduler.config.get("max_shift", 2.05), + ) + + # For now, duplicate the scheduler for use with the audio latents + audio_scheduler = copy.deepcopy(self.scheduler) + _, _ = retrieve_timesteps( + audio_scheduler, + num_inference_steps, + device, + timesteps, + sigmas=sigmas, + mu=mu, + ) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + timesteps, + sigmas=sigmas, + mu=mu, + ) + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + # 6. Prepare micro-conditions + rope_interpolation_scale = ( + self.vae_temporal_compression_ratio / frame_rate, + self.vae_spatial_compression_ratio, + self.vae_spatial_compression_ratio, + ) + # Pre-compute video and audio positional ids as they will be the same at each step of the denoising loop + video_coords = self.transformer.rope.prepare_video_coords( + latents.shape[0], latent_num_frames, latent_height, latent_width, latents.device, fps=frame_rate + ) + audio_coords = self.transformer.audio_rope.prepare_audio_coords( + audio_latents.shape[0], audio_num_frames, audio_latents.device + ) + # Duplicate the positional ids as well if using CFG + if self.do_classifier_free_guidance: + video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1)) # Repeat twice in batch dim + audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1)) + + # 7. Denoising loop + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + self._current_timestep = t + + latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents + latent_model_input = latent_model_input.to(prompt_embeds.dtype) + audio_latent_model_input = ( + torch.cat([audio_latents] * 2) if self.do_classifier_free_guidance else audio_latents + ) + audio_latent_model_input = audio_latent_model_input.to(prompt_embeds.dtype) + + timestep = t.expand(latent_model_input.shape[0]) + video_timestep = timestep.unsqueeze(-1) * (1 - conditioning_mask) + + with self.transformer.cache_context("cond_uncond"): + noise_pred_video, noise_pred_audio = self.transformer( + hidden_states=latent_model_input, + audio_hidden_states=audio_latent_model_input, + encoder_hidden_states=connector_prompt_embeds, + audio_encoder_hidden_states=connector_audio_prompt_embeds, + timestep=video_timestep, + audio_timestep=timestep, + encoder_attention_mask=connector_attention_mask, + audio_encoder_attention_mask=connector_attention_mask, + num_frames=latent_num_frames, + height=latent_height, + width=latent_width, + fps=frame_rate, + audio_num_frames=audio_num_frames, + video_coords=video_coords, + audio_coords=audio_coords, + # rope_interpolation_scale=rope_interpolation_scale, + attention_kwargs=attention_kwargs, + return_dict=False, + ) + noise_pred_video = noise_pred_video.float() + noise_pred_audio = noise_pred_audio.float() + + if self.do_classifier_free_guidance: + noise_pred_video_uncond, noise_pred_video_text = noise_pred_video.chunk(2) + noise_pred_video = noise_pred_video_uncond + self.guidance_scale * ( + noise_pred_video_text - noise_pred_video_uncond + ) + + noise_pred_audio_uncond, noise_pred_audio_text = noise_pred_audio.chunk(2) + noise_pred_audio = noise_pred_audio_uncond + self.guidance_scale * ( + noise_pred_audio_text - noise_pred_audio_uncond + ) + + if self.guidance_rescale > 0: + # Based on 3.4. in https://huggingface.co/papers/2305.08891 + noise_pred_video = rescale_noise_cfg( + noise_pred_video, noise_pred_video_text, guidance_rescale=self.guidance_rescale + ) + noise_pred_audio = rescale_noise_cfg( + noise_pred_audio, noise_pred_audio_text, guidance_rescale=self.guidance_rescale + ) + + # compute the previous noisy sample x_t -> x_t-1 + noise_pred_video = self._unpack_latents( + noise_pred_video, + latent_num_frames, + latent_height, + latent_width, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) + latents = self._unpack_latents( + latents, + latent_num_frames, + latent_height, + latent_width, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) + + noise_pred_video = noise_pred_video[:, :, 1:] + noise_latents = latents[:, :, 1:] + pred_latents = self.scheduler.step(noise_pred_video, t, noise_latents, return_dict=False)[0] + + latents = torch.cat([latents[:, :, :1], pred_latents], dim=2) + latents = self._pack_latents( + latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size + ) + + # NOTE: for now duplicate scheduler for audio latents in case self.scheduler sets internal state in + # the step method (such as _step_index) + audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0] + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + latents = self._unpack_latents( + latents, + latent_num_frames, + latent_height, + latent_width, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) + latents = self._denormalize_latents( + latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor + ) + + audio_latents = self._denormalize_audio_latents( + audio_latents, self.audio_vae.latents_mean, self.audio_vae.latents_std + ) + audio_latents = self._unpack_audio_latents(audio_latents, audio_num_frames, num_mel_bins=latent_mel_bins) + + if output_type == "latent": + video = latents + audio = audio_latents + else: + latents = latents.to(prompt_embeds.dtype) + + if not self.vae.config.timestep_conditioning: + timestep = None + else: + noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype) + if not isinstance(decode_timestep, list): + decode_timestep = [decode_timestep] * batch_size + if decode_noise_scale is None: + decode_noise_scale = decode_timestep + elif not isinstance(decode_noise_scale, list): + decode_noise_scale = [decode_noise_scale] * batch_size + + timestep = torch.tensor(decode_timestep, device=device, dtype=latents.dtype) + decode_noise_scale = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[ + :, None, None, None, None + ] + latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise + + latents = latents.to(self.vae.dtype) + video = self.vae.decode(latents, timestep, return_dict=False)[0] + video = self.video_processor.postprocess_video(video, output_type=output_type) + + audio_latents = audio_latents.to(self.audio_vae.dtype) + generated_mel_spectrograms = self.audio_vae.decode(audio_latents, return_dict=False)[0] + audio = self.vocoder(generated_mel_spectrograms) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (video, audio) + + return LTX2PipelineOutput(frames=video, audio=audio) diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_latent_upsample.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_latent_upsample.py new file mode 100644 index 000000000000..00d81dfd11c3 --- /dev/null +++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_latent_upsample.py @@ -0,0 +1,425 @@ +# Copyright 2025 Lightricks and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch + +from ...image_processor import PipelineImageInput +from ...models import AutoencoderKLLTX2Video +from ...utils import get_logger, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor +from ..ltx.pipeline_output import LTXPipelineOutput +from ..pipeline_utils import DiffusionPipeline +from .latent_upsampler import LTX2LatentUpsamplerModel + + +logger = get_logger(__name__) # pylint: disable=invalid-name + + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import LTX2ImageToVideoPipeline, LTX2LatentUpsamplePipeline + >>> from diffusers.pipelines.ltx2.export_utils import encode_video + >>> from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel + >>> from diffusers.utils import load_image + + >>> pipe = LTX2ImageToVideoPipeline.from_pretrained("Lightricks/LTX-2", torch_dtype=torch.bfloat16) + >>> pipe.enable_model_cpu_offload() + + >>> image = load_image( + ... "https://huggingface.co/datasets/a-r-r-o-w/tiny-meme-dataset-captioned/resolve/main/images/8.png" + ... ) + >>> prompt = "A young girl stands calmly in the foreground, looking directly at the camera, as a house fire rages in the background." + >>> negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted" + + >>> frame_rate = 24.0 + >>> video, audio = pipe( + ... image=image, + ... prompt=prompt, + ... negative_prompt=negative_prompt, + ... width=768, + ... height=512, + ... num_frames=121, + ... frame_rate=frame_rate, + ... num_inference_steps=40, + ... guidance_scale=4.0, + ... output_type="pil", + ... return_dict=False, + ... ) + + >>> latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained( + ... "Lightricks/LTX-2", subfolder="latent_upsampler", torch_dtype=torch.bfloat16 + ... ) + >>> upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=latent_upsampler) + >>> upsample_pipe.vae.enable_tiling() + >>> upsample_pipe.to(device="cuda", dtype=torch.bfloat16) + + >>> video = upsample_pipe( + ... video=video, + ... width=768, + ... height=512, + ... output_type="np", + ... return_dict=False, + ... )[0] + + >>> encode_video( + ... video[0], + ... fps=frame_rate, + ... audio=audio[0].float().cpu(), + ... audio_sample_rate=pipe.vocoder.config.output_sampling_rate, # should be 24000 + ... output_path="video.mp4", + ... ) + ``` +""" + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents +def retrieve_latents( + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" +): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +class LTX2LatentUpsamplePipeline(DiffusionPipeline): + model_cpu_offload_seq = "vae->latent_upsampler" + + def __init__( + self, + vae: AutoencoderKLLTX2Video, + latent_upsampler: LTX2LatentUpsamplerModel, + ) -> None: + super().__init__() + + self.register_modules(vae=vae, latent_upsampler=latent_upsampler) + + self.vae_spatial_compression_ratio = ( + self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 32 + ) + self.vae_temporal_compression_ratio = ( + self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 8 + ) + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio) + + def prepare_latents( + self, + video: torch.Tensor | None = None, + batch_size: int = 1, + num_frames: int = 121, + height: int = 512, + width: int = 768, + spatial_patch_size: int = 1, + temporal_patch_size: int = 1, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + if latents is not None: + if latents.ndim == 3: + # Convert token seq [B, S, D] to latent video [B, C, F, H, W] + latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 + latent_height = height // self.vae_spatial_compression_ratio + latent_width = width // self.vae_spatial_compression_ratio + latents = self._unpack_latents( + latents, latent_num_frames, latent_height, latent_width, spatial_patch_size, temporal_patch_size + ) + return latents.to(device=device, dtype=dtype) + + video = video.to(device=device, dtype=self.vae.dtype) + if isinstance(generator, list): + if len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + init_latents = [ + retrieve_latents(self.vae.encode(video[i].unsqueeze(0)), generator[i]) for i in range(batch_size) + ] + else: + init_latents = [retrieve_latents(self.vae.encode(vid.unsqueeze(0)), generator) for vid in video] + + init_latents = torch.cat(init_latents, dim=0).to(dtype) + # NOTE: latent upsampler operates on the unnormalized latents, so don't normalize here + # init_latents = self._normalize_latents(init_latents, self.vae.latents_mean, self.vae.latents_std) + return init_latents + + def adain_filter_latent(self, latents: torch.Tensor, reference_latents: torch.Tensor, factor: float = 1.0): + """ + Applies Adaptive Instance Normalization (AdaIN) to a latent tensor based on statistics from a reference latent + tensor. + + Args: + latent (`torch.Tensor`): + Input latents to normalize + reference_latents (`torch.Tensor`): + The reference latents providing style statistics. + factor (`float`): + Blending factor between original and transformed latent. Range: -10.0 to 10.0, Default: 1.0 + + Returns: + torch.Tensor: The transformed latent tensor + """ + result = latents.clone() + + for i in range(latents.size(0)): + for c in range(latents.size(1)): + r_sd, r_mean = torch.std_mean(reference_latents[i, c], dim=None) # index by original dim order + i_sd, i_mean = torch.std_mean(result[i, c], dim=None) + + result[i, c] = ((result[i, c] - i_mean) / i_sd) * r_sd + r_mean + + result = torch.lerp(latents, result, factor) + return result + + def tone_map_latents(self, latents: torch.Tensor, compression: float) -> torch.Tensor: + """ + Applies a non-linear tone-mapping function to latent values to reduce their dynamic range in a perceptually + smooth way using a sigmoid-based compression. + + This is useful for regularizing high-variance latents or for conditioning outputs during generation, especially + when controlling dynamic behavior with a `compression` factor. + + Args: + latents : torch.Tensor + Input latent tensor with arbitrary shape. Expected to be roughly in [-1, 1] or [0, 1] range. + compression : float + Compression strength in the range [0, 1]. + - 0.0: No tone-mapping (identity transform) + - 1.0: Full compression effect + + Returns: + torch.Tensor + The tone-mapped latent tensor of the same shape as input. + """ + # Remap [0-1] to [0-0.75] and apply sigmoid compression in one shot + scale_factor = compression * 0.75 + abs_latents = torch.abs(latents) + + # Sigmoid compression: sigmoid shifts large values toward 0.2, small values stay ~1.0 + # When scale_factor=0, sigmoid term vanishes, when scale_factor=0.75, full effect + sigmoid_term = torch.sigmoid(4.0 * scale_factor * (abs_latents - 1.0)) + scales = 1.0 - 0.8 * scale_factor * sigmoid_term + + filtered = latents * scales + return filtered + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._denormalize_latents + def _denormalize_latents( + latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 + ) -> torch.Tensor: + # Denormalize latents across the channel dimension [B, C, F, H, W] + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents = latents * latents_std / scaling_factor + latents_mean + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._unpack_latents + def _unpack_latents( + latents: torch.Tensor, num_frames: int, height: int, width: int, patch_size: int = 1, patch_size_t: int = 1 + ) -> torch.Tensor: + # Packed latents of shape [B, S, D] (S is the effective video sequence length, D is the effective feature dimensions) + # are unpacked and reshaped into a video tensor of shape [B, C, F, H, W]. This is the inverse operation of + # what happens in the `_pack_latents` method. + batch_size = latents.size(0) + latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size) + latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3) + return latents + + def check_inputs(self, video, height, width, latents, tone_map_compression_ratio): + if height % self.vae_spatial_compression_ratio != 0 or width % self.vae_spatial_compression_ratio != 0: + raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.") + + if video is not None and latents is not None: + raise ValueError("Only one of `video` or `latents` can be provided.") + if video is None and latents is None: + raise ValueError("One of `video` or `latents` has to be provided.") + + if not (0 <= tone_map_compression_ratio <= 1): + raise ValueError("`tone_map_compression_ratio` must be in the range [0, 1]") + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + video: list[PipelineImageInput] | None = None, + height: int = 512, + width: int = 768, + num_frames: int = 121, + spatial_patch_size: int = 1, + temporal_patch_size: int = 1, + latents: torch.Tensor | None = None, + latents_normalized: bool = False, + decode_timestep: float | list[float] = 0.0, + decode_noise_scale: float | list[float] | None = None, + adain_factor: float = 0.0, + tone_map_compression_ratio: float = 0.0, + generator: torch.Generator | list[torch.Generator] | None = None, + output_type: str | None = "pil", + return_dict: bool = True, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + video (`list[PipelineImageInput]`, *optional*) + The video to be upsampled (such as a LTX 2.0 first stage output). If not supplied, `latents` should be + supplied. + height (`int`, *optional*, defaults to `512`): + The height in pixels of the input video (not the generated video, which will have a larger resolution). + width (`int`, *optional*, defaults to `768`): + The width in pixels of the input video (not the generated video, which will have a larger resolution). + num_frames (`int`, *optional*, defaults to `121`): + The number of frames in the input video. + spatial_patch_size (`int`, *optional*, defaults to `1`): + The spatial patch size of the video latents. Used when `latents` is supplied if unpacking is necessary. + temporal_patch_size (`int`, *optional*, defaults to `1`): + The temporal patch size of the video latents. Used when `latents` is supplied if unpacking is + necessary. + latents (`torch.Tensor`, *optional*): + Pre-generated video latents. This can be supplied in place of the `video` argument. Can either be a + patch sequence of shape `(batch_size, seq_len, hidden_dim)` or a video latent of shape `(batch_size, + latent_channels, latent_frames, latent_height, latent_width)`. + latents_normalized (`bool`, *optional*, defaults to `False`) + If `latents` are supplied, whether the `latents` are normalized using the VAE latent mean and std. If + `True`, the `latents` will be denormalized before being supplied to the latent upsampler. + decode_timestep (`float`, defaults to `0.0`): + The timestep at which generated video is decoded. + decode_noise_scale (`float`, defaults to `None`): + The interpolation factor between random noise and denoised latents at the decode timestep. + adain_factor (`float`, *optional*, defaults to `0.0`): + Adaptive Instance Normalization (AdaIN) blending factor between the upsampled and original latents. + Should be in [-10.0, 10.0]; supplying 0.0 (the default) means that AdaIN is not performed. + tone_map_compression_ratio (`float`, *optional*, defaults to `0.0`): + The compression strength for tone mapping, which will reduce the dynamic range of the latent values. + This is useful for regularizing high-variance latents or for conditioning outputs during generation. + Should be in [0, 1], where 0.0 (the default) means tone mapping is not applied and 1.0 corresponds to + the full compression effect. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.ltx.LTXPipelineOutput`] instead of a plain tuple. + + Examples: + + Returns: + [`~pipelines.ltx.LTXPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`~pipelines.ltx.LTXPipelineOutput`] is returned, otherwise a `tuple` is + returned where the first element is the upsampled video. + """ + + self.check_inputs( + video=video, + height=height, + width=width, + latents=latents, + tone_map_compression_ratio=tone_map_compression_ratio, + ) + + if video is not None: + # Batched video input is not yet tested/supported. TODO: take a look later + batch_size = 1 + else: + batch_size = latents.shape[0] + device = self._execution_device + + if video is not None: + num_frames = len(video) + if num_frames % self.vae_temporal_compression_ratio != 1: + num_frames = ( + num_frames // self.vae_temporal_compression_ratio * self.vae_temporal_compression_ratio + 1 + ) + video = video[:num_frames] + logger.warning( + f"Video length expected to be of the form `k * {self.vae_temporal_compression_ratio} + 1` but is {len(video)}. Truncating to {num_frames} frames." + ) + video = self.video_processor.preprocess_video(video, height=height, width=width) + video = video.to(device=device, dtype=torch.float32) + + latents_supplied = latents is not None + latents = self.prepare_latents( + video=video, + batch_size=batch_size, + num_frames=num_frames, + height=height, + width=width, + spatial_patch_size=spatial_patch_size, + temporal_patch_size=temporal_patch_size, + dtype=torch.float32, + device=device, + generator=generator, + latents=latents, + ) + + if latents_supplied and latents_normalized: + latents = self._denormalize_latents( + latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor + ) + latents = latents.to(self.latent_upsampler.dtype) + latents_upsampled = self.latent_upsampler(latents) + + if adain_factor > 0.0: + latents = self.adain_filter_latent(latents_upsampled, latents, adain_factor) + else: + latents = latents_upsampled + + if tone_map_compression_ratio > 0.0: + latents = self.tone_map_latents(latents, tone_map_compression_ratio) + + if output_type == "latent": + video = latents + else: + if not self.vae.config.timestep_conditioning: + timestep = None + else: + noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype) + if not isinstance(decode_timestep, list): + decode_timestep = [decode_timestep] * batch_size + if decode_noise_scale is None: + decode_noise_scale = decode_timestep + elif not isinstance(decode_noise_scale, list): + decode_noise_scale = [decode_noise_scale] * batch_size + + timestep = torch.tensor(decode_timestep, device=device, dtype=latents.dtype) + decode_noise_scale = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[ + :, None, None, None, None + ] + latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise + + video = self.vae.decode(latents, timestep, return_dict=False)[0] + video = self.video_processor.postprocess_video(video, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (video,) + + return LTXPipelineOutput(frames=video) diff --git a/src/diffusers/pipelines/ltx2/pipeline_output.py b/src/diffusers/pipelines/ltx2/pipeline_output.py new file mode 100644 index 000000000000..02891219ee67 --- /dev/null +++ b/src/diffusers/pipelines/ltx2/pipeline_output.py @@ -0,0 +1,23 @@ +from dataclasses import dataclass + +import torch + +from diffusers.utils import BaseOutput + + +@dataclass +class LTX2PipelineOutput(BaseOutput): + r""" + Output class for LTX pipelines. + + Args: + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): + List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape + `(batch_size, num_frames, channels, height, width)`. + audio (`torch.Tensor`, `np.ndarray`): + TODO + """ + + frames: torch.Tensor + audio: torch.Tensor diff --git a/src/diffusers/pipelines/ltx2/utils.py b/src/diffusers/pipelines/ltx2/utils.py new file mode 100644 index 000000000000..f80469817fe6 --- /dev/null +++ b/src/diffusers/pipelines/ltx2/utils.py @@ -0,0 +1,6 @@ +# Pre-trained sigma values for distilled model are taken from +# https://github.com/Lightricks/LTX-2/blob/main/packages/ltx-pipelines/src/ltx_pipelines/utils/constants.py +DISTILLED_SIGMA_VALUES = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875] + +# Reduced schedule for super-resolution stage 2 (subset of distilled values) +STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875] diff --git a/src/diffusers/pipelines/ltx2/vocoder.py b/src/diffusers/pipelines/ltx2/vocoder.py new file mode 100644 index 000000000000..551c3ac5980f --- /dev/null +++ b/src/diffusers/pipelines/ltx2/vocoder.py @@ -0,0 +1,158 @@ +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...models.modeling_utils import ModelMixin + + +class ResBlock(nn.Module): + def __init__( + self, + channels: int, + kernel_size: int = 3, + stride: int = 1, + dilations: tuple[int, ...] = (1, 3, 5), + leaky_relu_negative_slope: float = 0.1, + padding_mode: str = "same", + ): + super().__init__() + self.dilations = dilations + self.negative_slope = leaky_relu_negative_slope + + self.convs1 = nn.ModuleList( + [ + nn.Conv1d(channels, channels, kernel_size, stride=stride, dilation=dilation, padding=padding_mode) + for dilation in dilations + ] + ) + + self.convs2 = nn.ModuleList( + [ + nn.Conv1d(channels, channels, kernel_size, stride=stride, dilation=1, padding=padding_mode) + for _ in range(len(dilations)) + ] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for conv1, conv2 in zip(self.convs1, self.convs2): + xt = F.leaky_relu(x, negative_slope=self.negative_slope) + xt = conv1(xt) + xt = F.leaky_relu(xt, negative_slope=self.negative_slope) + xt = conv2(xt) + x = x + xt + return x + + +class LTX2Vocoder(ModelMixin, ConfigMixin): + r""" + LTX 2.0 vocoder for converting generated mel spectrograms back to audio waveforms. + """ + + @register_to_config + def __init__( + self, + in_channels: int = 128, + hidden_channels: int = 1024, + out_channels: int = 2, + upsample_kernel_sizes: list[int] = [16, 15, 8, 4, 4], + upsample_factors: list[int] = [6, 5, 2, 2, 2], + resnet_kernel_sizes: list[int] = [3, 7, 11], + resnet_dilations: list[list[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + leaky_relu_negative_slope: float = 0.1, + output_sampling_rate: int = 24000, + ): + super().__init__() + self.num_upsample_layers = len(upsample_kernel_sizes) + self.resnets_per_upsample = len(resnet_kernel_sizes) + self.out_channels = out_channels + self.total_upsample_factor = math.prod(upsample_factors) + self.negative_slope = leaky_relu_negative_slope + + if self.num_upsample_layers != len(upsample_factors): + raise ValueError( + f"`upsample_kernel_sizes` and `upsample_factors` should be lists of the same length but are length" + f" {self.num_upsample_layers} and {len(upsample_factors)}, respectively." + ) + + if self.resnets_per_upsample != len(resnet_dilations): + raise ValueError( + f"`resnet_kernel_sizes` and `resnet_dilations` should be lists of the same length but are length" + f" {len(self.resnets_per_upsample)} and {len(resnet_dilations)}, respectively." + ) + + self.conv_in = nn.Conv1d(in_channels, hidden_channels, kernel_size=7, stride=1, padding=3) + + self.upsamplers = nn.ModuleList() + self.resnets = nn.ModuleList() + input_channels = hidden_channels + for i, (stride, kernel_size) in enumerate(zip(upsample_factors, upsample_kernel_sizes)): + output_channels = input_channels // 2 + self.upsamplers.append( + nn.ConvTranspose1d( + input_channels, # hidden_channels // (2 ** i) + output_channels, # hidden_channels // (2 ** (i + 1)) + kernel_size, + stride=stride, + padding=(kernel_size - stride) // 2, + ) + ) + + for kernel_size, dilations in zip(resnet_kernel_sizes, resnet_dilations): + self.resnets.append( + ResBlock( + output_channels, + kernel_size, + dilations=dilations, + leaky_relu_negative_slope=leaky_relu_negative_slope, + ) + ) + input_channels = output_channels + + self.conv_out = nn.Conv1d(output_channels, out_channels, 7, stride=1, padding=3) + + def forward(self, hidden_states: torch.Tensor, time_last: bool = False) -> torch.Tensor: + r""" + Forward pass of the vocoder. + + Args: + hidden_states (`torch.Tensor`): + Input Mel spectrogram tensor of shape `(batch_size, num_channels, time, num_mel_bins)` if `time_last` + is `False` (the default) or shape `(batch_size, num_channels, num_mel_bins, time)` if `time_last` is + `True`. + time_last (`bool`, *optional*, defaults to `False`): + Whether the last dimension of the input is the time/frame dimension or the Mel bins dimension. + + Returns: + `torch.Tensor`: + Audio waveform tensor of shape (batch_size, out_channels, audio_length) + """ + + # Ensure that the time/frame dimension is last + if not time_last: + hidden_states = hidden_states.transpose(2, 3) + # Combine channels and frequency (mel bins) dimensions + hidden_states = hidden_states.flatten(1, 2) + + hidden_states = self.conv_in(hidden_states) + + for i in range(self.num_upsample_layers): + hidden_states = F.leaky_relu(hidden_states, negative_slope=self.negative_slope) + hidden_states = self.upsamplers[i](hidden_states) + + # Run all resnets in parallel on hidden_states + start = i * self.resnets_per_upsample + end = (i + 1) * self.resnets_per_upsample + resnet_outputs = torch.stack([self.resnets[j](hidden_states) for j in range(start, end)], dim=0) + + hidden_states = torch.mean(resnet_outputs, dim=0) + + # NOTE: unlike the first leaky ReLU, this leaky ReLU is set to use the default F.leaky_relu negative slope of + # 0.01 (whereas the others usually use a slope of 0.1). Not sure if this is intended + hidden_states = F.leaky_relu(hidden_states, negative_slope=0.01) + hidden_states = self.conv_out(hidden_states) + hidden_states = torch.tanh(hidden_states) + + return hidden_states diff --git a/src/diffusers/pipelines/lucy/pipeline_lucy_edit.py b/src/diffusers/pipelines/lucy/pipeline_lucy_edit.py index 69f69d5768a8..392af492b702 100644 --- a/src/diffusers/pipelines/lucy/pipeline_lucy_edit.py +++ b/src/diffusers/pipelines/lucy/pipeline_lucy_edit.py @@ -14,10 +14,10 @@ # limitations under the License. # # Modifications by Decart AI Team: -# - Based on pipeline_wan.py, but with supports recieving a condition video appended to the channel dimension. +# - Based on pipeline_wan.py, but with supports receiving a condition video appended to the channel dimension. import html -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import regex as re import torch @@ -51,7 +51,7 @@ EXAMPLE_DOC_STRING = """ Examples: ```python - >>> from typing import List + >>> from typing import list >>> import torch >>> from PIL import Image @@ -69,7 +69,7 @@ >>> # Load video - >>> def convert_video(video: List[Image.Image]) -> List[Image.Image]: + >>> def convert_video(video: list[Image.Image]) -> list[Image.Image]: ... video = load_video(url)[:num_frames] ... video = [video[i].resize((width, height)) for i in range(num_frames)] ... return video @@ -119,7 +119,7 @@ def prompt_clean(text): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -172,9 +172,9 @@ def __init__( text_encoder: UMT5EncoderModel, vae: AutoencoderKLWan, scheduler: FlowMatchEulerDiscreteScheduler, - transformer: Optional[WanTransformer3DModel] = None, - transformer_2: Optional[WanTransformer3DModel] = None, - boundary_ratio: Optional[float] = None, + transformer: WanTransformer3DModel | None = None, + transformer_2: WanTransformer3DModel | None = None, + boundary_ratio: float | None = None, expand_timesteps: bool = False, # Wan2.2 ti2v ): super().__init__() @@ -196,11 +196,11 @@ def __init__( # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -238,23 +238,23 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -368,16 +368,16 @@ def check_inputs( def prepare_latents( self, - video: Optional[torch.Tensor] = None, + video: torch.Tensor | None = None, batch_size: int = 1, num_channels_latents: int = 16, height: int = 480, width: int = 832, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[torch.Generator] = None, - latents: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" @@ -451,38 +451,36 @@ def attention_kwargs(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - video: List[Image.Image], - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + video: list[Image.Image], + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, height: int = 480, width: int = 832, num_frames: int = 81, num_inference_steps: int = 50, guidance_scale: float = 5.0, - guidance_scale_2: Optional[float] = None, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "np", + guidance_scale_2: float | None = None, + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "np", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" The call function to the pipeline for generation. Args: - video (`List[Image.Image]`): + video (`list[Image.Image]`): The video to use as the condition for the video generation. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, pass `prompt_embeds` instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to avoid during image generation. If not defined, pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale` < `1`). height (`int`, defaults to `480`): @@ -506,7 +504,7 @@ def __call__( and the pipeline's `boundary_ratio` are not None. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -529,7 +527,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/lucy/pipeline_output.py b/src/diffusers/pipelines/lucy/pipeline_output.py index cf9ea91fd106..197ce194f475 100644 --- a/src/diffusers/pipelines/lucy/pipeline_output.py +++ b/src/diffusers/pipelines/lucy/pipeline_output.py @@ -11,8 +11,8 @@ class LucyPipelineOutput(BaseOutput): Output class for Lucy pipelines. Args: - frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): - List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): + list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape `(batch_size, num_frames, channels, height, width)`. """ diff --git a/src/diffusers/pipelines/lumina/pipeline_lumina.py b/src/diffusers/pipelines/lumina/pipeline_lumina.py index b59c265646cd..cc123218f4ee 100644 --- a/src/diffusers/pipelines/lumina/pipeline_lumina.py +++ b/src/diffusers/pipelines/lumina/pipeline_lumina.py @@ -17,7 +17,7 @@ import math import re import urllib.parse as ul -from typing import Callable, Dict, List, Optional, Tuple, Union +from typing import Callable import torch from transformers import GemmaPreTrainedModel, GemmaTokenizer, GemmaTokenizerFast @@ -76,10 +76,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -94,15 +94,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -182,7 +182,7 @@ def __init__( scheduler: FlowMatchEulerDiscreteScheduler, vae: AutoencoderKL, text_encoder: GemmaPreTrainedModel, - tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast], + tokenizer: GemmaTokenizer | GemmaTokenizerFast, ): super().__init__() @@ -205,11 +205,11 @@ def __init__( def _get_gemma_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - clean_caption: Optional[bool] = False, - max_length: Optional[int] = None, + device: torch.device | None = None, + clean_caption: bool | None = False, + max_length: int | None = None, ): device = device or self._execution_device prompt = [prompt] if isinstance(prompt, str) else prompt @@ -261,15 +261,15 @@ def _get_gemma_prompt_embeds( # Adapted from diffusers.pipelines.deepfloyd_if.pipeline_if.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], do_classifier_free_guidance: bool = True, - negative_prompt: Union[str, List[str]] = None, + negative_prompt: str | list[str] = None, num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, + device: torch.device | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, clean_caption: bool = False, **kwargs, ): @@ -277,9 +277,9 @@ def encode_prompt( Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For Lumina-T2I, this should be "". @@ -633,46 +633,44 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - width: Optional[int] = None, - height: Optional[int] = None, + prompt: str | list[str] = None, + width: int | None = None, + height: int | None = None, num_inference_steps: int = 30, guidance_scale: float = 4.0, - negative_prompt: Union[str, List[str]] = None, - sigmas: List[float] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + negative_prompt: str | list[str] = None, + sigmas: list[float] = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, clean_caption: bool = True, max_sequence_length: int = 256, - scaling_watershed: Optional[float] = 1.0, - proportional_attn: Optional[bool] = True, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], - ) -> Union[ImagePipelineOutput, Tuple]: + scaling_watershed: float | None = 1.0, + proportional_attn: bool | None = True, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + ) -> ImagePipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_inference_steps (`int`, *optional*, defaults to 30): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -691,7 +689,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -723,7 +721,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -799,7 +797,13 @@ def __call__( prompt_attention_mask = torch.cat([prompt_attention_mask, negative_prompt_attention_mask], dim=0) # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, sigmas=sigmas) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, sigmas=sigmas + ) # 5. Prepare latents. latent_channels = self.transformer.config.in_channels @@ -940,7 +944,7 @@ def __init__( scheduler: FlowMatchEulerDiscreteScheduler, vae: AutoencoderKL, text_encoder: GemmaPreTrainedModel, - tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast], + tokenizer: GemmaTokenizer | GemmaTokenizerFast, ): deprecation_message = "`LuminaText2ImgPipeline` has been renamed to `LuminaPipeline` and will be removed in a future version. Please use `LuminaPipeline` instead." deprecate("diffusers.pipelines.lumina.pipeline_lumina.LuminaText2ImgPipeline", "0.34", deprecation_message) diff --git a/src/diffusers/pipelines/lumina2/pipeline_lumina2.py b/src/diffusers/pipelines/lumina2/pipeline_lumina2.py index 937803edbcbc..576d3e8d9486 100644 --- a/src/diffusers/pipelines/lumina2/pipeline_lumina2.py +++ b/src/diffusers/pipelines/lumina2/pipeline_lumina2.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import torch @@ -77,10 +77,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -95,15 +95,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -164,7 +164,7 @@ def __init__( scheduler: FlowMatchEulerDiscreteScheduler, vae: AutoencoderKL, text_encoder: Gemma2PreTrainedModel, - tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast], + tokenizer: GemmaTokenizer | GemmaTokenizerFast, ): super().__init__() @@ -191,10 +191,10 @@ def __init__( def _get_gemma_prompt_embeds( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + device: torch.device | None = None, max_sequence_length: int = 256, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: device = device or self._execution_device prompt = [prompt] if isinstance(prompt, str) else prompt text_inputs = self.tokenizer( @@ -237,25 +237,25 @@ def _get_gemma_prompt_embeds( # Adapted from diffusers.pipelines.deepfloyd_if.pipeline_if.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], do_classifier_free_guidance: bool = True, - negative_prompt: Union[str, List[str]] = None, + negative_prompt: str | list[str] = None, num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - system_prompt: Optional[str] = None, + device: torch.device | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + system_prompt: str | None = None, max_sequence_length: int = 256, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For Lumina-T2I, this should be "". @@ -525,45 +525,45 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - width: Optional[int] = None, - height: Optional[int] = None, + prompt: str | list[str] = None, + width: int | None = None, + height: int | None = None, num_inference_steps: int = 30, guidance_scale: float = 4.0, - negative_prompt: Union[str, List[str]] = None, - sigmas: List[float] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + negative_prompt: str | list[str] = None, + sigmas: list[float] = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], - system_prompt: Optional[str] = None, + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + system_prompt: str | None = None, cfg_trunc_ratio: float = 1.0, cfg_normalization: bool = True, max_sequence_length: int = 256, - ) -> Union[ImagePipelineOutput, Tuple]: + ) -> ImagePipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_inference_steps (`int`, *optional*, defaults to 30): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -582,7 +582,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -612,7 +612,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -704,10 +704,14 @@ def __call__( self.scheduler.config.get("base_shift", 0.5), self.scheduler.config.get("max_shift", 1.15), ) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, - device, + timestep_device, sigmas=sigmas, mu=mu, ) @@ -801,7 +805,7 @@ def __init__( scheduler: FlowMatchEulerDiscreteScheduler, vae: AutoencoderKL, text_encoder: Gemma2PreTrainedModel, - tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast], + tokenizer: GemmaTokenizer | GemmaTokenizerFast, ): deprecation_message = "`Lumina2Text2ImgPipeline` has been renamed to `Lumina2Pipeline` and will be removed in a future version. Please use `Lumina2Pipeline` instead." deprecate("diffusers.pipelines.lumina2.pipeline_lumina2.Lumina2Text2ImgPipeline", "0.34", deprecation_message) diff --git a/src/diffusers/pipelines/marigold/marigold_image_processing.py b/src/diffusers/pipelines/marigold/marigold_image_processing.py index 5130a876606a..4a76ef9892e4 100644 --- a/src/diffusers/pipelines/marigold/marigold_image_processing.py +++ b/src/diffusers/pipelines/marigold/marigold_image_processing.py @@ -16,7 +16,7 @@ # More information and citation instructions are available on the # Marigold project website: https://marigoldcomputervision.github.io # -------------------------------------------------------------------------- -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import numpy as np import PIL @@ -47,7 +47,7 @@ def __init__( super().__init__() @staticmethod - def expand_tensor_or_array(images: Union[torch.Tensor, np.ndarray]) -> Union[torch.Tensor, np.ndarray]: + def expand_tensor_or_array(images: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray: """ Expand a tensor or array to a specified number of images. """ @@ -90,7 +90,7 @@ def numpy_to_pt(images: np.ndarray) -> torch.Tensor: @staticmethod def resize_antialias( - image: torch.Tensor, size: Tuple[int, int], mode: str, is_aa: Optional[bool] = None + image: torch.Tensor, size: tuple[int, int], mode: str, is_aa: bool | None = None ) -> torch.Tensor: if not torch.is_tensor(image): raise ValueError(f"Invalid input type={type(image)}.") @@ -126,7 +126,7 @@ def resize_to_max_edge(image: torch.Tensor, max_edge_sz: int, mode: str) -> torc return image @staticmethod - def pad_image(image: torch.Tensor, align: int) -> Tuple[torch.Tensor, Tuple[int, int]]: + def pad_image(image: torch.Tensor, align: int) -> tuple[torch.Tensor, tuple[int, int]]: if not torch.is_tensor(image): raise ValueError(f"Invalid input type={type(image)}.") if not torch.is_floating_point(image): @@ -142,7 +142,7 @@ def pad_image(image: torch.Tensor, align: int) -> Tuple[torch.Tensor, Tuple[int, return image, (ph, pw) @staticmethod - def unpad_image(image: torch.Tensor, padding: Tuple[int, int]) -> torch.Tensor: + def unpad_image(image: torch.Tensor, padding: tuple[int, int]) -> torch.Tensor: if not torch.is_tensor(image): raise ValueError(f"Invalid input type={type(image)}.") if not torch.is_floating_point(image): @@ -160,10 +160,10 @@ def unpad_image(image: torch.Tensor, padding: Tuple[int, int]) -> torch.Tensor: @staticmethod def load_image_canonical( - image: Union[torch.Tensor, np.ndarray, Image.Image], + image: torch.Tensor | np.ndarray | Image.Image, device: torch.device = torch.device("cpu"), dtype: torch.dtype = torch.float32, - ) -> Tuple[torch.Tensor, int]: + ) -> tuple[torch.Tensor, int]: if isinstance(image, Image.Image): image = np.array(image) @@ -216,7 +216,7 @@ def check_image_values_range(image: torch.Tensor) -> None: def preprocess( self, image: PipelineImageInput, - processing_resolution: Optional[int] = None, + processing_resolution: int | None = None, resample_method_input: str = "bilinear", device: torch.device = torch.device("cpu"), dtype: torch.dtype = torch.float32, @@ -256,11 +256,11 @@ def preprocess( @staticmethod def colormap( - image: Union[np.ndarray, torch.Tensor], + image: np.ndarray | torch.Tensor, cmap: str = "Spectral", bytes: bool = False, - _force_method: Optional[str] = None, - ) -> Union[np.ndarray, torch.Tensor]: + _force_method: str | None = None, + ) -> np.ndarray | torch.Tensor: """ Converts a monochrome image into an RGB image by applying the specified colormap. This function mimics the behavior of matplotlib.colormaps, but allows the user to use the most discriminative color maps ("Spectral", @@ -386,30 +386,28 @@ def method_custom(image, cmap, bytes=False): @staticmethod def visualize_depth( - depth: Union[ - PIL.Image.Image, - np.ndarray, - torch.Tensor, - List[PIL.Image.Image], - List[np.ndarray], - List[torch.Tensor], - ], + depth: PIL.Image.Image + | np.ndarray + | torch.Tensor + | list[PIL.Image.Image] + | list[np.ndarray] + | list[torch.Tensor], val_min: float = 0.0, val_max: float = 1.0, color_map: str = "Spectral", - ) -> List[PIL.Image.Image]: + ) -> list[PIL.Image.Image]: """ Visualizes depth maps, such as predictions of the `MarigoldDepthPipeline`. Args: - depth (`Union[PIL.Image.Image, np.ndarray, torch.Tensor, List[PIL.Image.Image], List[np.ndarray], - List[torch.Tensor]]`): Depth maps. + depth (`PIL.Image.Image | np.ndarray | torch.Tensor | list[PIL.Image.Image, list[np.ndarray], + list[torch.Tensor]]`): Depth maps. val_min (`float`, *optional*, defaults to `0.0`): Minimum value of the visualized depth range. val_max (`float`, *optional*, defaults to `1.0`): Maximum value of the visualized depth range. color_map (`str`, *optional*, defaults to `"Spectral"`): Color map used to convert a single-channel depth prediction into colored representation. - Returns: `List[PIL.Image.Image]` with depth maps visualization. + Returns: `list[PIL.Image.Image]` with depth maps visualization. """ if val_max <= val_min: raise ValueError(f"Invalid values range: [{val_min}, {val_max}].") @@ -451,10 +449,10 @@ def visualize_depth_one(img, idx=None): @staticmethod def export_depth_to_16bit_png( - depth: Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]], + depth: np.ndarray | torch.Tensor | list[np.ndarray] | list[torch.Tensor], val_min: float = 0.0, val_max: float = 1.0, - ) -> List[PIL.Image.Image]: + ) -> list[PIL.Image.Image]: def export_depth_to_16bit_png_one(img, idx=None): prefix = "Depth" + (f"[{idx}]" if idx else "") if not isinstance(img, np.ndarray) and not torch.is_tensor(img): @@ -487,21 +485,16 @@ def export_depth_to_16bit_png_one(img, idx=None): @staticmethod def visualize_normals( - normals: Union[ - np.ndarray, - torch.Tensor, - List[np.ndarray], - List[torch.Tensor], - ], + normals: np.ndarray | torch.Tensor | list[np.ndarray] | list[torch.Tensor], flip_x: bool = False, flip_y: bool = False, flip_z: bool = False, - ) -> List[PIL.Image.Image]: + ) -> list[PIL.Image.Image]: """ Visualizes surface normals, such as predictions of the `MarigoldNormalsPipeline`. Args: - normals (`Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]]`): + normals (`np.ndarray | torch.Tensor | list[np.ndarray, list[torch.Tensor]]`): Surface normals. flip_x (`bool`, *optional*, defaults to `False`): Flips the X axis of the normals frame of reference. Default direction is right. @@ -510,7 +503,7 @@ def visualize_normals( flip_z (`bool`, *optional*, defaults to `False`): Flips the Z axis of the normals frame of reference. Default direction is facing the observer. - Returns: `List[PIL.Image.Image]` with surface normals visualization. + Returns: `list[PIL.Image.Image]` with surface normals visualization. """ flip_vec = None if any((flip_x, flip_y, flip_z)): @@ -548,30 +541,25 @@ def visualize_normals_one(img, idx=None): @staticmethod def visualize_intrinsics( - prediction: Union[ - np.ndarray, - torch.Tensor, - List[np.ndarray], - List[torch.Tensor], - ], - target_properties: Dict[str, Any], - color_map: Union[str, Dict[str, str]] = "binary", - ) -> List[Dict[str, PIL.Image.Image]]: + prediction: np.ndarray | torch.Tensor | list[np.ndarray] | list[torch.Tensor], + target_properties: dict[str, Any], + color_map: str | dict[str, str] = "binary", + ) -> list[dict[str, PIL.Image.Image]]: """ Visualizes intrinsic image decomposition, such as predictions of the `MarigoldIntrinsicsPipeline`. Args: - prediction (`Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]]`): + prediction (`np.ndarray | torch.Tensor | list[np.ndarray, list[torch.Tensor]]`): Intrinsic image decomposition. - target_properties (`Dict[str, Any]`): - Decomposition properties. Expected entries: `target_names: List[str]` and a dictionary with keys - `prediction_space: str`, `sub_target_names: List[Union[str, Null]]` (must have 3 entries, null for - missing modalities), `up_to_scale: bool`, one for each target and sub-target. - color_map (`Union[str, Dict[str, str]]`, *optional*, defaults to `"Spectral"`): + target_properties (`dict[str, Any]`): + Decomposition properties. Expected entries: `target_names: list[str]` and a dictionary with keys + `prediction_space: str`, `sub_target_names: list[str | Null]` (must have 3 entries, null for missing + modalities), `up_to_scale: bool`, one for each target and sub-target. + color_map (`str | dict[str, str]`, *optional*, defaults to `"Spectral"`): Color map used to convert a single-channel predictions into colored representations. When a dictionary is passed, each modality can be colored with its own color map. - Returns: `List[Dict[str, PIL.Image.Image]]` with intrinsic image decomposition visualization. + Returns: `list[dict[str, PIL.Image.Image]]` with intrinsic image decomposition visualization. """ if "target_names" not in target_properties: raise ValueError("Missing `target_names` in target_properties") @@ -641,25 +629,20 @@ def visualize_targets_one(images, idx=None): @staticmethod def visualize_uncertainty( - uncertainty: Union[ - np.ndarray, - torch.Tensor, - List[np.ndarray], - List[torch.Tensor], - ], + uncertainty: np.ndarray | torch.Tensor | list[np.ndarray] | list[torch.Tensor], saturation_percentile=95, - ) -> List[PIL.Image.Image]: + ) -> list[PIL.Image.Image]: """ Visualizes dense uncertainties, such as produced by `MarigoldDepthPipeline`, `MarigoldNormalsPipeline`, or `MarigoldIntrinsicsPipeline`. Args: - uncertainty (`Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]]`): + uncertainty (`np.ndarray | torch.Tensor | list[np.ndarray, list[torch.Tensor]]`): Uncertainty maps. saturation_percentile (`int`, *optional*, defaults to `95`): Specifies the percentile uncertainty value visualized with maximum intensity. - Returns: `List[PIL.Image.Image]` with uncertainty visualization. + Returns: `list[PIL.Image.Image]` with uncertainty visualization. """ def visualize_uncertainty_one(img, idx=None): diff --git a/src/diffusers/pipelines/marigold/pipeline_marigold_depth.py b/src/diffusers/pipelines/marigold/pipeline_marigold_depth.py index 92ec16fd455b..a81d1c51742c 100644 --- a/src/diffusers/pipelines/marigold/pipeline_marigold_depth.py +++ b/src/diffusers/pipelines/marigold/pipeline_marigold_depth.py @@ -18,7 +18,7 @@ # -------------------------------------------------------------------------- from dataclasses import dataclass from functools import partial -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import numpy as np import torch @@ -96,9 +96,9 @@ class MarigoldDepthOutput(BaseOutput): The shape is `numimages * numensemble × 4 × latentheight × latentwidth`. """ - prediction: Union[np.ndarray, torch.Tensor] - uncertainty: Union[None, np.ndarray, torch.Tensor] - latent: Union[None, torch.Tensor] + prediction: np.ndarray | torch.Tensor + uncertainty: None | np.ndarray | torch.Tensor + latent: None | torch.Tensor class MarigoldDepthPipeline(DiffusionPipeline): @@ -150,14 +150,14 @@ def __init__( self, unet: UNet2DConditionModel, vae: AutoencoderKL, - scheduler: Union[DDIMScheduler, LCMScheduler], + scheduler: DDIMScheduler | LCMScheduler, text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, - prediction_type: Optional[str] = None, - scale_invariant: Optional[bool] = True, - shift_invariant: Optional[bool] = True, - default_denoising_steps: Optional[int] = None, - default_processing_resolution: Optional[int] = None, + prediction_type: str | None = None, + scale_invariant: bool | None = True, + shift_invariant: bool | None = True, + default_denoising_steps: int | None = None, + default_processing_resolution: int | None = None, ): super().__init__() @@ -202,9 +202,9 @@ def check_inputs( resample_method_input: str, resample_method_output: str, batch_size: int, - ensembling_kwargs: Optional[Dict[str, Any]], - latents: Optional[torch.Tensor], - generator: Optional[Union[torch.Generator, List[torch.Generator]]], + ensembling_kwargs: dict[str, Any] | None, + latents: torch.Tensor | None, + generator: torch.Generator | list[torch.Generator] | None, output_type: str, output_uncertainty: bool, ) -> int: @@ -349,16 +349,16 @@ def progress_bar(self, iterable=None, total=None, desc=None, leave=True): def __call__( self, image: PipelineImageInput, - num_inference_steps: Optional[int] = None, + num_inference_steps: int | None = None, ensemble_size: int = 1, - processing_resolution: Optional[int] = None, + processing_resolution: int | None = None, match_input_resolution: bool = True, resample_method_input: str = "bilinear", resample_method_output: str = "bilinear", batch_size: int = 1, - ensembling_kwargs: Optional[Dict[str, Any]] = None, - latents: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + ensembling_kwargs: dict[str, Any] | None = None, + latents: torch.Tensor | list[torch.Tensor] | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, output_type: str = "np", output_uncertainty: bool = False, output_latent: bool = False, @@ -368,8 +368,8 @@ def __call__( Function invoked when calling the pipeline. Args: - image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`), - `List[torch.Tensor]`: An input image or images used as an input for the depth estimation task. For + image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`), + `list[torch.Tensor]`: An input image or images used as an input for the depth estimation task. For arrays and tensors, the expected value range is between `[0, 1]`. Passing a batch of images is possible by providing a four-dimensional array or a tensor. Additionally, a list of images of two- or three-dimensional arrays or tensors can be passed. In the latter case, all list elements must have the @@ -406,10 +406,10 @@ def __call__( tolerance is reached. - max_res (`int`, *optional*, defaults to `None`): Resolution at which the alignment is performed; `None` matches the `processing_resolution`. - latents (`torch.Tensor`, or `List[torch.Tensor]`, *optional*, defaults to `None`): + latents (`torch.Tensor`, or `list[torch.Tensor]`, *optional*, defaults to `None`): Latent noise tensors to replace the random initialization. These can be taken from the previous function call's output. - generator (`torch.Generator`, or `List[torch.Generator]`, *optional*, defaults to `None`): + generator (`torch.Generator`, or `list[torch.Generator]`, *optional*, defaults to `None`): Random number generator object to ensure reproducibility. output_type (`str`, *optional*, defaults to `"np"`): Preferred format of the output's `prediction` and the optional `uncertainty` fields. The accepted @@ -621,11 +621,11 @@ def __call__( def prepare_latents( self, image: torch.Tensor, - latents: Optional[torch.Tensor], - generator: Optional[torch.Generator], + latents: torch.Tensor | None, + generator: torch.Generator | None, ensemble_size: int, batch_size: int, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: def retrieve_latents(encoder_output): if hasattr(encoder_output, "latent_dist"): return encoder_output.latent_dist.mode() @@ -680,7 +680,7 @@ def ensemble_depth( max_iter: int = 2, tol: float = 1e-3, max_res: int = 1024, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + ) -> tuple[torch.Tensor, torch.Tensor | None]: """ Ensembles the depth maps represented by the `depth` tensor with expected shape `(B, 1, H, W)`, where B is the number of ensemble members for a given prediction of size `(H x W)`. Even though the function is designed for @@ -754,7 +754,7 @@ def align(depth: torch.Tensor, param: np.ndarray) -> torch.Tensor: def ensemble( depth_aligned: torch.Tensor, return_uncertainty: bool = False - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + ) -> tuple[torch.Tensor, torch.Tensor | None]: uncertainty = None if reduction == "mean": prediction = torch.mean(depth_aligned, dim=0, keepdim=True) diff --git a/src/diffusers/pipelines/marigold/pipeline_marigold_intrinsics.py b/src/diffusers/pipelines/marigold/pipeline_marigold_intrinsics.py index bef9ca77c708..9488d8f5c9b8 100644 --- a/src/diffusers/pipelines/marigold/pipeline_marigold_intrinsics.py +++ b/src/diffusers/pipelines/marigold/pipeline_marigold_intrinsics.py @@ -17,7 +17,7 @@ # Marigold project website: https://marigoldcomputervision.github.io # -------------------------------------------------------------------------- from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import numpy as np import torch @@ -112,9 +112,9 @@ class MarigoldIntrinsicsOutput(BaseOutput): The shape is `(numimages * numensemble) × (numtargets * 4) × latentheight × latentwidth`. """ - prediction: Union[np.ndarray, torch.Tensor] - uncertainty: Union[None, np.ndarray, torch.Tensor] - latent: Union[None, torch.Tensor] + prediction: np.ndarray | torch.Tensor + uncertainty: None | np.ndarray | torch.Tensor + latent: None | torch.Tensor class MarigoldIntrinsicsPipeline(DiffusionPipeline): @@ -139,8 +139,8 @@ class MarigoldIntrinsicsPipeline(DiffusionPipeline): CLIP tokenizer. prediction_type (`str`, *optional*): Type of predictions made by the model. - target_properties (`Dict[str, Any]`, *optional*): - Properties of the predicted modalities, such as `target_names`, a `List[str]` used to define the number, + target_properties (`dict[str, Any]`, *optional*): + Properties of the predicted modalities, such as `target_names`, a `list[str]` used to define the number, order and names of the predicted modalities, and any other metadata that may be required to interpret the predictions. default_denoising_steps (`int`, *optional*): @@ -163,13 +163,13 @@ def __init__( self, unet: UNet2DConditionModel, vae: AutoencoderKL, - scheduler: Union[DDIMScheduler, LCMScheduler], + scheduler: DDIMScheduler | LCMScheduler, text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, - prediction_type: Optional[str] = None, - target_properties: Optional[Dict[str, Any]] = None, - default_denoising_steps: Optional[int] = None, - default_processing_resolution: Optional[int] = None, + prediction_type: str | None = None, + target_properties: dict[str, Any] | None = None, + default_denoising_steps: int | None = None, + default_processing_resolution: int | None = None, ): super().__init__() @@ -216,9 +216,9 @@ def check_inputs( resample_method_input: str, resample_method_output: str, batch_size: int, - ensembling_kwargs: Optional[Dict[str, Any]], - latents: Optional[torch.Tensor], - generator: Optional[Union[torch.Generator, List[torch.Generator]]], + ensembling_kwargs: dict[str, Any] | None, + latents: torch.Tensor | None, + generator: torch.Generator | list[torch.Generator] | None, output_type: str, output_uncertainty: bool, ) -> int: @@ -361,16 +361,16 @@ def progress_bar(self, iterable=None, total=None, desc=None, leave=True): def __call__( self, image: PipelineImageInput, - num_inference_steps: Optional[int] = None, + num_inference_steps: int | None = None, ensemble_size: int = 1, - processing_resolution: Optional[int] = None, + processing_resolution: int | None = None, match_input_resolution: bool = True, resample_method_input: str = "bilinear", resample_method_output: str = "bilinear", batch_size: int = 1, - ensembling_kwargs: Optional[Dict[str, Any]] = None, - latents: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + ensembling_kwargs: dict[str, Any] | None = None, + latents: torch.Tensor | list[torch.Tensor] | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, output_type: str = "np", output_uncertainty: bool = False, output_latent: bool = False, @@ -380,8 +380,8 @@ def __call__( Function invoked when calling the pipeline. Args: - image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`), - `List[torch.Tensor]`: An input image or images used as an input for the intrinsic decomposition task. + image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`), + `list[torch.Tensor]`: An input image or images used as an input for the intrinsic decomposition task. For arrays and tensors, the expected value range is between `[0, 1]`. Passing a batch of images is possible by providing a four-dimensional array or a tensor. Additionally, a list of images of two- or three-dimensional arrays or tensors can be passed. In the latter case, all list elements must have the @@ -413,7 +413,7 @@ def __call__( latents (`torch.Tensor`, *optional*, defaults to `None`): Latent noise tensors to replace the random initialization. These can be taken from the previous function call's output. - generator (`torch.Generator`, or `List[torch.Generator]`, *optional*, defaults to `None`): + generator (`torch.Generator`, or `list[torch.Generator]`, *optional*, defaults to `None`): Random number generator object to ensure reproducibility. output_type (`str`, *optional*, defaults to `"np"`): Preferred format of the output's `prediction` and the optional `uncertainty` fields. The accepted @@ -627,11 +627,11 @@ def __call__( def prepare_latents( self, image: torch.Tensor, - latents: Optional[torch.Tensor], - generator: Optional[torch.Generator], + latents: torch.Tensor | None, + generator: torch.Generator | None, ensemble_size: int, batch_size: int, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: def retrieve_latents(encoder_output): if hasattr(encoder_output, "latent_dist"): return encoder_output.latent_dist.mode() @@ -680,7 +680,7 @@ def ensemble_intrinsics( targets: torch.Tensor, output_uncertainty: bool = False, reduction: str = "median", - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + ) -> tuple[torch.Tensor, torch.Tensor | None]: """ Ensembles the intrinsic decomposition represented by the `targets` tensor with expected shape `(B, T, 3, H, W)`, where B is the number of ensemble members for a given prediction of size `(H x W)`, and T is the number of diff --git a/src/diffusers/pipelines/marigold/pipeline_marigold_normals.py b/src/diffusers/pipelines/marigold/pipeline_marigold_normals.py index 485a39c995ec..3f94ce441232 100644 --- a/src/diffusers/pipelines/marigold/pipeline_marigold_normals.py +++ b/src/diffusers/pipelines/marigold/pipeline_marigold_normals.py @@ -17,7 +17,7 @@ # Marigold project website: https://marigoldcomputervision.github.io # -------------------------------------------------------------------------- from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import numpy as np import torch @@ -91,9 +91,9 @@ class MarigoldNormalsOutput(BaseOutput): The shape is `numimages * numensemble × 4 × latentheight × latentwidth`. """ - prediction: Union[np.ndarray, torch.Tensor] - uncertainty: Union[None, np.ndarray, torch.Tensor] - latent: Union[None, torch.Tensor] + prediction: np.ndarray | torch.Tensor + uncertainty: None | np.ndarray | torch.Tensor + latent: None | torch.Tensor class MarigoldNormalsPipeline(DiffusionPipeline): @@ -140,13 +140,13 @@ def __init__( self, unet: UNet2DConditionModel, vae: AutoencoderKL, - scheduler: Union[DDIMScheduler, LCMScheduler], + scheduler: DDIMScheduler | LCMScheduler, text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, - prediction_type: Optional[str] = None, - use_full_z_range: Optional[bool] = True, - default_denoising_steps: Optional[int] = None, - default_processing_resolution: Optional[int] = None, + prediction_type: str | None = None, + use_full_z_range: bool | None = True, + default_denoising_steps: int | None = None, + default_processing_resolution: int | None = None, ): super().__init__() @@ -189,9 +189,9 @@ def check_inputs( resample_method_input: str, resample_method_output: str, batch_size: int, - ensembling_kwargs: Optional[Dict[str, Any]], - latents: Optional[torch.Tensor], - generator: Optional[Union[torch.Generator, List[torch.Generator]]], + ensembling_kwargs: dict[str, Any] | None, + latents: torch.Tensor | None, + generator: torch.Generator | list[torch.Generator] | None, output_type: str, output_uncertainty: bool, ) -> int: @@ -334,16 +334,16 @@ def progress_bar(self, iterable=None, total=None, desc=None, leave=True): def __call__( self, image: PipelineImageInput, - num_inference_steps: Optional[int] = None, + num_inference_steps: int | None = None, ensemble_size: int = 1, - processing_resolution: Optional[int] = None, + processing_resolution: int | None = None, match_input_resolution: bool = True, resample_method_input: str = "bilinear", resample_method_output: str = "bilinear", batch_size: int = 1, - ensembling_kwargs: Optional[Dict[str, Any]] = None, - latents: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + ensembling_kwargs: dict[str, Any] | None = None, + latents: torch.Tensor | list[torch.Tensor] | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, output_type: str = "np", output_uncertainty: bool = False, output_latent: bool = False, @@ -353,8 +353,8 @@ def __call__( Function invoked when calling the pipeline. Args: - image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`), - `List[torch.Tensor]`: An input image or images used as an input for the normals estimation task. For + image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`), + `list[torch.Tensor]`: An input image or images used as an input for the normals estimation task. For arrays and tensors, the expected value range is between `[0, 1]`. Passing a batch of images is possible by providing a four-dimensional array or a tensor. Additionally, a list of images of two- or three-dimensional arrays or tensors can be passed. In the latter case, all list elements must have the @@ -386,7 +386,7 @@ def __call__( latents (`torch.Tensor`, *optional*, defaults to `None`): Latent noise tensors to replace the random initialization. These can be taken from the previous function call's output. - generator (`torch.Generator`, or `List[torch.Generator]`, *optional*, defaults to `None`): + generator (`torch.Generator`, or `list[torch.Generator]`, *optional*, defaults to `None`): Random number generator object to ensure reproducibility. output_type (`str`, *optional*, defaults to `"np"`): Preferred format of the output's `prediction` and the optional `uncertainty` fields. The accepted @@ -595,11 +595,11 @@ def __call__( def prepare_latents( self, image: torch.Tensor, - latents: Optional[torch.Tensor], - generator: Optional[torch.Generator], + latents: torch.Tensor | None, + generator: torch.Generator | None, ensemble_size: int, batch_size: int, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: def retrieve_latents(encoder_output): if hasattr(encoder_output, "latent_dist"): return encoder_output.latent_dist.mode() @@ -660,7 +660,7 @@ def normalize_normals(normals: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: @staticmethod def ensemble_normals( normals: torch.Tensor, output_uncertainty: bool, reduction: str = "closest" - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + ) -> tuple[torch.Tensor, torch.Tensor | None]: """ Ensembles the normals maps represented by the `normals` tensor with expected shape `(B, 3, H, W)`, where B is the number of ensemble members for a given prediction of size `(H x W)`. diff --git a/src/diffusers/pipelines/mochi/pipeline_mochi.py b/src/diffusers/pipelines/mochi/pipeline_mochi.py index 5874a92c6f2f..e8acc0a75e4d 100644 --- a/src/diffusers/pipelines/mochi/pipeline_mochi.py +++ b/src/diffusers/pipelines/mochi/pipeline_mochi.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -78,10 +78,10 @@ def linear_quadratic_schedule(num_steps, threshold_noise, linear_steps=None): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -96,15 +96,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -196,11 +196,11 @@ def __init__( def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 256, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -253,25 +253,25 @@ def _get_t5_prompt_embeds( # Adapted from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, max_sequence_length: int = 256, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -498,33 +498,33 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, num_frames: int = 19, num_inference_steps: int = 64, - timesteps: List[int] = None, + timesteps: list[int] = None, guidance_scale: float = 4.5, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 256, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. height (`int`, *optional*, defaults to `self.default_height`): @@ -536,7 +536,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. @@ -548,7 +548,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of videos to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -579,7 +579,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -668,10 +668,14 @@ def __call__( sigmas = linear_quadratic_schedule(num_inference_steps, threshold_noise) sigmas = np.array(sigmas) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, - device, + timestep_device, timesteps, sigmas, ) diff --git a/src/diffusers/pipelines/mochi/pipeline_output.py b/src/diffusers/pipelines/mochi/pipeline_output.py index d15827bc0084..5068cf930aaa 100644 --- a/src/diffusers/pipelines/mochi/pipeline_output.py +++ b/src/diffusers/pipelines/mochi/pipeline_output.py @@ -11,8 +11,8 @@ class MochiPipelineOutput(BaseOutput): Output class for Mochi pipelines. Args: - frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): - List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): + list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape `(batch_size, num_frames, channels, height, width)`. """ diff --git a/src/diffusers/pipelines/musicldm/pipeline_musicldm.py b/src/diffusers/pipelines/musicldm/pipeline_musicldm.py index c909e5eb0d26..e7747a4f8c3d 100644 --- a/src/diffusers/pipelines/musicldm/pipeline_musicldm.py +++ b/src/diffusers/pipelines/musicldm/pipeline_musicldm.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -106,9 +106,9 @@ class MusicLDMPipeline(DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusi def __init__( self, vae: AutoencoderKL, - text_encoder: Union[ClapTextModelWithProjection, ClapModel], - tokenizer: Union[RobertaTokenizer, RobertaTokenizerFast], - feature_extractor: Optional[ClapFeatureExtractor], + text_encoder: ClapTextModelWithProjection | ClapModel, + tokenizer: RobertaTokenizer | RobertaTokenizerFast, + feature_extractor: ClapFeatureExtractor | None, unet: UNet2DConditionModel, scheduler: KarrasDiffusionSchedulers, vocoder: SpeechT5HifiGan, @@ -133,14 +133,14 @@ def _encode_prompt( num_waveforms_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device (`torch.device`): torch device @@ -148,7 +148,7 @@ def _encode_prompt( number of waveforms that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the audio generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -207,7 +207,7 @@ def _encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -434,28 +434,28 @@ def enable_model_cpu_offload(self, gpu_id=0): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - audio_length_in_s: Optional[float] = None, + prompt: str | list[str] = None, + audio_length_in_s: float | None = None, num_inference_steps: int = 200, guidance_scale: float = 2.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_waveforms_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_waveforms_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: Optional[int] = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - output_type: Optional[str] = "np", + callback: Callable[[int, int, torch.Tensor], None] | None = None, + callback_steps: int | None = 1, + cross_attention_kwargs: dict[str, Any] | None = None, + output_type: str | None = "np", ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide audio generation. If not defined, you need to pass `prompt_embeds`. audio_length_in_s (`int`, *optional*, defaults to 10.24): The length of the generated audio sample in seconds. @@ -465,7 +465,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 2.0): A higher guidance scale value encourages the model to generate audio that is closely linked to the text `prompt` at the expense of lower sound quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in audio generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_waveforms_per_prompt (`int`, *optional*, defaults to 1): @@ -477,7 +477,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/omnigen/pipeline_omnigen.py b/src/diffusers/pipelines/omnigen/pipeline_omnigen.py index 090cb46aace4..bdf4e30c6619 100644 --- a/src/diffusers/pipelines/omnigen/pipeline_omnigen.py +++ b/src/diffusers/pipelines/omnigen/pipeline_omnigen.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import numpy as np import torch @@ -59,10 +59,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -77,15 +77,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -170,9 +170,9 @@ def __init__( def encode_input_images( self, - input_pixel_values: List[torch.Tensor], - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + input_pixel_values: list[torch.Tensor], + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): """ get the continue embedding of input images by VAE @@ -331,32 +331,32 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]], - input_images: Union[PipelineImageInput, List[PipelineImageInput]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str], + input_images: PipelineImageInput | list[PipelineImageInput] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, max_input_image_size: int = 1024, - timesteps: List[int] = None, + timesteps: list[int] = None, guidance_scale: float = 2.5, img_guidance_scale: float = 1.6, use_input_image_size_as_output: bool = False, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If the input includes images, need to add placeholders `<|image_i|>` in the prompt to indicate the position of the i-th images. - input_images (`PipelineImageInput` or `List[PipelineImageInput]`, *optional*): + input_images (`PipelineImageInput` or `list[PipelineImageInput]`, *optional*): The list of input images. We will replace the "<|image_i|>" in prompt with the i-th image in list. height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): The height in pixels of the generated image. This is set to 1024 by default for the best results. @@ -367,7 +367,7 @@ def __call__( expense of slower inference. max_input_image_size (`int`, *optional*, defaults to 1024): the maximum size of input image, which will be used to crop the input image to the maximum size - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. @@ -384,7 +384,7 @@ def __call__( e.g., image editing task num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -401,7 +401,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -459,8 +459,12 @@ def __call__( # 5. Prepare timesteps sigmas = np.linspace(1, 0, num_inference_steps + 1)[:num_inference_steps] + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas=sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas=sigmas ) self._num_timesteps = len(timesteps) diff --git a/src/diffusers/pipelines/omnigen/processor_omnigen.py b/src/diffusers/pipelines/omnigen/processor_omnigen.py index 7ed11871bb2a..8a7f53a77136 100644 --- a/src/diffusers/pipelines/omnigen/processor_omnigen.py +++ b/src/diffusers/pipelines/omnigen/processor_omnigen.py @@ -13,7 +13,6 @@ # limitations under the License. import re -from typing import Dict, List import numpy as np import torch @@ -132,8 +131,8 @@ def add_prefix_instruction(self, prompt): def __call__( self, - instructions: List[str], - input_images: List[List[str]] = None, + instructions: list[str], + input_images: list[list[str]] = None, height: int = 1024, width: int = 1024, negative_prompt: str = "low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers.", @@ -141,7 +140,7 @@ def __call__( separate_cfg_input: bool = False, use_input_image_size_as_output: bool = False, num_images_per_prompt: int = 1, - ) -> Dict: + ) -> dict: if isinstance(instructions, str): instructions = [instructions] input_images = [input_images] diff --git a/src/diffusers/pipelines/onnx_utils.py b/src/diffusers/pipelines/onnx_utils.py index 74e9f0b97800..657750f90add 100644 --- a/src/diffusers/pipelines/onnx_utils.py +++ b/src/diffusers/pipelines/onnx_utils.py @@ -14,11 +14,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations import os import shutil from pathlib import Path -from typing import Optional, Union import numpy as np from huggingface_hub import hf_hub_download @@ -61,7 +61,7 @@ def __call__(self, **kwargs): return self.model.run(None, inputs) @staticmethod - def load_model(path: Union[str, Path], provider=None, sess_options=None, provider_options=None): + def load_model(path: str | Path, provider=None, sess_options=None, provider_options=None): """ Loads an ONNX Inference session with an ExecutionProvider. Default provider is `CPUExecutionProvider` @@ -84,7 +84,7 @@ def load_model(path: Union[str, Path], provider=None, sess_options=None, provide path, providers=[provider], sess_options=sess_options, provider_options=provider_options ) - def _save_pretrained(self, save_directory: Union[str, Path], file_name: Optional[str] = None, **kwargs): + def _save_pretrained(self, save_directory: str | Path, file_name: str | None = None, **kwargs): """ Save a model and its configuration file to a directory, so that it can be re-loaded using the [`~optimum.onnxruntime.modeling_ort.ORTModel.from_pretrained`] class method. It will always save the @@ -117,7 +117,7 @@ def _save_pretrained(self, save_directory: Union[str, Path], file_name: Optional def save_pretrained( self, - save_directory: Union[str, os.PathLike], + save_directory: str | os.PathLike, **kwargs, ): """ @@ -141,14 +141,14 @@ def save_pretrained( @validate_hf_hub_args def _from_pretrained( cls, - model_id: Union[str, Path], - token: Optional[Union[bool, str, None]] = None, - revision: Optional[Union[str, None]] = None, + model_id: str | Path, + token: bool | str | None | None = None, + revision: str | None | None = None, force_download: bool = False, - cache_dir: Optional[str] = None, - file_name: Optional[str] = None, - provider: Optional[str] = None, - sess_options: Optional["ort.SessionOptions"] = None, + cache_dir: str | None = None, + file_name: str | None = None, + provider: str | None = None, + sess_options: "ort.SessionOptions" | None = None, **kwargs, ): """ @@ -161,7 +161,7 @@ def _from_pretrained( Is needed to load models from a private or gated repository revision (`str`): Revision is the specific model version to use. It can be a branch name, a tag name, or a commit id - cache_dir (`Union[str, Path]`, *optional*): + cache_dir (`str | Path`, *optional*): Path to a directory in which a downloaded pretrained model configuration should be cached if the standard cache should not be used. force_download (`bool`, *optional*, defaults to `False`): @@ -210,10 +210,10 @@ def _from_pretrained( @validate_hf_hub_args def from_pretrained( cls, - model_id: Union[str, Path], + model_id: str | Path, force_download: bool = True, - token: Optional[str] = None, - cache_dir: Optional[str] = None, + token: str | None = None, + cache_dir: str | None = None, **model_kwargs, ): revision = None diff --git a/src/diffusers/pipelines/ovis_image/pipeline_output.py b/src/diffusers/pipelines/ovis_image/pipeline_output.py index 160c5b73a917..4daecf3b2197 100644 --- a/src/diffusers/pipelines/ovis_image/pipeline_output.py +++ b/src/diffusers/pipelines/ovis_image/pipeline_output.py @@ -13,7 +13,6 @@ # limitations under the License. from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -27,9 +26,9 @@ class OvisImagePipelineOutput(BaseOutput): Output class for Ovis-Image pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) + images (`list[PIL.Image.Image]` or `np.ndarray`) List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image, np.ndarray] diff --git a/src/diffusers/pipelines/ovis_image/pipeline_ovis_image.py b/src/diffusers/pipelines/ovis_image/pipeline_ovis_image.py index 94d6cee93d7e..c8ff8227f27e 100644 --- a/src/diffusers/pipelines/ovis_image/pipeline_ovis_image.py +++ b/src/diffusers/pipelines/ovis_image/pipeline_ovis_image.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -69,10 +69,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -87,15 +87,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -181,7 +181,7 @@ def __init__( def _get_messages( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, ): prompt = [prompt] if isinstance(prompt, str) else prompt messages = [] @@ -200,10 +200,10 @@ def _get_messages( def _get_ovis_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -239,10 +239,10 @@ def _get_ovis_prompt_embeds( def encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_embeds: torch.FloatTensor | None = None, ): r""" @@ -413,33 +413,33 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = "", + prompt: str | list[str] = None, + negative_prompt: str | list[str] = "", guidance_scale: float = 5.0, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + sigmas: list[float] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 256, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is not greater than `1`). @@ -453,13 +453,13 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): diff --git a/src/diffusers/pipelines/pag/pag_utils.py b/src/diffusers/pipelines/pag/pag_utils.py index 8a56961f321c..41395ece7421 100644 --- a/src/diffusers/pipelines/pag/pag_utils.py +++ b/src/diffusers/pipelines/pag/pag_utils.py @@ -13,7 +13,6 @@ # limitations under the License. import re -from typing import Dict, List, Tuple, Union import torch import torch.nn as nn @@ -112,7 +111,7 @@ def _apply_perturbed_attention_guidance( return_pred_text (bool): Whether to return the text noise prediction. Returns: - Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: The updated noise prediction tensor after applying + torch.Tensor | tuple[torch.Tensor, torch.Tensor]: The updated noise prediction tensor after applying perturbed attention guidance and the text noise prediction. """ pag_scale = self._get_pag_scale(t) @@ -151,8 +150,8 @@ def _prepare_perturbed_attention_guidance(self, cond, uncond, do_classifier_free def set_pag_applied_layers( self, - pag_applied_layers: Union[str, List[str]], - pag_attn_processors: Tuple[AttentionProcessor, AttentionProcessor] = ( + pag_applied_layers: str | list[str], + pag_attn_processors: tuple[AttentionProcessor, AttentionProcessor] = ( PAGCFGIdentitySelfAttnProcessor2_0(), PAGIdentitySelfAttnProcessor2_0(), ), @@ -161,7 +160,7 @@ def set_pag_applied_layers( Set the self-attention layers to apply PAG. Raise ValueError if the input is invalid. Args: - pag_applied_layers (`str` or `List[str]`): + pag_applied_layers (`str` or `list[str]`): One or more strings identifying the layer names, or a simple regex for matching multiple layers, where PAG is to be applied. A few ways of expected usage are as follows: - Single layers specified as - "blocks.{layer_index}" @@ -169,7 +168,7 @@ def set_pag_applied_layers( - Multiple layers as a block name - "mid" - Multiple layers as regex - "blocks.({layer_index_1}|{layer_index_2})" pag_attn_processors: - (`Tuple[AttentionProcessor, AttentionProcessor]`, defaults to `(PAGCFGIdentitySelfAttnProcessor2_0(), + (`tuple[AttentionProcessor, AttentionProcessor]`, defaults to `(PAGCFGIdentitySelfAttnProcessor2_0(), PAGIdentitySelfAttnProcessor2_0())`): A tuple of two attention processors. The first attention processor is for PAG with Classifier-free guidance enabled (conditional and unconditional). The second attention processor is for PAG with CFG disabled (unconditional only). @@ -214,7 +213,7 @@ def do_perturbed_attention_guidance(self) -> bool: return self._pag_scale > 0 and len(self.pag_applied_layers) > 0 @property - def pag_attn_processors(self) -> Dict[str, AttentionProcessor]: + def pag_attn_processors(self) -> dict[str, AttentionProcessor]: r""" Returns: `dict` of PAG attention processors: A dictionary contains all PAG attention processors used in the model diff --git a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd.py b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd.py index 1abef014301a..807c42d21bb4 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd.py @@ -14,7 +14,7 @@ import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -108,10 +108,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -126,15 +126,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -196,7 +196,7 @@ class StableDiffusionControlNetPAGPipeline( A `CLIPTokenizer` to tokenize text. unet ([`UNet2DConditionModel`]): A `UNet2DConditionModel` to denoise the encoded image latents. - controlnet ([`ControlNetModel`] or `List[ControlNetModel]`): + controlnet ([`ControlNetModel`] or `list[ControlNetModel]`): Provides additional conditioning to the `unet` during the denoising process. If you set multiple ControlNets as a list, the outputs from each ControlNet are added together to create one combined additional conditioning. @@ -222,13 +222,13 @@ def __init__( text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, unet: UNet2DConditionModel, - controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel], + controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel, scheduler: KarrasDiffusionSchedulers, safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, image_encoder: CLIPVisionModelWithProjection = None, requires_safety_checker: bool = True, - pag_applied_layers: Union[str, List[str]] = "mid", + pag_applied_layers: str | list[str] = "mid", ): super().__init__() @@ -279,16 +279,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -296,7 +296,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -395,7 +395,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -864,35 +864,33 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, + cross_attention_kwargs: dict[str, Any] | None = None, + controlnet_conditioning_scale: float | list[float] = 1.0, guess_mode: bool = False, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], pag_scale: float = 3.0, pag_adaptive_scale: float = 0.0, ): @@ -900,10 +898,10 @@ def __call__( The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: - `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,: + `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`): The ControlNet input condition to provide guidance to the `unet` for generation. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or @@ -919,18 +917,18 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -938,7 +936,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -952,7 +950,7 @@ def __call__( Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -965,16 +963,16 @@ def __call__( cross_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set the corresponding scale as a list. guess_mode (`bool`, *optional*, defaults to `False`): The ControlNet encoder tries to recognize the content of the input image even if you remove all prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended. - control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the ControlNet starts applying. - control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the ControlNet stops applying. clip_skip (`int`, *optional*): Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that @@ -984,7 +982,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1131,8 +1129,12 @@ def __call__( assert False # 5. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) self._num_timesteps = len(timesteps) diff --git a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py index 2781af789018..ebc2e882868c 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py @@ -15,7 +15,7 @@ # This model implementation is heavily inspired by https://github.com/haofanwang/ControlNet-for-Diffusers/ import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -119,7 +119,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -170,7 +170,7 @@ class StableDiffusionControlNetPAGInpaintPipeline( A `CLIPTokenizer` to tokenize text. unet ([`UNet2DConditionModel`]): A `UNet2DConditionModel` to denoise the encoded image latents. - controlnet ([`ControlNetModel`] or `List[ControlNetModel]`): + controlnet ([`ControlNetModel`] or `list[ControlNetModel]`): Provides additional conditioning to the `unet` during the denoising process. If you set multiple ControlNets as a list, the outputs from each ControlNet are added together to create one combined additional conditioning. @@ -196,13 +196,13 @@ def __init__( text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, unet: UNet2DConditionModel, - controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel], + controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel, scheduler: KarrasDiffusionSchedulers, safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, image_encoder: CLIPVisionModelWithProjection = None, requires_safety_checker: bool = True, - pag_applied_layers: Union[str, List[str]] = "mid", + pag_applied_layers: str | list[str] = "mid", ): super().__init__() @@ -255,16 +255,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -272,7 +272,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -371,7 +371,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -973,36 +973,34 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, image: PipelineImageInput = None, mask_image: PipelineImageInput = None, control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, - padding_mask_crop: Optional[int] = None, + height: int | None = None, + width: int | None = None, + padding_mask_crop: int | None = None, strength: float = 1.0, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 0.5, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + cross_attention_kwargs: dict[str, Any] | None = None, + controlnet_conditioning_scale: float | list[float] = 0.5, + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], pag_scale: float = 3.0, pag_adaptive_scale: float = 0.0, ): @@ -1010,25 +1008,25 @@ def __call__( The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, - `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, + `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, NumPy array or tensor representing an image batch to be used as the starting point. For both NumPy array and PyTorch tensor, the expected value range is between `[0, 1]`. If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a NumPy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)`. It can also accept image latents as `image`, but if passing latents directly it is not encoded again. - mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, - `List[PIL.Image.Image]`, or `List[np.ndarray]`): + mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, + `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, NumPy array or tensor representing an image batch to mask `image`. White pixels in the mask are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a single channel (luminance) before use. If it's a NumPy array or PyTorch tensor, it should contain one color channel (L) instead of 3, so the expected shape for PyTorch tensor would be `(B, 1, H, W)`, `(B, H, W)`, `(1, H, W)`, `(H, W)`. And for NumPy array, it would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W, 1)`, or `(H, W)`. - control_image (`torch.Tensor`, `PIL.Image.Image`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, - `List[List[torch.Tensor]]`, or `List[List[PIL.Image.Image]]`): + control_image (`torch.Tensor`, `PIL.Image.Image`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, + `list[list[torch.Tensor]]`, or `list[list[PIL.Image.Image]]`): The ControlNet input condition to provide guidance to the `unet` for generation. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or @@ -1058,7 +1056,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -1066,7 +1064,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -1080,7 +1078,7 @@ def __call__( Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -1093,13 +1091,13 @@ def __call__( cross_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 0.5): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 0.5): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set the corresponding scale as a list. - control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the ControlNet starts applying. - control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the ControlNet stops applying. clip_skip (`int`, *optional*): Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that @@ -1109,7 +1107,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl.py b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl.py index 381352ccc5d4..b0e2c03faed7 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl.py @@ -14,7 +14,7 @@ import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -125,10 +125,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -143,15 +143,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -218,7 +218,7 @@ class StableDiffusionXLControlNetPAGPipeline( A `CLIPTokenizer` to tokenize text. unet ([`UNet2DConditionModel`]): A `UNet2DConditionModel` to denoise the encoded image latents. - controlnet ([`ControlNetModel`] or `List[ControlNetModel]`): + controlnet ([`ControlNetModel`] or `list[ControlNetModel]`): Provides additional conditioning to the `unet` during the denoising process. If you set multiple ControlNets as a list, the outputs from each ControlNet are added together to create one combined additional conditioning. @@ -262,13 +262,13 @@ def __init__( tokenizer: CLIPTokenizer, tokenizer_2: CLIPTokenizer, unet: UNet2DConditionModel, - controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel], + controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel, scheduler: KarrasDiffusionSchedulers, force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, + add_watermarker: bool | None = None, feature_extractor: CLIPImageProcessor = None, image_encoder: CLIPVisionModelWithProjection = None, - pag_applied_layers: Union[str, List[str]] = "mid", # ["down.block_2", "up.block_1.attentions_0"], "mid" + pag_applied_layers: str | list[str] = "mid", # ["down.block_2", "up.block_1.attentions_0"], "mid" ): super().__init__() @@ -306,26 +306,26 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, + prompt_2: str | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders device: (`torch.device`): @@ -334,11 +334,11 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -456,7 +456,7 @@ def encode_prompt( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) - uncond_tokens: List[str] + uncond_tokens: list[str] if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -1002,45 +1002,43 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_end: Optional[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, + denoising_end: float | None = None, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, - original_size: Tuple[int, int] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Tuple[int, int] = None, - negative_original_size: Optional[Tuple[int, int]] = None, - negative_crops_coords_top_left: Tuple[int, int] = (0, 0), - negative_target_size: Optional[Tuple[int, int]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + cross_attention_kwargs: dict[str, Any] | None = None, + controlnet_conditioning_scale: float | list[float] = 1.0, + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, + original_size: tuple[int, int] = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] = None, + negative_original_size: tuple[int, int] | None = None, + negative_crops_coords_top_left: tuple[int, int] = (0, 0), + negative_target_size: tuple[int, int] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], pag_scale: float = 3.0, pag_adaptive_scale: float = 0.0, ): @@ -1048,13 +1046,13 @@ def __call__( The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: - `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,: + `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`): The ControlNet input condition to provide guidance to the `unet` for generation. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or @@ -1072,11 +1070,11 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -1090,10 +1088,10 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 5.0): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. This is sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders. num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -1101,7 +1099,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -1122,7 +1120,7 @@ def __call__( weighting). If not provided, pooled `negative_prompt_embeds` are generated from `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -1135,39 +1133,39 @@ def __call__( cross_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set the corresponding scale as a list. - control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the ControlNet starts applying. - control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the ControlNet stops applying. - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a specific image resolution. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a target image resolution. It should be as same as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more @@ -1180,7 +1178,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1329,8 +1327,12 @@ def __call__( assert False # 5. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) self._num_timesteps = len(timesteps) diff --git a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl_img2img.py b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl_img2img.py index df5b3f5c10a5..8967e50251b9 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl_img2img.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl_img2img.py @@ -14,7 +14,7 @@ import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -85,7 +85,6 @@ >>> from diffusers import ControlNetModel, StableDiffusionXLControlNetPAGImg2ImgPipeline, AutoencoderKL >>> from diffusers.utils import load_image - >>> depth_estimator = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to("cuda") >>> feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-hybrid-midas") >>> controlnet = ControlNetModel.from_pretrained( @@ -150,7 +149,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -203,7 +202,7 @@ class StableDiffusionXLControlNetPAGImg2ImgPipeline( Second Tokenizer of class [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents. - controlnet ([`ControlNetModel`] or `List[ControlNetModel]`): + controlnet ([`ControlNetModel`] or `list[ControlNetModel]`): Provides additional conditioning to the unet during the denoising process. If you set multiple ControlNets as a list, the outputs from each ControlNet are added together to create one combined additional conditioning. @@ -251,14 +250,14 @@ def __init__( tokenizer: CLIPTokenizer, tokenizer_2: CLIPTokenizer, unet: UNet2DConditionModel, - controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel], + controlnet: ControlNetModel | list[ControlNetModel] | tuple[ControlNetModel] | MultiControlNetModel, scheduler: KarrasDiffusionSchedulers, requires_aesthetics_score: bool = False, force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, + add_watermarker: bool | None = None, feature_extractor: CLIPImageProcessor = None, image_encoder: CLIPVisionModelWithProjection = None, - pag_applied_layers: Union[str, List[str]] = "mid", # ["mid"], ["down.block_1", "up.block_0.attentions_0"] + pag_applied_layers: str | list[str] = "mid", # ["mid"], ["down.block_1", "up.block_0.attentions_0"] ): super().__init__() @@ -298,26 +297,26 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, + prompt_2: str | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders device: (`torch.device`): @@ -326,11 +325,11 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -448,7 +447,7 @@ def encode_prompt( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) - uncond_tokens: List[str] + uncond_tokens: list[str] if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -1080,47 +1079,45 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, image: PipelineImageInput = None, control_image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, strength: float = 0.8, num_inference_steps: int = 50, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 0.8, + cross_attention_kwargs: dict[str, Any] | None = None, + controlnet_conditioning_scale: float | list[float] = 0.8, guess_mode: bool = False, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, - original_size: Tuple[int, int] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Tuple[int, int] = None, - negative_original_size: Optional[Tuple[int, int]] = None, - negative_crops_coords_top_left: Tuple[int, int] = (0, 0), - negative_target_size: Optional[Tuple[int, int]] = None, + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, + original_size: tuple[int, int] = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] = None, + negative_original_size: tuple[int, int] | None = None, + negative_crops_coords_top_left: tuple[int, int] = (0, 0), + negative_target_size: tuple[int, int] | None = None, aesthetic_score: float = 6.0, negative_aesthetic_score: float = 2.5, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], pag_scale: float = 3.0, pag_adaptive_scale: float = 0.0, ): @@ -1128,18 +1125,18 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: - `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,: + `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`): The initial image will be used as the starting point for the image generation process. Can also accept image latents as `image`, if passing latents directly, it will not be encoded again. - control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: - `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,: + `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`): The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height @@ -1169,11 +1166,11 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -1181,7 +1178,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -1203,7 +1200,7 @@ def __call__( weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -1218,42 +1215,42 @@ def __call__( A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original unet. If multiple ControlNets are specified in init, you can set the corresponding scale as a list. guess_mode (`bool`, *optional*, defaults to `False`): In this mode, the ControlNet encoder will try best to recognize the content of the input image even if you remove all prompts. The `guidance_scale` between 3.0 and 5.0 is recommended. - control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0): + control_guidance_start (`float` or `list[float]`, *optional*, defaults to 0.0): The percentage of total steps at which the controlnet starts applying. - control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0): + control_guidance_end (`float` or `list[float]`, *optional*, defaults to 1.0): The percentage of total steps at which the controlnet stops applying. - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a specific image resolution. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a target image resolution. It should be as same as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more @@ -1274,7 +1271,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/pag/pipeline_pag_hunyuandit.py b/src/diffusers/pipelines/pag/pipeline_pag_hunyuandit.py index 6704924b2512..15ac665acd2b 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_hunyuandit.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_hunyuandit.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Callable, Dict, List, Optional, Tuple, Union +from typing import Callable import numpy as np import torch @@ -164,10 +164,10 @@ class HunyuanDiTPAGPipeline(DiffusionPipeline, PAGMixin): vae ([`AutoencoderKL`]): Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. We use `sdxl-vae-fp16-fix`. - text_encoder (Optional[`~transformers.BertModel`, `~transformers.CLIPTextModel`]): + text_encoder (`~transformers.BertModel`, `~transformers.CLIPTextModel` | None): Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). HunyuanDiT uses a fine-tuned [bilingual CLIP]. - tokenizer (Optional[`~transformers.BertTokenizer`, `~transformers.CLIPTokenizer`]): + tokenizer (`~transformers.BertTokenizer`, `~transformers.CLIPTokenizer` | None): A `BertTokenizer` or `CLIPTokenizer` to tokenize text. transformer ([`HunyuanDiT2DModel`]): The HunyuanDiT model designed by Tencent Hunyuan. @@ -204,12 +204,12 @@ def __init__( tokenizer: BertTokenizer, transformer: HunyuanDiT2DModel, scheduler: DDPMScheduler, - safety_checker: Optional[StableDiffusionSafetyChecker] = None, - feature_extractor: Optional[CLIPImageProcessor] = None, + safety_checker: StableDiffusionSafetyChecker | None = None, + feature_extractor: CLIPImageProcessor | None = None, requires_safety_checker: bool = True, - text_encoder_2: Optional[T5EncoderModel] = None, - tokenizer_2: Optional[T5Tokenizer] = None, - pag_applied_layers: Union[str, List[str]] = "blocks.1", # "blocks.16.attn1", "blocks.16", "16", 16 + text_encoder_2: T5EncoderModel | None = None, + tokenizer_2: T5Tokenizer | None = None, + pag_applied_layers: str | list[str] = "blocks.1", # "blocks.16.attn1", "blocks.16", "16", 16 ): super().__init__() @@ -262,19 +262,19 @@ def encode_prompt( dtype: torch.dtype = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - max_sequence_length: Optional[int] = None, + negative_prompt: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + max_sequence_length: int | None = None, text_encoder_index: int = 0, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -284,7 +284,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -373,7 +373,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -580,34 +580,35 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: Optional[int] = 50, - guidance_scale: Optional[float] = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - eta: Optional[float] = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_2: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_2: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - prompt_attention_mask_2: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask_2: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, + num_inference_steps: int = 50, + guidance_scale: float = 5.0, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int = 1, + eta: float = 0.0, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_2: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_2: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + prompt_attention_mask_2: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask_2: torch.Tensor | None = None, + output_type: str = "pil", return_dict: bool = True, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int, dict], None] + | PipelineCallback + | MultiPipelineCallbacks + | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], guidance_rescale: float = 0.0, - original_size: Optional[Tuple[int, int]] = (1024, 1024), - target_size: Optional[Tuple[int, int]] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), + original_size: tuple[int, int] = (1024, 1024), + target_size: tuple[int, int] = None, + crops_coords_top_left: tuple[int, int] = (0, 0), use_resolution_binning: bool = True, pag_scale: float = 3.0, pag_adaptive_scale: float = 0.0, @@ -616,7 +617,7 @@ def __call__( The call function to the pipeline for generation with HunyuanDiT. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`): The height in pixels of the generated image. @@ -628,7 +629,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -636,7 +637,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. prompt_embeds (`torch.Tensor`, *optional*): @@ -664,19 +665,19 @@ def __call__( return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a plain tuple. - callback_on_step_end (`Callable[[int, int, Dict], None]`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): + callback_on_step_end (`Callable[[int, int], None]`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): A callback function or a list of callback functions to be called at the end of each denoising step. - callback_on_step_end_tensor_inputs (`List[str]`, *optional*): + callback_on_step_end_tensor_inputs (`list[str]`, *optional*): A list of tensor inputs that should be passed to the callback function. If not defined, all tensor inputs will be passed. guidance_rescale (`float`, *optional*, defaults to 0.0): Rescale the noise_cfg according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). See Section 3.4 - original_size (`Tuple[int, int]`, *optional*, defaults to `(1024, 1024)`): + original_size (`tuple[int, int]`, *optional*, defaults to `(1024, 1024)`): The original size of the image. Used to calculate the time ids. - target_size (`Tuple[int, int]`, *optional*): + target_size (`tuple[int, int]`, *optional*): The target size of the image. Used to calculate the time ids. - crops_coords_top_left (`Tuple[int, int]`, *optional*, defaults to `(0, 0)`): + crops_coords_top_left (`tuple[int, int]`, *optional*, defaults to `(0, 0)`): The top left coordinates of the crop. Used to calculate the time ids. use_resolution_binning (`bool`, *optional*, defaults to `True`): Whether to use resolution binning or not. If `True`, the input resolution will be mapped to the closest diff --git a/src/diffusers/pipelines/pag/pipeline_pag_kolors.py b/src/diffusers/pipelines/pag/pipeline_pag_kolors.py index 1403be03a620..4f138d91d9c6 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_kolors.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_kolors.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection @@ -68,10 +68,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -86,15 +86,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -154,7 +154,7 @@ class KolorsPAGPipeline( force_zeros_for_empty_prompt (`bool`, *optional*, defaults to `"False"`): Whether the negative prompt embeddings shall be forced to always be set to 0. Also see the config of `Kwai-Kolors/Kolors-diffusers`. - pag_applied_layers (`str` or `List[str]``, *optional*, defaults to `"mid"`): + pag_applied_layers (`str` or `list[str]``, *optional*, defaults to `"mid"`): Set the transformer attention layers where to apply the perturbed attention guidance. Can be a string or a list of strings with "down", "mid", "up", a whole transformer block or specific transformer block attention layers, e.g.: @@ -187,7 +187,7 @@ def __init__( image_encoder: CLIPVisionModelWithProjection = None, feature_extractor: CLIPImageProcessor = None, force_zeros_for_empty_prompt: bool = False, - pag_applied_layers: Union[str, List[str]] = "mid", + pag_applied_layers: str | list[str] = "mid", ): super().__init__() @@ -216,21 +216,21 @@ def __init__( def encode_prompt( self, prompt, - device: Optional[torch.device] = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, negative_prompt=None, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 256, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -238,7 +238,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -308,7 +308,7 @@ def encode_prompt( if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt: negative_prompt_embeds = torch.zeros_like(prompt_embeds) elif do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -666,38 +666,36 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_end: Optional[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, + denoising_end: float | None = None, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - original_size: Optional[Tuple[int, int]] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Optional[Tuple[int, int]] = None, - negative_original_size: Optional[Tuple[int, int]] = None, - negative_crops_coords_top_left: Tuple[int, int] = (0, 0), - negative_target_size: Optional[Tuple[int, int]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + cross_attention_kwargs: dict[str, Any] | None = None, + original_size: tuple[int, int] | None = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] | None = None, + negative_original_size: tuple[int, int] | None = None, + negative_crops_coords_top_left: tuple[int, int] = (0, 0), + negative_target_size: tuple[int, int] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], pag_scale: float = 3.0, pag_adaptive_scale: float = 0.0, max_sequence_length: int = 256, @@ -706,7 +704,7 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): @@ -722,11 +720,11 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -743,7 +741,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -752,7 +750,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -774,7 +772,7 @@ def __call__( weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -788,31 +786,31 @@ def __call__( A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a specific image resolution. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a target image resolution. It should be as same as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more @@ -822,7 +820,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -905,8 +903,12 @@ def __call__( ) # 4. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) # 5. Prepare latent variables diff --git a/src/diffusers/pipelines/pag/pipeline_pag_pixart_sigma.py b/src/diffusers/pipelines/pag/pipeline_pag_pixart_sigma.py index 9031877b5b8d..2e7e8ab9814f 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_pixart_sigma.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_pixart_sigma.py @@ -16,7 +16,7 @@ import inspect import re import urllib.parse as ul -from typing import Callable, List, Optional, Tuple, Union +from typing import Callable import torch from transformers import T5EncoderModel, T5Tokenizer @@ -84,10 +84,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -102,15 +102,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -173,7 +173,7 @@ def __init__( vae: AutoencoderKL, transformer: PixArtTransformer2DModel, scheduler: KarrasDiffusionSchedulers, - pag_applied_layers: Union[str, List[str]] = "blocks.1", # 1st transformer block + pag_applied_layers: str | list[str] = "blocks.1", # 1st transformer block ): super().__init__() @@ -189,15 +189,15 @@ def __init__( # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.encode_prompt with 120->300 def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], do_classifier_free_guidance: bool = True, negative_prompt: str = "", num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, + device: torch.device | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, clean_caption: bool = False, max_sequence_length: int = 300, **kwargs, @@ -206,9 +206,9 @@ def encode_prompt( Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For PixArt-Alpha, this should be "". @@ -575,51 +575,51 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, negative_prompt: str = "", num_inference_steps: int = 20, - timesteps: List[int] = None, - sigmas: List[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, guidance_scale: float = 4.5, - num_images_per_prompt: Optional[int] = 1, - height: Optional[int] = None, - width: Optional[int] = None, + num_images_per_prompt: int | None = 1, + height: int | None = None, + width: int | None = None, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, clean_caption: bool = True, use_resolution_binning: bool = True, max_sequence_length: int = 300, pag_scale: float = 3.0, pag_adaptive_scale: float = 0.0, - ) -> Union[ImagePipelineOutput, Tuple]: + ) -> ImagePipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_inference_steps (`int`, *optional*, defaults to 100): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -638,7 +638,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -764,8 +764,12 @@ def __call__( prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) # 4. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) # 5. Prepare latents. diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sana.py b/src/diffusers/pipelines/pag/pipeline_pag_sana.py index 9e91ccbe8006..71861f366477 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sana.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sana.py @@ -17,7 +17,7 @@ import re import urllib.parse as ul import warnings -from typing import Callable, Dict, List, Optional, Tuple, Union +from typing import Callable import torch from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast @@ -88,10 +88,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -106,15 +106,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -161,12 +161,12 @@ class SanaPAGPipeline(DiffusionPipeline, PAGMixin): def __init__( self, - tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast], + tokenizer: GemmaTokenizer | GemmaTokenizerFast, text_encoder: Gemma2PreTrainedModel, vae: AutoencoderDC, transformer: SanaTransformer2DModel, scheduler: FlowMatchEulerDiscreteScheduler, - pag_applied_layers: Union[str, List[str]] = "transformer_blocks.0", + pag_applied_layers: str | list[str] = "transformer_blocks.0", ): super().__init__() @@ -241,26 +241,26 @@ def disable_vae_tiling(self): def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], do_classifier_free_guidance: bool = True, negative_prompt: str = "", num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, + device: torch.device | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, clean_caption: bool = False, max_sequence_length: int = 300, - complex_human_instruction: Optional[List[str]] = None, + complex_human_instruction: list[str] | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For PixArt-Alpha, this should be "". @@ -649,30 +649,30 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, negative_prompt: str = "", num_inference_steps: int = 20, - timesteps: List[int] = None, - sigmas: List[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, guidance_scale: float = 4.5, - num_images_per_prompt: Optional[int] = 1, + num_images_per_prompt: int | None = 1, height: int = 1024, width: int = 1024, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, clean_caption: bool = False, use_resolution_binning: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 300, - complex_human_instruction: List[str] = [ + complex_human_instruction: list[str] = [ "Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:", "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.", "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.", @@ -684,26 +684,26 @@ def __call__( ], pag_scale: float = 3.0, pag_adaptive_scale: float = 0.0, - ) -> Union[ImagePipelineOutput, Tuple]: + ) -> ImagePipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_inference_steps (`int`, *optional*, defaults to 20): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -722,7 +722,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -756,12 +756,12 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. max_sequence_length (`int` defaults to 300): Maximum sequence length to use with the `prompt`. - complex_human_instruction (`List[str]`, *optional*): + complex_human_instruction (`list[str]`, *optional*): Instructions for complex human attention: https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55. pag_scale (`float`, *optional*, defaults to 3.0): @@ -856,8 +856,12 @@ def __call__( prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) # 4. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) # 5. Prepare latents. diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd.py b/src/diffusers/pipelines/pag/pipeline_pag_sd.py index ea64f8be2c50..26ea717556c5 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch from packaging import version @@ -97,10 +97,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -115,15 +115,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -212,7 +212,7 @@ def __init__( feature_extractor: CLIPImageProcessor, image_encoder: CLIPVisionModelWithProjection = None, requires_safety_checker: bool = True, - pag_applied_layers: Union[str, List[str]] = "mid", + pag_applied_layers: str | list[str] = "mid", ): super().__init__() @@ -308,16 +308,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -325,7 +325,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -424,7 +424,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -746,29 +746,29 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, guidance_rescale: float = 0.0, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], pag_scale: float = 3.0, pag_adaptive_scale: float = 0.0, ): @@ -776,7 +776,7 @@ def __call__( The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -785,18 +785,18 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -804,7 +804,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -818,7 +818,7 @@ def __call__( Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -843,7 +843,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -952,8 +952,12 @@ def __call__( ip_adapter_image_embeds[i] = image_embeds # 4. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) # 5. Prepare latent variables diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py index 941b675099b9..f0fbef29b699 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_3.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch from transformers import ( @@ -76,10 +76,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -94,15 +94,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -185,7 +185,7 @@ def __init__( tokenizer_2: CLIPTokenizer, text_encoder_3: T5EncoderModel, tokenizer_3: T5TokenizerFast, - pag_applied_layers: Union[str, List[str]] = "blocks.1", # 1st transformer block + pag_applied_layers: str | list[str] = "blocks.1", # 1st transformer block ): super().__init__() @@ -221,11 +221,11 @@ def __init__( # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 256, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -278,10 +278,10 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - clip_skip: Optional[int] = None, + device: torch.device | None = None, + clip_skip: int | None = None, clip_model_index: int = 0, ): device = device or self._execution_device @@ -334,32 +334,32 @@ def _get_clip_prompt_embeds( # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], - prompt_3: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + prompt_2: str | list[str], + prompt_3: str | list[str], + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt_3: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_3: str | list[str] | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + clip_skip: int | None = None, max_sequence_length: int = 256, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in all text-encoders - prompt_3 (`str` or `List[str]`, *optional*): + prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is used in all text-encoders device: (`torch.device`): @@ -368,14 +368,14 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders. - negative_prompt_3 (`str` or `List[str]`, *optional*): + negative_prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `negative_prompt` is used in all the text-encoders. prompt_embeds (`torch.FloatTensor`, *optional*): @@ -685,30 +685,30 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - prompt_3: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + prompt_3: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 28, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 7.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt_3: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_3: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 256, pag_scale: float = 3.0, pag_adaptive_scale: float = 0.0, @@ -717,13 +717,13 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead - prompt_3 (`str` or `List[str]`, *optional*): + prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is will be used instead height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): @@ -733,7 +733,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -743,19 +743,19 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used instead - negative_prompt_3 (`str` or `List[str]`, *optional*): + negative_prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `negative_prompt` is used instead num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -791,7 +791,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -888,7 +888,13 @@ def __call__( pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0) # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, sigmas=sigmas) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, sigmas=sigmas + ) num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) self._num_timesteps = len(timesteps) diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py index f40dd52fc244..84b727dc0613 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import PIL.Image import torch @@ -77,7 +77,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -92,10 +92,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -110,15 +110,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -201,7 +201,7 @@ def __init__( tokenizer_2: CLIPTokenizer, text_encoder_3: T5EncoderModel, tokenizer_3: T5TokenizerFast, - pag_applied_layers: Union[str, List[str]] = "blocks.1", # 1st transformer block + pag_applied_layers: str | list[str] = "blocks.1", # 1st transformer block ): super().__init__() @@ -237,11 +237,11 @@ def __init__( # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 256, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -294,10 +294,10 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - clip_skip: Optional[int] = None, + device: torch.device | None = None, + clip_skip: int | None = None, clip_model_index: int = 0, ): device = device or self._execution_device @@ -350,32 +350,32 @@ def _get_clip_prompt_embeds( # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], - prompt_3: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + prompt_2: str | list[str], + prompt_3: str | list[str], + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt_3: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_3: str | list[str] | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + clip_skip: int | None = None, max_sequence_length: int = 256, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in all text-encoders - prompt_3 (`str` or `List[str]`, *optional*): + prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is used in all text-encoders device: (`torch.device`): @@ -384,14 +384,14 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders. - negative_prompt_3 (`str` or `List[str]`, *optional*): + negative_prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `negative_prompt` is used in all the text-encoders. prompt_embeds (`torch.FloatTensor`, *optional*): @@ -736,32 +736,32 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - prompt_3: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + prompt_3: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, image: PipelineImageInput = None, strength: float = 0.6, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 7.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt_3: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_3: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 256, pag_scale: float = 3.0, pag_adaptive_scale: float = 0.0, @@ -770,16 +770,16 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead - prompt_3 (`str` or `List[str]`, *optional*): + prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is will be used instead - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a @@ -794,7 +794,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -804,19 +804,19 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used instead - negative_prompt_3 (`str` or `List[str]`, *optional*): + negative_prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `negative_prompt` is used instead num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -852,7 +852,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -951,7 +951,13 @@ def __call__( image = self.image_processor.preprocess(image, height=height, width=width) # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, sigmas=sigmas) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, sigmas=sigmas + ) timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) # 5. Prepare latent variables diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_animatediff.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_animatediff.py index de13be9c4d22..62d1c912283f 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_animatediff.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_animatediff.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection @@ -135,12 +135,12 @@ def __init__( vae: AutoencoderKL, text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, - unet: Union[UNet2DConditionModel, UNetMotionModel], + unet: UNet2DConditionModel | UNetMotionModel, motion_adapter: MotionAdapter, scheduler: KarrasDiffusionSchedulers, feature_extractor: CLIPImageProcessor = None, image_encoder: CLIPVisionModelWithProjection = None, - pag_applied_layers: Union[str, List[str]] = "mid_block.*attn1", # ["mid"], ["down_blocks.1"] + pag_applied_layers: str | list[str] = "mid_block.*attn1", # ["mid"], ["down_blocks.1"] ): super().__init__() if isinstance(unet, UNet2DConditionModel): @@ -169,16 +169,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -186,7 +186,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -285,7 +285,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -576,27 +576,27 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Optional[Union[str, List[str]]] = None, - num_frames: Optional[int] = 16, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] | None = None, + num_frames: int | None = 16, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_videos_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_videos_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + cross_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], decode_chunk_size: int = 16, pag_scale: float = 3.0, pag_adaptive_scale: float = 0.0, @@ -605,7 +605,7 @@ def __call__( The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated video. @@ -620,13 +620,13 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -642,7 +642,7 @@ def __call__( not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -663,7 +663,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_img2img.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_img2img.py index 8351112ce409..822483eca995 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_img2img.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_img2img.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import PIL.Image import torch @@ -77,7 +77,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -92,10 +92,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -110,15 +110,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -207,7 +207,7 @@ def __init__( feature_extractor: CLIPImageProcessor, image_encoder: CLIPVisionModelWithProjection = None, requires_safety_checker: bool = True, - pag_applied_layers: Union[str, List[str]] = "mid", # ["mid"], ["down.block_1", "up.block_0.attentions_0"] + pag_applied_layers: str | list[str] = "mid", # ["mid"], ["down.block_1", "up.block_0.attentions_0"] ): super().__init__() @@ -303,16 +303,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -320,7 +320,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -419,7 +419,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -783,29 +783,27 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, image: PipelineImageInput = None, strength: float = 0.8, - num_inference_steps: Optional[int] = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - guidance_scale: Optional[float] = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - eta: Optional[float] = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + num_inference_steps: int | None = 50, + timesteps: list[int] = None, + sigmas: list[float] = None, + guidance_scale: float | None = 7.5, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + eta: float | None = 0.0, + generator: torch.Generator | list[torch.Generator] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, clip_skip: int = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], pag_scale: float = 3.0, pag_adaptive_scale: float = 0.0, ): @@ -813,9 +811,9 @@ def __call__( The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a @@ -830,18 +828,18 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. This parameter is modulated by `strength`. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -849,7 +847,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. prompt_embeds (`torch.Tensor`, *optional*): @@ -859,7 +857,7 @@ def __call__( Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -880,7 +878,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -986,8 +984,12 @@ def __call__( image = self.image_processor.preprocess(image) # 5. set timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_inpaint.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_inpaint.py index 6b1b294e10f5..0f6fbbd9ae16 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_inpaint.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_inpaint.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import PIL.Image import torch @@ -82,7 +82,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -124,10 +124,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -142,15 +142,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -239,7 +239,7 @@ def __init__( feature_extractor: CLIPImageProcessor, image_encoder: CLIPVisionModelWithProjection = None, requires_safety_checker: bool = True, - pag_applied_layers: Union[str, List[str]] = "mid", + pag_applied_layers: str | list[str] = "mid", ): super().__init__() @@ -338,16 +338,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -355,7 +355,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -454,7 +454,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -911,34 +911,34 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, image: PipelineImageInput = None, mask_image: PipelineImageInput = None, masked_image_latents: torch.Tensor = None, - height: Optional[int] = None, - width: Optional[int] = None, - padding_mask_crop: Optional[int] = None, + height: int | None = None, + width: int | None = None, + padding_mask_crop: int | None = None, strength: float = 0.9999, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, guidance_rescale: float = 0.0, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], pag_scale: float = 3.0, pag_adaptive_scale: float = 0.0, ): @@ -946,7 +946,7 @@ def __call__( The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -955,18 +955,18 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -974,7 +974,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -988,7 +988,7 @@ def __call__( Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -1013,7 +1013,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1094,8 +1094,12 @@ def __call__( ) # 4. set timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) timesteps, num_inference_steps = self.get_timesteps( num_inference_steps=num_inference_steps, strength=strength, device=device diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_xl.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_xl.py index a69f06536a55..2987c90626ef 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_xl.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_xl.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from transformers import ( @@ -112,10 +112,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -130,15 +130,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -254,8 +254,8 @@ def __init__( image_encoder: CLIPVisionModelWithProjection = None, feature_extractor: CLIPImageProcessor = None, force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, - pag_applied_layers: Union[str, List[str]] = "mid", # ["mid"],["down.block_1"],["up.block_0.attentions_0"] + add_watermarker: bool | None = None, + pag_applied_layers: str | list[str] = "mid", # ["mid"],["down.block_1"],["up.block_0.attentions_0"] ): super().__init__() @@ -293,26 +293,26 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, + prompt_2: str | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders device: (`torch.device`): @@ -321,11 +321,11 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -443,7 +443,7 @@ def encode_prompt( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) - uncond_tokens: List[str] + uncond_tokens: list[str] if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -835,40 +835,40 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_end: Optional[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, + denoising_end: float | None = None, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, guidance_rescale: float = 0.0, - original_size: Optional[Tuple[int, int]] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Optional[Tuple[int, int]] = None, - negative_original_size: Optional[Tuple[int, int]] = None, - negative_crops_coords_top_left: Tuple[int, int] = (0, 0), - negative_target_size: Optional[Tuple[int, int]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + original_size: tuple[int, int] | None = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] | None = None, + negative_original_size: tuple[int, int] | None = None, + negative_crops_coords_top_left: tuple[int, int] = (0, 0), + negative_target_size: tuple[int, int] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], pag_scale: float = 3.0, pag_adaptive_scale: float = 0.0, ): @@ -876,10 +876,10 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): @@ -895,11 +895,11 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -916,11 +916,11 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -928,7 +928,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -950,7 +950,7 @@ def __call__( weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -971,31 +971,31 @@ def __call__( [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when using zero terminal SNR. - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a specific image resolution. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a target image resolution. It should be as same as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more @@ -1005,7 +1005,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1095,8 +1095,12 @@ def __call__( ) # 4. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) # 5. Prepare latent variables diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_img2img.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_img2img.py index 416d9e5677b4..433b9edc69b7 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_img2img.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_img2img.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import PIL.Image import torch @@ -116,7 +116,7 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -131,10 +131,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -149,15 +149,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -277,8 +277,8 @@ def __init__( feature_extractor: CLIPImageProcessor = None, requires_aesthetics_score: bool = False, force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, - pag_applied_layers: Union[str, List[str]] = "mid", # ["mid"], ["down.block_1", "up.block_0.attentions_0"] + add_watermarker: bool | None = None, + pag_applied_layers: str | list[str] = "mid", # ["mid"], ["down.block_1", "up.block_0.attentions_0"] ): super().__init__() @@ -311,26 +311,26 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, + prompt_2: str | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders device: (`torch.device`): @@ -339,11 +339,11 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -461,7 +461,7 @@ def encode_prompt( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) - uncond_tokens: List[str] + uncond_tokens: list[str] if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -988,45 +988,43 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, image: PipelineImageInput = None, strength: float = 0.3, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_start: Optional[float] = None, - denoising_end: Optional[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, + denoising_start: float | None = None, + denoising_end: float | None = None, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, guidance_rescale: float = 0.0, - original_size: Tuple[int, int] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Tuple[int, int] = None, - negative_original_size: Optional[Tuple[int, int]] = None, - negative_crops_coords_top_left: Tuple[int, int] = (0, 0), - negative_target_size: Optional[Tuple[int, int]] = None, + original_size: tuple[int, int] = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] = None, + negative_original_size: tuple[int, int] | None = None, + negative_crops_coords_top_left: tuple[int, int] = (0, 0), + negative_target_size: tuple[int, int] | None = None, aesthetic_score: float = 6.0, negative_aesthetic_score: float = 2.5, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], pag_scale: float = 3.0, pag_adaptive_scale: float = 0.0, ): @@ -1034,13 +1032,13 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders - image (`torch.Tensor` or `PIL.Image.Image` or `np.ndarray` or `List[torch.Tensor]` or `List[PIL.Image.Image]` or `List[np.ndarray]`): + image (`torch.Tensor` or `PIL.Image.Image` or `np.ndarray` or `list[torch.Tensor]` or `list[PIL.Image.Image]` or `list[np.ndarray]`): The image(s) to modify with the pipeline. strength (`float`, *optional*, defaults to 0.3): Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1. `image` @@ -1052,11 +1050,11 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -1081,11 +1079,11 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -1093,7 +1091,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -1115,7 +1113,7 @@ def __call__( weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -1136,31 +1134,31 @@ def __call__( [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when using zero terminal SNR. - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a specific image resolution. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a target image resolution. It should be as same as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more @@ -1181,7 +1179,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1272,8 +1270,12 @@ def __call__( def denoising_value_valid(dnv): return isinstance(dnv, float) and 0 < dnv < 1 + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) timesteps, num_inference_steps = self.get_timesteps( num_inference_steps, diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py b/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py index 6be341e07b1a..9caf50e5e333 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import PIL.Image import torch @@ -129,7 +129,7 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -144,10 +144,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -162,15 +162,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -293,8 +293,8 @@ def __init__( feature_extractor: CLIPImageProcessor = None, requires_aesthetics_score: bool = False, force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, - pag_applied_layers: Union[str, List[str]] = "mid", # ["mid"], ["down.block_1", "up.block_0.attentions_0"] + add_watermarker: bool | None = None, + pag_applied_layers: str | list[str] = "mid", # ["mid"], ["down.block_1", "up.block_0.attentions_0"] ): super().__init__() @@ -401,26 +401,26 @@ def prepare_ip_adapter_image_embeds( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, + prompt_2: str | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders device: (`torch.device`): @@ -429,11 +429,11 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -551,7 +551,7 @@ def encode_prompt( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) - uncond_tokens: List[str] + uncond_tokens: list[str] if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -1079,50 +1079,48 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, image: PipelineImageInput = None, mask_image: PipelineImageInput = None, masked_image_latents: torch.Tensor = None, - height: Optional[int] = None, - width: Optional[int] = None, - padding_mask_crop: Optional[int] = None, + height: int | None = None, + width: int | None = None, + padding_mask_crop: int | None = None, strength: float = 0.9999, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_start: Optional[float] = None, - denoising_end: Optional[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, + denoising_start: float | None = None, + denoising_end: float | None = None, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, guidance_rescale: float = 0.0, - original_size: Tuple[int, int] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Tuple[int, int] = None, - negative_original_size: Optional[Tuple[int, int]] = None, - negative_crops_coords_top_left: Tuple[int, int] = (0, 0), - negative_target_size: Optional[Tuple[int, int]] = None, + original_size: tuple[int, int] = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] = None, + negative_original_size: tuple[int, int] | None = None, + negative_crops_coords_top_left: tuple[int, int] = (0, 0), + negative_target_size: tuple[int, int] | None = None, aesthetic_score: float = 6.0, negative_aesthetic_score: float = 2.5, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], pag_scale: float = 3.0, pag_adaptive_scale: float = 0.0, ): @@ -1130,10 +1128,10 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders image (`PIL.Image.Image`): @@ -1172,11 +1170,11 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -1201,11 +1199,11 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -1223,7 +1221,7 @@ def __call__( weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -1250,31 +1248,31 @@ def __call__( A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a specific image resolution. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a target image resolution. It should be as same as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more @@ -1295,7 +1293,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1392,8 +1390,12 @@ def __call__( def denoising_value_valid(dnv): return isinstance(dnv, float) and 0 < dnv < 1 + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) timesteps, num_inference_steps = self.get_timesteps( num_inference_steps, diff --git a/src/diffusers/pipelines/paint_by_example/__init__.py b/src/diffusers/pipelines/paint_by_example/__init__.py index aaa775f690c3..d2906b540c6e 100644 --- a/src/diffusers/pipelines/paint_by_example/__init__.py +++ b/src/diffusers/pipelines/paint_by_example/__init__.py @@ -1,5 +1,5 @@ from dataclasses import dataclass -from typing import TYPE_CHECKING, List, Optional, Union +from typing import TYPE_CHECKING import numpy as np import PIL diff --git a/src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py b/src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py index c09992befbcb..aa7dbaa720e5 100644 --- a/src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py +++ b/src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Callable, List, Optional, Union +from typing import Callable import numpy as np import PIL.Image @@ -43,7 +43,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -65,7 +65,7 @@ def prepare_mask_and_masked_image(image, mask): binarized (``mask > 0.5``) and cast to ``torch.float32`` too. Args: - image (Union[np.array, PIL.Image, torch.Tensor]): The image to inpaint. + image (np.array | PIL.Image | torch.Tensor): The image to inpaint. It can be a ``PIL.Image``, or a ``height x width x 3`` ``np.array`` or a ``channels x height x width`` ``torch.Tensor`` or a ``batch x channels x height x width`` ``torch.Tensor``. mask (_type_): The mask to apply to the image, i.e. regions to inpaint. @@ -198,7 +198,7 @@ def __init__( vae: AutoencoderKL, image_encoder: PaintByExampleImageEncoder, unet: UNet2DConditionModel, - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], + scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler, safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, requires_safety_checker: bool = False, @@ -270,7 +270,7 @@ def check_inputs(self, image, height, width, callback_steps): and not isinstance(image, list) ): raise ValueError( - "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is" + "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]` but is" f" {type(image)}" ) @@ -397,33 +397,33 @@ def _encode_image(self, image, device, num_images_per_prompt, do_classifier_free @torch.no_grad() def __call__( self, - example_image: Union[torch.Tensor, PIL.Image.Image], - image: Union[torch.Tensor, PIL.Image.Image], - mask_image: Union[torch.Tensor, PIL.Image.Image], - height: Optional[int] = None, - width: Optional[int] = None, + example_image: torch.Tensor | PIL.Image.Image, + image: torch.Tensor | PIL.Image.Image, + mask_image: torch.Tensor | PIL.Image.Image, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, ): r""" The call function to the pipeline for generation. Args: - example_image (`torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]`): + example_image (`torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]`): An example image to guide image generation. - image (`torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]`): + image (`torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]`): `Image` or tensor representing an image batch to be inpainted (parts of the image are masked out with `mask_image` and repainted according to `prompt`). - mask_image (`torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]`): + mask_image (`torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]`): `Image` or tensor representing an image batch to mask `image`. White pixels in the mask are repainted, while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a single channel (luminance) before use. If it's a tensor, it should contain one color channel (L) instead of 3, so the @@ -438,7 +438,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -446,7 +446,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/pia/pipeline_pia.py b/src/diffusers/pipelines/pia/pipeline_pia.py index dfc6e83fbd7c..d108deb9c5da 100644 --- a/src/diffusers/pipelines/pia/pipeline_pia.py +++ b/src/diffusers/pipelines/pia/pipeline_pia.py @@ -14,7 +14,7 @@ import inspect from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL @@ -85,7 +85,7 @@ ``` """ -RANGE_LIST = [ +RANGE_list = [ [1.0, 0.9, 0.85, 0.85, 0.85, 0.8], # 0 Small Motion [1.0, 0.8, 0.8, 0.8, 0.79, 0.78, 0.75], # Moderate Motion [1.0, 0.8, 0.7, 0.7, 0.7, 0.7, 0.7, 0.7, 0.7, 0.7, 0.6, 0.5, 0.5], # Large Motion @@ -103,7 +103,7 @@ def prepare_mask_coef_by_statistics(num_frames: int, cond_frame: int, motion_sca assert num_frames > cond_frame, "video_length should be greater than cond_frame" - range_list = RANGE_LIST + range_list = RANGE_list assert motion_scale < len(range_list), f"motion_scale type{motion_scale} not implemented" @@ -122,13 +122,13 @@ class PIAPipelineOutput(BaseOutput): Output class for PIAPipeline. Args: - frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): Nested list of length `batch_size` with denoised PIL image sequences of length `num_frames`, NumPy array of shape `(batch_size, num_frames, channels, height, width, Torch tensor of shape `(batch_size, num_frames, channels, height, width)`. """ - frames: Union[torch.Tensor, np.ndarray, List[List[PIL.Image.Image]]] + frames: torch.Tensor | np.ndarray | list[list[PIL.Image.Image]] class PIAPipeline( @@ -179,16 +179,14 @@ def __init__( vae: AutoencoderKL, text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, - unet: Union[UNet2DConditionModel, UNetMotionModel], - scheduler: Union[ - DDIMScheduler, - PNDMScheduler, - LMSDiscreteScheduler, - EulerDiscreteScheduler, - EulerAncestralDiscreteScheduler, - DPMSolverMultistepScheduler, - ], - motion_adapter: Optional[MotionAdapter] = None, + unet: UNet2DConditionModel | UNetMotionModel, + scheduler: DDIMScheduler + | PNDMScheduler + | LMSDiscreteScheduler + | EulerDiscreteScheduler + | EulerAncestralDiscreteScheduler + | DPMSolverMultistepScheduler, + motion_adapter: MotionAdapter | None = None, feature_extractor: CLIPImageProcessor = None, image_encoder: CLIPVisionModelWithProjection = None, ): @@ -217,16 +215,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -234,7 +232,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -333,7 +331,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -674,29 +672,29 @@ def num_timesteps(self): def __call__( self, image: PipelineImageInput, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, strength: float = 1.0, - num_frames: Optional[int] = 16, - height: Optional[int] = None, - width: Optional[int] = None, + num_frames: int | None = 16, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_videos_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_videos_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, motion_scale: int = 0, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + cross_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], ): r""" The call function to the pipeline for generation. @@ -704,7 +702,7 @@ def __call__( Args: image (`PipelineImageInput`): The input image to be used for video generation. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. strength (`float`, *optional*, defaults to 1.0): Indicates extent to transform the reference `image`. Must be between 0 and 1. @@ -721,13 +719,13 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -743,7 +741,7 @@ def __call__( not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -769,7 +767,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/pipeline_flax_utils.py b/src/diffusers/pipelines/pipeline_flax_utils.py index 2724c764c771..51dcf9a2fecf 100644 --- a/src/diffusers/pipelines/pipeline_flax_utils.py +++ b/src/diffusers/pipelines/pipeline_flax_utils.py @@ -17,7 +17,7 @@ import importlib import inspect import os -from typing import Any, Dict, List, Optional, Union +from typing import Any import flax import numpy as np @@ -90,12 +90,12 @@ class FlaxImagePipelineOutput(BaseOutput): Output class for image pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, num_channels)`. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image] | np.ndarray class FlaxDiffusionPipeline(ConfigMixin, PushToHubMixin): @@ -150,8 +150,8 @@ def register_modules(self, **kwargs): def save_pretrained( self, - save_directory: Union[str, os.PathLike], - params: Union[Dict, FrozenDict], + save_directory: str | os.PathLike, + params: dict | FrozenDict, push_to_hub: bool = False, **kwargs, ): @@ -168,7 +168,7 @@ class implements both a save and loading method. The pipeline is easily reloaded Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the repository you want to push to with `repo_id` (will default to the name of `save_directory` in your namespace). - kwargs (`Dict[str, Any]`, *optional*): + kwargs (`dict[str, Any]`, *optional*): Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. """ self.save_config(save_directory) @@ -228,7 +228,7 @@ class implements both a save and loading method. The pipeline is easily reloaded @classmethod @validate_hf_hub_args - def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs): + def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None, **kwargs): r""" Instantiate a Flax-based diffusion pipeline from pretrained pipeline weights. @@ -254,7 +254,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. output_loading_info(`bool`, *optional*, defaults to `False`): @@ -544,7 +544,7 @@ def _get_signature_keys(cls, obj): return expected_modules, optional_parameters @property - def components(self) -> Dict[str, Any]: + def components(self) -> dict[str, Any]: r""" The `self.components` property can be useful to run different pipelines with the same weights and diff --git a/src/diffusers/pipelines/pipeline_loading_utils.py b/src/diffusers/pipelines/pipeline_loading_utils.py index 8868e942ce3d..b2564d25505e 100644 --- a/src/diffusers/pipelines/pipeline_loading_utils.py +++ b/src/diffusers/pipelines/pipeline_loading_utils.py @@ -17,7 +17,7 @@ import re import warnings from pathlib import Path -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import httpx import requests @@ -210,7 +210,7 @@ def filter_with_regex(filenames, pattern_re): return {f for f in filenames if pattern_re.match(f.split("/")[-1]) is not None} -def variant_compatible_siblings(filenames, variant=None, ignore_patterns=None) -> Union[List[os.PathLike], str]: +def variant_compatible_siblings(filenames, variant=None, ignore_patterns=None) -> list[os.PathLike] | str: weight_names = [ WEIGHTS_NAME, SAFETENSORS_WEIGHTS_NAME, @@ -525,12 +525,12 @@ def _get_pipeline_class( def _load_empty_model( library_name: str, class_name: str, - importable_classes: List[Any], + importable_classes: list[Any], pipelines: Any, is_pipeline_module: bool, name: str, - torch_dtype: Union[str, torch.dtype], - cached_folder: Union[str, os.PathLike], + torch_dtype: str | torch.dtype, + cached_folder: str | os.PathLike, **kwargs, ): # retrieve class objects. @@ -607,7 +607,7 @@ def _load_empty_model( def _assign_components_to_devices( - module_sizes: Dict[str, float], device_memory: Dict[str, float], device_mapping_strategy: str = "balanced" + module_sizes: dict[str, float], device_memory: dict[str, float], device_mapping_strategy: str = "balanced" ): device_ids = list(device_memory.keys()) device_cycle = device_ids + device_ids[::-1] @@ -738,27 +738,28 @@ def _get_final_device_map(device_map, pipeline_class, passed_class_obj, init_dic def load_sub_model( library_name: str, class_name: str, - importable_classes: List[Any], + importable_classes: list[Any], pipelines: Any, is_pipeline_module: bool, pipeline_class: Any, torch_dtype: torch.dtype, provider: Any, sess_options: Any, - device_map: Optional[Union[Dict[str, torch.device], str]], - max_memory: Optional[Dict[Union[int, str], Union[int, str]]], - offload_folder: Optional[Union[str, os.PathLike]], + device_map: dict[str, torch.device] | str | None, + max_memory: dict[int | str, int | str] | None, + offload_folder: str | os.PathLike | None, offload_state_dict: bool, - model_variants: Dict[str, str], + model_variants: dict[str, str], name: str, from_flax: bool, variant: str, low_cpu_mem_usage: bool, - cached_folder: Union[str, os.PathLike], + cached_folder: str | os.PathLike, use_safetensors: bool, - dduf_entries: Optional[Dict[str, DDUFEntry]], + dduf_entries: dict[str, DDUFEntry] | None, provider_options: Any, - quantization_config: Optional[Any] = None, + disable_mmap: bool, + quantization_config: Any | None = None, ): """Helper method to load the module `name` from `library_name` and `class_name`""" from ..quantizers import PipelineQuantizationConfig @@ -801,12 +802,6 @@ def load_sub_model( # add kwargs to loading method diffusers_module = importlib.import_module(__name__.split(".")[0]) loading_kwargs = {} - if issubclass(class_obj, torch.nn.Module): - loading_kwargs["torch_dtype"] = torch_dtype - if issubclass(class_obj, diffusers_module.OnnxRuntimeModel): - loading_kwargs["provider"] = provider - loading_kwargs["sess_options"] = sess_options - loading_kwargs["provider_options"] = provider_options is_diffusers_model = issubclass(class_obj, diffusers_module.ModelMixin) @@ -821,6 +816,17 @@ def load_sub_model( and transformers_version >= version.parse("4.20.0") ) + # For transformers models >= 4.56.0, use 'dtype' instead of 'torch_dtype' to avoid deprecation warnings + if issubclass(class_obj, torch.nn.Module): + if is_transformers_model and transformers_version >= version.parse("4.56.0"): + loading_kwargs["dtype"] = torch_dtype + else: + loading_kwargs["torch_dtype"] = torch_dtype + if issubclass(class_obj, diffusers_module.OnnxRuntimeModel): + loading_kwargs["provider"] = provider + loading_kwargs["sess_options"] = sess_options + loading_kwargs["provider_options"] = provider_options + # When loading a transformers model, if the device_map is None, the weights will be initialized as opposed to diffusers. # To make default loading faster we set the `low_cpu_mem_usage=low_cpu_mem_usage` flag which is `True` by default. # This makes sure that the weights won't be initialized which significantly speeds up loading. @@ -854,6 +860,9 @@ def load_sub_model( else: loading_kwargs["low_cpu_mem_usage"] = False + if is_diffusers_model: + loading_kwargs["disable_mmap"] = disable_mmap + if is_transformers_model and is_transformers_version(">=", "4.57.0"): loading_kwargs.pop("offload_state_dict") @@ -1041,10 +1050,10 @@ def get_connected_passed_kwargs(prefix): def _get_custom_components_and_folders( pretrained_model_name: str, - config_dict: Dict[str, Any], - filenames: Optional[List[str]] = None, - variant_filenames: Optional[List[str]] = None, - variant: Optional[str] = None, + config_dict: dict[str, Any], + filenames: list[str] | None = None, + variant_filenames: list[str] | None = None, + variant: str | None = None, ): config_dict = config_dict.copy() @@ -1077,15 +1086,15 @@ def _get_custom_components_and_folders( def _get_ignore_patterns( passed_components, - model_folder_names: List[str], - model_filenames: List[str], + model_folder_names: list[str], + model_filenames: list[str], use_safetensors: bool, from_flax: bool, allow_pickle: bool, use_onnx: bool, is_onnx: bool, - variant: Optional[str] = None, -) -> List[str]: + variant: str | None = None, +) -> list[str]: if ( use_safetensors and not allow_pickle diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index 392d5fb3feb4..d675f1de04a7 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -19,9 +19,10 @@ import os import re import sys +import types from dataclasses import dataclass from pathlib import Path -from typing import Any, Callable, Dict, List, Optional, Union, get_args, get_origin +from typing import Any, Callable, Dict, List, Union, get_args, get_origin, get_type_hints import httpx import numpy as np @@ -60,6 +61,7 @@ deprecate, is_accelerate_available, is_accelerate_version, + is_bitsandbytes_version, is_hpu_available, is_torch_npu_available, is_torch_version, @@ -67,6 +69,7 @@ logging, numpy_to_pil, ) +from ..utils.distributed_utils import is_torch_dist_rank_zero from ..utils.hub_utils import _check_legacy_sharding_variant_format, load_or_create_model_card, populate_model_card from ..utils.torch_utils import empty_device_cache, get_device, is_compiled_module @@ -109,7 +112,7 @@ for library in LOADABLE_CLASSES: LIBRARIES.append(library) -SUPPORTED_DEVICE_MAP = ["balanced"] + [get_device()] +SUPPORTED_DEVICE_MAP = ["balanced"] + [get_device(), "cpu"] logger = logging.get_logger(__name__) @@ -125,7 +128,7 @@ class ImagePipelineOutput(BaseOutput): num_channels)`. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image] | np.ndarray @dataclass @@ -236,10 +239,10 @@ def __setattr__(self, name: str, value: Any): def save_pretrained( self, - save_directory: Union[str, os.PathLike], + save_directory: str | os.PathLike, safe_serialization: bool = True, - variant: Optional[str] = None, - max_shard_size: Optional[Union[int, str]] = None, + variant: str | None = None, + max_shard_size: int | str | None = None, push_to_hub: bool = False, **kwargs, ): @@ -338,6 +341,7 @@ def is_saveable_module(name, value): save_method_accept_safe = "safe_serialization" in save_method_signature.parameters save_method_accept_variant = "variant" in save_method_signature.parameters save_method_accept_max_shard_size = "max_shard_size" in save_method_signature.parameters + save_method_accept_peft_format = "save_peft_format" in save_method_signature.parameters save_kwargs = {} if save_method_accept_safe: @@ -347,6 +351,11 @@ def is_saveable_module(name, value): if save_method_accept_max_shard_size and max_shard_size is not None: # max_shard_size is expected to not be None in ModelMixin save_kwargs["max_shard_size"] = max_shard_size + if save_method_accept_peft_format: + # Set save_peft_format=False for transformers>=5.0.0 compatibility + # In transformers 5.0.0+, the default save_peft_format=True adds "base_model.model" prefix + # to adapter keys, but from_pretrained expects keys without this prefix + save_kwargs["save_peft_format"] = False save_method(os.path.join(save_directory, pipeline_component_name), **save_kwargs) @@ -443,7 +452,10 @@ def module_is_sequentially_offloaded(module): _, _, is_loaded_in_8bit_bnb = _check_bnb_status(module) - if is_loaded_in_8bit_bnb: + # https://github.com/huggingface/accelerate/pull/3907 + if is_loaded_in_8bit_bnb and ( + is_bitsandbytes_version("<", "0.48.0") or is_accelerate_version("<", "1.13.0.dev0") + ): return False return hasattr(module, "_hf_hook") and ( @@ -462,8 +474,7 @@ def module_is_offloaded(module): pipeline_is_sequentially_offloaded = any( module_is_sequentially_offloaded(module) for _, module in self.components.items() ) - - is_pipeline_device_mapped = self.hf_device_map is not None and len(self.hf_device_map) > 1 + is_pipeline_device_mapped = self._is_pipeline_device_mapped() if is_pipeline_device_mapped: raise ValueError( "It seems like you have activated a device mapping strategy on the pipeline which doesn't allow explicit device placement using `to()`. You can call `reset_device_map()` to remove the existing device map from the pipeline." @@ -522,9 +533,10 @@ def module_is_offloaded(module): f"The module '{module.__class__.__name__}' has been loaded in `bitsandbytes` {'4bit' if is_loaded_in_4bit_bnb else '8bit'} and conversion to {dtype} is not supported. Module is still in {'4bit' if is_loaded_in_4bit_bnb else '8bit'} precision." ) - if is_loaded_in_8bit_bnb and device is not None: + if is_loaded_in_8bit_bnb and device is not None and is_bitsandbytes_version("<", "0.48.0"): logger.warning( f"The module '{module.__class__.__name__}' has been loaded in `bitsandbytes` 8bit and moving it to {device} via `.to()` is not supported. Module is still on {module.device}." + "You need to upgrade bitsandbytes to at least 0.48.0" ) # Note: we also handle this at the ModelMixin level. The reason for doing it here too is that modeling @@ -541,6 +553,14 @@ def module_is_offloaded(module): # https://github.com/huggingface/transformers/pull/33122. So, we guard this accordingly. if is_loaded_in_4bit_bnb and device is not None and is_transformers_version(">", "4.44.0"): module.to(device=device) + # added here https://github.com/huggingface/transformers/pull/43258 + if ( + is_loaded_in_8bit_bnb + and device is not None + and is_transformers_version(">", "4.58.0") + and is_bitsandbytes_version(">=", "0.48.0") + ): + module.to(device=device) elif not is_loaded_in_4bit_bnb and not is_loaded_in_8bit_bnb and not is_group_offloaded: module.to(device, dtype) @@ -591,7 +611,7 @@ def dtype(self) -> torch.dtype: @classmethod @validate_hf_hub_args - def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs) -> Self: + def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike, **kwargs) -> Self: r""" Instantiate a PyTorch diffusion pipeline from pretrained pipeline weights. @@ -707,6 +727,9 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P loading `from_flax`. dduf_file(`str`, *optional*): Load weights from the specified dduf file. + disable_mmap ('bool', *optional*, defaults to 'False'): + Whether to disable mmap when loading a Safetensors model. This option can perform better when the model + is on a network mount or hard drive, which may not handle the seeky-ness of mmap very well. > [!TIP] > To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with `hf > auth login`. @@ -758,6 +781,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P use_onnx = kwargs.pop("use_onnx", None) load_connected_pipeline = kwargs.pop("load_connected_pipeline", False) quantization_config = kwargs.pop("quantization_config", None) + disable_mmap = kwargs.pop("disable_mmap", False) if torch_dtype is not None and not isinstance(torch_dtype, dict) and not isinstance(torch_dtype, torch.dtype): torch_dtype = torch.float32 @@ -982,7 +1006,11 @@ def load_module(name, value): # 7. Load each module in the pipeline current_device_map = None _maybe_warn_for_wrong_component_in_quant_config(init_dict, quantization_config) - for name, (library_name, class_name) in logging.tqdm(init_dict.items(), desc="Loading pipeline components..."): + logging_tqdm_kwargs = {"desc": "Loading pipeline components..."} + if not is_torch_dist_rank_zero(): + logging_tqdm_kwargs["disable"] = True + + for name, (library_name, class_name) in logging.tqdm(init_dict.items(), **logging_tqdm_kwargs): # 7.1 device_map shenanigans if final_device_map is not None: if isinstance(final_device_map, dict) and len(final_device_map) > 0: @@ -1041,6 +1069,7 @@ def load_module(name, value): use_safetensors=use_safetensors, dduf_entries=dduf_entries, provider_options=provider_options, + disable_mmap=disable_mmap, quantization_config=quantization_config, ) logger.info( @@ -1147,7 +1176,7 @@ def remove_all_hooks(self): accelerate.hooks.remove_hook_from_module(model, recurse=True) self._all_hooks = [] - def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None): + def enable_model_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None): r""" Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the accelerator when its @@ -1164,7 +1193,7 @@ def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[t """ self._maybe_raise_error_if_group_offload_active(raise_error=True) - is_pipeline_device_mapped = self.hf_device_map is not None and len(self.hf_device_map) > 1 + is_pipeline_device_mapped = self._is_pipeline_device_mapped() if is_pipeline_device_mapped: raise ValueError( "It seems like you have activated a device mapping strategy on the pipeline so calling `enable_model_cpu_offload() isn't allowed. You can call `reset_device_map()` first and then call `enable_model_cpu_offload()`." @@ -1218,7 +1247,9 @@ def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[t # This is because the model would already be placed on a CUDA device. _, _, is_loaded_in_8bit_bnb = _check_bnb_status(model) - if is_loaded_in_8bit_bnb: + if is_loaded_in_8bit_bnb and ( + is_transformers_version("<", "4.58.0") or is_bitsandbytes_version("<", "0.48.0") + ): logger.info( f"Skipping the hook placement for the {model.__class__.__name__} as it is loaded in `bitsandbytes` 8bit." ) @@ -1263,7 +1294,7 @@ def maybe_free_model_hooks(self): # make sure the model is in the same state as before calling it self.enable_model_cpu_offload(device=getattr(self, "_offload_device", "cuda")) - def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None): + def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None): r""" Offloads all models to CPU using 🤗 Accelerate, significantly reducing memory usage. When called, the state dicts of all `torch.nn.Module` components (except those in `self._exclude_from_cpu_offload`) are saved to CPU @@ -1286,7 +1317,7 @@ def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Un raise ImportError("`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher") self.remove_all_hooks() - is_pipeline_device_mapped = self.hf_device_map is not None and len(self.hf_device_map) > 1 + is_pipeline_device_mapped = self._is_pipeline_device_mapped() if is_pipeline_device_mapped: raise ValueError( "It seems like you have activated a device mapping strategy on the pipeline so calling `enable_sequential_cpu_offload() isn't allowed. You can call `reset_device_map()` first and then call `enable_sequential_cpu_offload()`." @@ -1335,13 +1366,13 @@ def enable_group_offload( onload_device: torch.device, offload_device: torch.device = torch.device("cpu"), offload_type: str = "block_level", - num_blocks_per_group: Optional[int] = None, + num_blocks_per_group: int | None = None, non_blocking: bool = False, use_stream: bool = False, record_stream: bool = False, low_cpu_mem_usage=False, - offload_to_disk_path: Optional[str] = None, - exclude_modules: Optional[Union[str, List[str]]] = None, + offload_to_disk_path: str | None = None, + exclude_modules: str | list[str] | None = None, ) -> None: r""" Applies group offloading to the internal layers of a torch.nn.Module. To understand what group offloading is, @@ -1472,7 +1503,7 @@ def reset_device_map(self): @classmethod @validate_hf_hub_args - def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]: + def download(cls, pretrained_model_name, **kwargs) -> str | os.PathLike: r""" Download and cache a PyTorch diffusion pipeline from pretrained pipeline weights. @@ -1568,7 +1599,7 @@ def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]: use_onnx = kwargs.pop("use_onnx", None) load_connected_pipeline = kwargs.pop("load_connected_pipeline", False) trust_remote_code = kwargs.pop("trust_remote_code", False) - dduf_file: Optional[Dict[str, DDUFEntry]] = kwargs.pop("dduf_file", None) + dduf_file: dict[str, DDUFEntry] | None = kwargs.pop("dduf_file", None) if dduf_file: if custom_pipeline: @@ -1592,7 +1623,7 @@ def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]: allow_patterns = None ignore_patterns = None - model_info_call_error: Optional[Exception] = None + model_info_call_error: Exception | None = None if not local_files_only: try: info = model_info(pretrained_model_name, token=token, revision=revision) @@ -1813,19 +1844,48 @@ def _get_signature_keys(cls, obj): @classmethod def _get_signature_types(cls): signature_types = {} - for k, v in inspect.signature(cls.__init__).parameters.items(): - if inspect.isclass(v.annotation): - signature_types[k] = (v.annotation,) - elif get_origin(v.annotation) == Union: - signature_types[k] = get_args(v.annotation) - elif get_origin(v.annotation) in [List, Dict, list, dict]: - signature_types[k] = (v.annotation,) + # Use get_type_hints to properly resolve string annotations (from __future__ import annotations) + try: + type_hints = get_type_hints(cls.__init__) + except Exception: + # Fallback to direct annotation access if get_type_hints fails + type_hints = {} + for k, v in inspect.signature(cls.__init__).parameters.items(): + if v.annotation != inspect.Parameter.empty: + type_hints[k] = v.annotation + + # Get all parameters from the signature to ensure we don't miss any + all_params = inspect.signature(cls.__init__).parameters + + for param_name, param in all_params.items(): + # Skip 'self' parameter + if param_name == "self": + continue + + # If we have type hints, use them + if param_name in type_hints: + annotation = type_hints[param_name] + if inspect.isclass(annotation): + signature_types[param_name] = (annotation,) + elif get_origin(annotation) == Union: + signature_types[param_name] = get_args(annotation) + elif isinstance(annotation, types.UnionType): + # Handle PEP 604 union syntax (X | Y) introduced in Python 3.10+ + signature_types[param_name] = get_args(annotation) + elif get_origin(annotation) in [List, Dict, list, dict]: + signature_types[param_name] = (annotation,) + else: + logger.warning(f"cannot get type annotation for Parameter {param_name} of {cls}.") + # Still add it with empty signature so it's in expected_types + signature_types[param_name] = (inspect.Signature.empty,) else: - logger.warning(f"cannot get type annotation for Parameter {k} of {cls}.") + # No type annotation found - add with empty signature + signature_types[param_name] = (inspect.Signature.empty,) + return signature_types @property - def parameters(self) -> Dict[str, Any]: + def parameters(self) -> dict[str, Any]: r""" The `self.parameters` property can be useful to run different pipelines with the same weights and configurations without reallocating additional memory. @@ -1855,7 +1915,7 @@ def parameters(self) -> Dict[str, Any]: return pipeline_parameters @property - def components(self) -> Dict[str, Any]: + def components(self) -> dict[str, Any]: r""" The `self.components` property can be useful to run different pipelines with the same weights and configurations without reallocating additional memory. @@ -1908,17 +1968,21 @@ def progress_bar(self, iterable=None, total=None): f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}." ) + progress_bar_config = dict(self._progress_bar_config) + if "disable" not in progress_bar_config: + progress_bar_config["disable"] = not is_torch_dist_rank_zero() + if iterable is not None: - return tqdm(iterable, **self._progress_bar_config) + return tqdm(iterable, **progress_bar_config) elif total is not None: - return tqdm(total=total, **self._progress_bar_config) + return tqdm(total=total, **progress_bar_config) else: raise ValueError("Either `total` or `iterable` has to be defined.") def set_progress_bar_config(self, **kwargs): self._progress_bar_config = kwargs - def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None): + def enable_xformers_memory_efficient_attention(self, attention_op: Callable | None = None): r""" Enable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/). When this option is enabled, you should observe lower GPU memory usage and a potential speed up during inference. Speed @@ -1955,9 +2019,7 @@ def disable_xformers_memory_efficient_attention(self): """ self.set_use_memory_efficient_attention_xformers(False) - def set_use_memory_efficient_attention_xformers( - self, valid: bool, attention_op: Optional[Callable] = None - ) -> None: + def set_use_memory_efficient_attention_xformers(self, valid: bool, attention_op: Callable | None = None) -> None: # Recursively walk through all the children. # Any children which exposes the set_use_memory_efficient_attention_xformers method # gets the message @@ -1975,7 +2037,7 @@ def fn_recursive_set_mem_eff(module: torch.nn.Module): for module in modules: fn_recursive_set_mem_eff(module) - def enable_attention_slicing(self, slice_size: Optional[Union[str, int]] = "auto"): + def enable_attention_slicing(self, slice_size: str | int = "auto"): r""" Enable sliced attention computation. When this option is enabled, the attention module splits the input tensor in slices to compute attention in several steps. For more than one attention head, the computation is performed @@ -2020,7 +2082,7 @@ def disable_attention_slicing(self): # set slice_size = `None` to disable `attention slicing` self.enable_attention_slicing(None) - def set_attention_slice(self, slice_size: Optional[int]): + def set_attention_slice(self, slice_size: int | None): module_names, _ = self._get_signature_keys(self) modules = [getattr(self, n, None) for n in module_names] modules = [m for m in modules if isinstance(m, torch.nn.Module) and hasattr(m, "set_attention_slice")] @@ -2154,7 +2216,7 @@ def from_pipe(cls, pipeline, **kwargs): return new_pipeline def _maybe_raise_error_if_group_offload_active( - self, raise_error: bool = False, module: Optional[torch.nn.Module] = None + self, raise_error: bool = False, module: torch.nn.Module | None = None ) -> bool: from ..hooks.group_offloading import _is_group_offload_enabled @@ -2171,6 +2233,21 @@ def _maybe_raise_error_if_group_offload_active( return True return False + def _is_pipeline_device_mapped(self): + # We support passing `device_map="cuda"`, for example. This is helpful, in case + # users want to pass `device_map="cpu"` when initializing a pipeline. This explicit declaration is desirable + # in limited VRAM environments because quantized models often initialize directly on the accelerator. + device_map = self.hf_device_map + is_device_type_map = False + if isinstance(device_map, str): + try: + torch.device(device_map) + is_device_type_map = True + except RuntimeError: + pass + + return not is_device_type_map and isinstance(device_map, dict) and len(device_map) > 1 + class StableDiffusionMixin: r""" diff --git a/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py b/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py index 1d718a4852a4..604e51d88583 100644 --- a/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py +++ b/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py @@ -16,7 +16,7 @@ import inspect import re import urllib.parse as ul -from typing import Callable, List, Optional, Tuple, Union +from typing import Callable import torch from transformers import T5EncoderModel, T5Tokenizer @@ -181,10 +181,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -199,15 +199,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -302,15 +302,15 @@ def __init__( # Adapted from diffusers.pipelines.deepfloyd_if.pipeline_if.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], do_classifier_free_guidance: bool = True, negative_prompt: str = "", num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, + device: torch.device | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, clean_caption: bool = False, max_sequence_length: int = 120, **kwargs, @@ -319,9 +319,9 @@ def encode_prompt( Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For PixArt-Alpha, this should be "". @@ -687,50 +687,50 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, negative_prompt: str = "", num_inference_steps: int = 20, - timesteps: List[int] = None, - sigmas: List[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, guidance_scale: float = 4.5, - num_images_per_prompt: Optional[int] = 1, - height: Optional[int] = None, - width: Optional[int] = None, + num_images_per_prompt: int | None = 1, + height: int | None = None, + width: int | None = None, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, clean_caption: bool = True, use_resolution_binning: bool = True, max_sequence_length: int = 120, **kwargs, - ) -> Union[ImagePipelineOutput, Tuple]: + ) -> ImagePipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_inference_steps (`int`, *optional*, defaults to 100): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -749,7 +749,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -862,8 +862,12 @@ def __call__( prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) # 4. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) # 5. Prepare latents. diff --git a/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py b/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py index bb169ac5c443..286695aa8eb9 100644 --- a/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py +++ b/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py @@ -16,7 +16,7 @@ import inspect import re import urllib.parse as ul -from typing import Callable, List, Optional, Tuple, Union +from typing import Callable import torch from transformers import T5EncoderModel, T5Tokenizer @@ -125,10 +125,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -143,15 +143,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -246,15 +246,15 @@ def __init__( # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.encode_prompt with 120->300 def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], do_classifier_free_guidance: bool = True, negative_prompt: str = "", num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, + device: torch.device | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, clean_caption: bool = False, max_sequence_length: int = 300, **kwargs, @@ -263,9 +263,9 @@ def encode_prompt( Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For PixArt-Alpha, this should be "". @@ -632,50 +632,50 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, negative_prompt: str = "", num_inference_steps: int = 20, - timesteps: List[int] = None, - sigmas: List[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, guidance_scale: float = 4.5, - num_images_per_prompt: Optional[int] = 1, - height: Optional[int] = None, - width: Optional[int] = None, + num_images_per_prompt: int | None = 1, + height: int | None = None, + width: int | None = None, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, clean_caption: bool = True, use_resolution_binning: bool = True, max_sequence_length: int = 300, **kwargs, - ) -> Union[ImagePipelineOutput, Tuple]: + ) -> ImagePipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_inference_steps (`int`, *optional*, defaults to 100): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -694,7 +694,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -806,8 +806,12 @@ def __call__( prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) # 4. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) # 5. Prepare latents. diff --git a/src/diffusers/pipelines/prx/__init__.py b/src/diffusers/pipelines/prx/__init__.py index 87aaefbd1368..ad2948e92e04 100644 --- a/src/diffusers/pipelines/prx/__init__.py +++ b/src/diffusers/pipelines/prx/__init__.py @@ -24,14 +24,25 @@ else: _import_structure["pipeline_prx"] = ["PRXPipeline"] -# Import T5GemmaEncoder for pipeline loading compatibility +# Wrap T5GemmaEncoder to pass config.encoder (T5GemmaModuleConfig) instead of the +# composite T5GemmaConfig, which lacks flat attributes expected by T5GemmaEncoder.__init__. try: if is_transformers_available(): import transformers - from transformers.models.t5gemma.modeling_t5gemma import T5GemmaEncoder + from transformers.models.t5gemma.modeling_t5gemma import T5GemmaEncoder as _T5GemmaEncoder + + class T5GemmaEncoder(_T5GemmaEncoder): + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): + if "config" not in kwargs: + from transformers.models.t5gemma.configuration_t5gemma import T5GemmaConfig + + config = T5GemmaConfig.from_pretrained(pretrained_model_name_or_path) + if hasattr(config, "encoder"): + kwargs["config"] = config.encoder + return super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs) _additional_imports["T5GemmaEncoder"] = T5GemmaEncoder - # Patch transformers module directly for serialization if not hasattr(transformers, "T5GemmaEncoder"): transformers.T5GemmaEncoder = T5GemmaEncoder except ImportError: diff --git a/src/diffusers/pipelines/prx/pipeline_output.py b/src/diffusers/pipelines/prx/pipeline_output.py index ea1bc9bf418a..81f92c294735 100644 --- a/src/diffusers/pipelines/prx/pipeline_output.py +++ b/src/diffusers/pipelines/prx/pipeline_output.py @@ -13,7 +13,6 @@ # limitations under the License. from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -27,9 +26,9 @@ class PRXPipelineOutput(BaseOutput): Output class for PRX pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image] | np.ndarray diff --git a/src/diffusers/pipelines/prx/pipeline_prx.py b/src/diffusers/pipelines/prx/pipeline_prx.py index 873f25316e6d..e14815b91c41 100644 --- a/src/diffusers/pipelines/prx/pipeline_prx.py +++ b/src/diffusers/pipelines/prx/pipeline_prx.py @@ -16,9 +16,8 @@ import inspect import re import urllib.parse as ul -from typing import Callable, Dict, List, Optional, Union +from typing import Callable -import ftfy import torch from transformers import ( AutoTokenizer, @@ -34,13 +33,13 @@ from diffusers.pipelines.pipeline_utils import DiffusionPipeline from diffusers.pipelines.prx.pipeline_output import PRXPipelineOutput from diffusers.schedulers import FlowMatchEulerDiscreteScheduler -from diffusers.utils import ( - logging, - replace_example_docstring, -) +from diffusers.utils import is_ftfy_available, logging, replace_example_docstring from diffusers.utils.torch_utils import randn_tensor +if is_ftfy_available(): + import ftfy + DEFAULT_RESOLUTION = 512 ASPECT_RATIO_256_BIN = { @@ -284,9 +283,9 @@ def __init__( transformer: PRXTransformer2DModel, scheduler: FlowMatchEulerDiscreteScheduler, text_encoder: T5GemmaEncoder, - tokenizer: Union[T5TokenizerFast, GemmaTokenizerFast, AutoTokenizer], - vae: Optional[Union[AutoencoderKL, AutoencoderDC]] = None, - default_sample_size: Optional[int] = DEFAULT_RESOLUTION, + tokenizer: T5TokenizerFast | GemmaTokenizerFast | AutoTokenizer, + vae: AutoencoderKL | AutoencoderDC | None = None, + default_sample_size: int | None = DEFAULT_RESOLUTION, ): super().__init__() @@ -352,8 +351,8 @@ def prepare_latents( width: int, dtype: torch.dtype, device: torch.device, - generator: Optional[torch.Generator] = None, - latents: Optional[torch.Tensor] = None, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, ): """Prepare initial latents for the diffusion process.""" if latents is None: @@ -370,15 +369,15 @@ def prepare_latents( def encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + device: torch.device | None = None, do_classifier_free_guidance: bool = True, negative_prompt: str = "", num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - prompt_attention_mask: Optional[torch.BoolTensor] = None, - negative_prompt_attention_mask: Optional[torch.BoolTensor] = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + prompt_attention_mask: torch.BoolTensor | None = None, + negative_prompt_attention_mask: torch.BoolTensor | None = None, ): """Encode text prompt using standard text encoder and tokenizer, or use precomputed embeddings.""" if device is None: @@ -420,7 +419,7 @@ def encode_prompt( negative_prompt_attention_mask if do_classifier_free_guidance else None, ) - def _tokenize_prompts(self, prompts: List[str], device: torch.device): + def _tokenize_prompts(self, prompts: list[str], device: torch.device): """Tokenize and clean prompts.""" cleaned = [self.text_preprocessor.clean_text(text) for text in prompts] tokens = self.tokenizer( @@ -435,7 +434,7 @@ def _tokenize_prompts(self, prompts: List[str], device: torch.device): def _encode_prompt_standard( self, - prompt: List[str], + prompt: list[str], device: torch.device, do_classifier_free_guidance: bool = True, negative_prompt: str = "", @@ -473,13 +472,13 @@ def _encode_prompt_standard( def check_inputs( self, - prompt: Union[str, List[str]], + prompt: str | list[str], height: int, width: int, guidance_scale: float, - callback_on_step_end_tensor_inputs: Optional[List[str]] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, + callback_on_step_end_tensor_inputs: list[str] | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, ): """Check that all inputs are in correct format.""" if prompt is not None and prompt_embeds is not None: @@ -527,31 +526,31 @@ def check_inputs( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, negative_prompt: str = "", - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 28, - timesteps: List[int] = None, + timesteps: list[int] = None, guidance_scale: float = 4.0, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - prompt_attention_mask: Optional[torch.BoolTensor] = None, - negative_prompt_attention_mask: Optional[torch.BoolTensor] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + prompt_attention_mask: torch.BoolTensor | None = None, + negative_prompt_attention_mask: torch.BoolTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, use_resolution_binning: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], ): """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds` instead. negative_prompt (`str`, *optional*, defaults to `""`): @@ -564,7 +563,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 28): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. @@ -576,7 +575,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -610,7 +609,7 @@ def __call__( with the following arguments: `callback_on_step_end(self, step, timestep, callback_kwargs)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include tensors that are listed in the `._callback_tensor_inputs` attribute. diff --git a/src/diffusers/pipelines/qwenimage/pipeline_output.py b/src/diffusers/pipelines/qwenimage/pipeline_output.py index eef4b60e3770..e4ed06856e6a 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_output.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -13,9 +12,9 @@ class QwenImagePipelineOutput(BaseOutput): Output class for Stable Diffusion pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image] | np.ndarray diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py index 33dc2039b986..1715aa4d4250 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -72,10 +72,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -90,15 +90,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -187,9 +187,9 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor def _get_qwen_prompt_embeds( self, - prompt: Union[str, List[str]] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt: str | list[str] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -225,17 +225,17 @@ def _get_qwen_prompt_embeds( def encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, max_sequence_length: int = 1024, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -254,13 +254,17 @@ def encode_prompt( prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, device) prompt_embeds = prompt_embeds[:, :max_sequence_length] - prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length] - _, seq_len, _ = prompt_embeds.shape prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) - prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + + if prompt_embeds_mask is not None: + prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length] + prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) + prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + + if prompt_embeds_mask.all(): + prompt_embeds_mask = None return prompt_embeds, prompt_embeds_mask @@ -307,15 +311,6 @@ def check_inputs( f" {negative_prompt_embeds}. Please make sure to only forward one of the two." ) - if prompt_embeds is not None and prompt_embeds_mask is None: - raise ValueError( - "If `prompt_embeds` are provided, `prompt_embeds_mask` also have to be passed. Make sure to generate `prompt_embeds_mask` from the same text encoder that was used to generate `prompt_embeds`." - ) - if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None: - raise ValueError( - "If `negative_prompt_embeds` are provided, `negative_prompt_embeds_mask` also have to be passed. Make sure to generate `negative_prompt_embeds_mask` from the same text encoder that was used to generate `negative_prompt_embeds`." - ) - if max_sequence_length is not None and max_sequence_length > 1024: raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}") @@ -452,36 +447,36 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, true_cfg_scale: float = 4.0, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, - guidance_scale: Optional[float] = None, + sigmas: list[float] | None = None, + guidance_scale: float | None = None, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is not greater than `1`). @@ -499,7 +494,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -515,7 +510,7 @@ def __call__( enable classifier-free guidance computations). num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -543,7 +538,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -672,11 +667,6 @@ def __call__( if self.attention_kwargs is None: self._attention_kwargs = {} - txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None - negative_txt_seq_lens = ( - negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None - ) - # 6. Denoising loop self.scheduler.set_begin_index(0) with self.progress_bar(total=num_inference_steps) as progress_bar: @@ -695,7 +685,6 @@ def __call__( encoder_hidden_states_mask=prompt_embeds_mask, encoder_hidden_states=prompt_embeds, img_shapes=img_shapes, - txt_seq_lens=txt_seq_lens, attention_kwargs=self.attention_kwargs, return_dict=False, )[0] @@ -709,7 +698,6 @@ def __call__( encoder_hidden_states_mask=negative_prompt_embeds_mask, encoder_hidden_states=negative_prompt_embeds, img_shapes=img_shapes, - txt_seq_lens=negative_txt_seq_lens, attention_kwargs=self.attention_kwargs, return_dict=False, )[0] diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet.py index 5111096d93c1..85a936f9ec24 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -117,7 +117,7 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -132,10 +132,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -150,15 +150,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -218,7 +218,7 @@ def __init__( text_encoder: Qwen2_5_VLForConditionalGeneration, tokenizer: Qwen2Tokenizer, transformer: QwenImageTransformer2DModel, - controlnet: Union[QwenImageControlNetModel, QwenImageMultiControlNetModel], + controlnet: QwenImageControlNetModel | QwenImageMultiControlNetModel, ): super().__init__() @@ -251,9 +251,9 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor # Coped from diffusers.pipelines.qwenimage.pipeline_qwenimage.get_qwen_prompt_embeds def _get_qwen_prompt_embeds( self, - prompt: Union[str, List[str]] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt: str | list[str] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -290,17 +290,17 @@ def _get_qwen_prompt_embeds( # Coped from diffusers.pipelines.qwenimage.pipeline_qwenimage.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, max_sequence_length: int = 1024, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -321,8 +321,13 @@ def encode_prompt( _, seq_len, _ = prompt_embeds.shape prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) - prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + + if prompt_embeds_mask is not None: + prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) + prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + + if prompt_embeds_mask.all(): + prompt_embeds_mask = None return prompt_embeds, prompt_embeds_mask @@ -369,15 +374,6 @@ def check_inputs( f" {negative_prompt_embeds}. Please make sure to only forward one of the two." ) - if prompt_embeds is not None and prompt_embeds_mask is None: - raise ValueError( - "If `prompt_embeds` are provided, `prompt_embeds_mask` also have to be passed. Make sure to generate `prompt_embeds_mask` from the same text encoder that was used to generate `prompt_embeds`." - ) - if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None: - raise ValueError( - "If `negative_prompt_embeds` are provided, `negative_prompt_embeds_mask` also have to be passed. Make sure to generate `negative_prompt_embeds_mask` from the same text encoder that was used to generate `negative_prompt_embeds`." - ) - if max_sequence_length is not None and max_sequence_length > 1024: raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}") @@ -552,40 +548,40 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, true_cfg_scale: float = 4.0, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, - guidance_scale: Optional[float] = None, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, + sigmas: list[float] | None = None, + guidance_scale: float | None = None, + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, control_image: PipelineImageInput = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, + controlnet_conditioning_scale: float | list[float] = 1.0, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is not greater than `1`). @@ -603,7 +599,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -619,7 +615,7 @@ def __call__( enable classifier-free guidance computations). num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -647,7 +643,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -909,7 +905,6 @@ def __call__( encoder_hidden_states=prompt_embeds, encoder_hidden_states_mask=prompt_embeds_mask, img_shapes=img_shapes, - txt_seq_lens=prompt_embeds_mask.sum(dim=1).tolist(), return_dict=False, ) @@ -920,7 +915,6 @@ def __call__( encoder_hidden_states=prompt_embeds, encoder_hidden_states_mask=prompt_embeds_mask, img_shapes=img_shapes, - txt_seq_lens=prompt_embeds_mask.sum(dim=1).tolist(), controlnet_block_samples=controlnet_block_samples, attention_kwargs=self.attention_kwargs, return_dict=False, @@ -935,7 +929,6 @@ def __call__( encoder_hidden_states_mask=negative_prompt_embeds_mask, encoder_hidden_states=negative_prompt_embeds, img_shapes=img_shapes, - txt_seq_lens=negative_prompt_embeds_mask.sum(dim=1).tolist(), controlnet_block_samples=controlnet_block_samples, attention_kwargs=self.attention_kwargs, return_dict=False, diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet_inpaint.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet_inpaint.py index 102a813ab582..b1da59cb4f6c 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet_inpaint.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet_inpaint.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -90,7 +90,7 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -105,10 +105,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -123,15 +123,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -233,9 +233,9 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor # Coped from diffusers.pipelines.qwenimage.pipeline_qwenimage.get_qwen_prompt_embeds def _get_qwen_prompt_embeds( self, - prompt: Union[str, List[str]] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt: str | list[str] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -272,16 +272,16 @@ def _get_qwen_prompt_embeds( # Coped from diffusers.pipelines.qwenimage.pipeline_qwenimage.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, max_sequence_length: int = 1024, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -305,6 +305,9 @@ def encode_prompt( prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + if prompt_embeds_mask is not None and prompt_embeds_mask.all(): + prompt_embeds_mask = None + return prompt_embeds, prompt_embeds_mask def check_inputs( @@ -592,41 +595,41 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, true_cfg_scale: float = 4.0, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 1.0, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, + control_guidance_start: float | list[float] = 0.0, + control_guidance_end: float | list[float] = 1.0, control_image: PipelineImageInput = None, control_mask: PipelineImageInput = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, + controlnet_conditioning_scale: float | list[float] = 1.0, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is not greater than `1`). @@ -639,7 +642,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -651,7 +654,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -679,7 +682,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -852,7 +855,6 @@ def __call__( encoder_hidden_states=prompt_embeds, encoder_hidden_states_mask=prompt_embeds_mask, img_shapes=img_shapes, - txt_seq_lens=prompt_embeds_mask.sum(dim=1).tolist(), return_dict=False, ) @@ -863,7 +865,6 @@ def __call__( encoder_hidden_states=prompt_embeds, encoder_hidden_states_mask=prompt_embeds_mask, img_shapes=img_shapes, - txt_seq_lens=prompt_embeds_mask.sum(dim=1).tolist(), controlnet_block_samples=controlnet_block_samples, attention_kwargs=self.attention_kwargs, return_dict=False, @@ -878,7 +879,6 @@ def __call__( encoder_hidden_states_mask=negative_prompt_embeds_mask, encoder_hidden_states=negative_prompt_embeds, img_shapes=img_shapes, - txt_seq_lens=negative_prompt_embeds_mask.sum(dim=1).tolist(), controlnet_block_samples=controlnet_block_samples, attention_kwargs=self.attention_kwargs, return_dict=False, diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py index ed37b238c8c9..15e72a010ce5 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py @@ -14,7 +14,7 @@ import inspect import math -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -81,10 +81,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -99,15 +99,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -140,7 +140,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -225,10 +225,10 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor def _get_qwen_prompt_embeds( self, - prompt: Union[str, List[str]] = None, - image: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt: str | list[str] = None, + image: torch.Tensor | None = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -272,18 +272,18 @@ def _get_qwen_prompt_embeds( def encode_prompt( self, - prompt: Union[str, List[str]], - image: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, + prompt: str | list[str], + image: torch.Tensor | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, max_sequence_length: int = 1024, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded image (`torch.Tensor`, *optional*): image to be encoded @@ -309,6 +309,9 @@ def encode_prompt( prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + if prompt_embeds_mask is not None and prompt_embeds_mask.all(): + prompt_embeds_mask = None + return prompt_embeds, prompt_embeds_mask def check_inputs( @@ -547,43 +550,43 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - image: Optional[PipelineImageInput] = None, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + image: PipelineImageInput | None = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, true_cfg_scale: float = 4.0, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, - guidance_scale: Optional[float] = None, + sigmas: list[float] | None = None, + guidance_scale: float | None = None, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image latents as `image`, but if passing latents directly it is not encoded again. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is not greater than `1`). @@ -601,7 +604,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -617,7 +620,7 @@ def __call__( enable classifier-free guidance computations). num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -645,7 +648,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -793,11 +796,6 @@ def __call__( if self.attention_kwargs is None: self._attention_kwargs = {} - txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None - negative_txt_seq_lens = ( - negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None - ) - # 6. Denoising loop self.scheduler.set_begin_index(0) with self.progress_bar(total=num_inference_steps) as progress_bar: @@ -821,7 +819,6 @@ def __call__( encoder_hidden_states_mask=prompt_embeds_mask, encoder_hidden_states=prompt_embeds, img_shapes=img_shapes, - txt_seq_lens=txt_seq_lens, attention_kwargs=self.attention_kwargs, return_dict=False, )[0] @@ -836,7 +833,6 @@ def __call__( encoder_hidden_states_mask=negative_prompt_embeds_mask, encoder_hidden_states=negative_prompt_embeds, img_shapes=img_shapes, - txt_seq_lens=negative_txt_seq_lens, attention_kwargs=self.attention_kwargs, return_dict=False, )[0] diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_inpaint.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_inpaint.py index d54d1881fa4e..20a2748bc7f9 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_inpaint.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_inpaint.py @@ -14,7 +14,7 @@ import inspect import math -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -82,10 +82,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -100,15 +100,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -141,7 +141,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -236,10 +236,10 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline._get_qwen_prompt_embeds def _get_qwen_prompt_embeds( self, - prompt: Union[str, List[str]] = None, - image: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt: str | list[str] = None, + image: torch.Tensor | None = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -284,18 +284,18 @@ def _get_qwen_prompt_embeds( # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - image: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, + prompt: str | list[str], + image: torch.Tensor | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, max_sequence_length: int = 1024, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded image (`torch.Tensor`, *optional*): image to be encoded @@ -321,6 +321,9 @@ def encode_prompt( prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + if prompt_embeds_mask is not None and prompt_embeds_mask.all(): + prompt_embeds_mask = None + return prompt_embeds, prompt_embeds_mask # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_inpaint.QwenImageInpaintPipeline.check_inputs @@ -375,14 +378,6 @@ def check_inputs( f" {negative_prompt_embeds}. Please make sure to only forward one of the two." ) - if prompt_embeds is not None and prompt_embeds_mask is None: - raise ValueError( - "If `prompt_embeds` are provided, `prompt_embeds_mask` also have to be passed. Make sure to generate `prompt_embeds_mask` from the same text encoder that was used to generate `prompt_embeds`." - ) - if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None: - raise ValueError( - "If `negative_prompt_embeds` are provided, `negative_prompt_embeds_mask` also have to be passed. Make sure to generate `negative_prompt_embeds_mask` from the same text encoder that was used to generate `negative_prompt_embeds`." - ) if padding_mask_crop is not None: if not isinstance(image, PIL.Image.Image): raise ValueError( @@ -680,47 +675,47 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - image: Optional[PipelineImageInput] = None, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + image: PipelineImageInput | None = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, mask_image: PipelineImageInput = None, masked_image_latents: PipelineImageInput = None, true_cfg_scale: float = 4.0, - height: Optional[int] = None, - width: Optional[int] = None, - padding_mask_crop: Optional[int] = None, + height: int | None = None, + width: int | None = None, + padding_mask_crop: int | None = None, strength: float = 0.6, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, - guidance_scale: Optional[float] = None, + sigmas: list[float] | None = None, + guidance_scale: float | None = None, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image latents as `image`, but if passing latents directly it is not encoded again. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is not greater than `1`). @@ -731,14 +726,14 @@ def __call__( enabled by setting `true_cfg_scale > 1` and a provided `negative_prompt`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one color channel (L) instead of 3, so the expected shape for pytorch tensor would be `(B, 1, H, W)`, `(B, H, W)`, `(1, H, W)`, `(H, W)`. And for numpy array would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W, 1)`, or `(H, W)`. - mask_image_latent (`torch.Tensor`, `List[torch.Tensor]`): + mask_image_latent (`torch.Tensor`, `list[torch.Tensor]`): `Tensor` representing an image batch to mask `image` generated by VAE. If not provided, the mask latents tensor will ge generated by `mask_image`. height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): @@ -761,7 +756,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -777,7 +772,7 @@ def __call__( enable classifier-free guidance computations). num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -805,7 +800,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1008,11 +1003,6 @@ def __call__( if self.attention_kwargs is None: self._attention_kwargs = {} - txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None - negative_txt_seq_lens = ( - negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None - ) - # 6. Denoising loop with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): @@ -1035,7 +1025,6 @@ def __call__( encoder_hidden_states_mask=prompt_embeds_mask, encoder_hidden_states=prompt_embeds, img_shapes=img_shapes, - txt_seq_lens=txt_seq_lens, attention_kwargs=self.attention_kwargs, return_dict=False, )[0] @@ -1050,7 +1039,6 @@ def __call__( encoder_hidden_states_mask=negative_prompt_embeds_mask, encoder_hidden_states=negative_prompt_embeds, img_shapes=img_shapes, - txt_seq_lens=negative_txt_seq_lens, attention_kwargs=self.attention_kwargs, return_dict=False, )[0] diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_plus.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_plus.py index ec203edf166c..588783458571 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_plus.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_plus.py @@ -14,7 +14,7 @@ import inspect import math -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -84,10 +84,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -102,15 +102,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -143,7 +143,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -228,10 +228,10 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor def _get_qwen_prompt_embeds( self, - prompt: Union[str, List[str]] = None, - image: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt: str | list[str] = None, + image: torch.Tensor | None = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -286,18 +286,18 @@ def _get_qwen_prompt_embeds( # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - image: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, + prompt: str | list[str], + image: torch.Tensor | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, max_sequence_length: int = 1024, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded image (`torch.Tensor`, *optional*): image to be encoded @@ -323,6 +323,9 @@ def encode_prompt( prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + if prompt_embeds_mask is not None and prompt_embeds_mask.all(): + prompt_embeds_mask = None + return prompt_embeds, prompt_embeds_mask # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline.check_inputs @@ -516,43 +519,43 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - image: Optional[PipelineImageInput] = None, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + image: PipelineImageInput | None = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, true_cfg_scale: float = 4.0, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, - guidance_scale: Optional[float] = None, + sigmas: list[float] | None = None, + guidance_scale: float | None = None, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image latents as `image`, but if passing latents directly it is not encoded again. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is not greater than `1`). @@ -570,7 +573,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -586,7 +589,7 @@ def __call__( enable classifier-free guidance computations). num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -614,7 +617,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -663,6 +666,13 @@ def __call__( else: batch_size = prompt_embeds.shape[0] + # QwenImageEditPlusPipeline does not currently support batch_size > 1 + if batch_size > 1: + raise ValueError( + f"QwenImageEditPlusPipeline currently only supports batch_size=1, but received batch_size={batch_size}. " + "Please process prompts one at a time." + ) + device = self._execution_device # 3. Preprocess image if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels): @@ -777,11 +787,6 @@ def __call__( if self.attention_kwargs is None: self._attention_kwargs = {} - txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None - negative_txt_seq_lens = ( - negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None - ) - # 6. Denoising loop self.scheduler.set_begin_index(0) with self.progress_bar(total=num_inference_steps) as progress_bar: @@ -805,7 +810,6 @@ def __call__( encoder_hidden_states_mask=prompt_embeds_mask, encoder_hidden_states=prompt_embeds, img_shapes=img_shapes, - txt_seq_lens=txt_seq_lens, attention_kwargs=self.attention_kwargs, return_dict=False, )[0] @@ -820,7 +824,6 @@ def __call__( encoder_hidden_states_mask=negative_prompt_embeds_mask, encoder_hidden_states=negative_prompt_embeds, img_shapes=img_shapes, - txt_seq_lens=negative_txt_seq_lens, attention_kwargs=self.attention_kwargs, return_dict=False, )[0] diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_img2img.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_img2img.py index cb4c5d8016bb..42e63f8919a2 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_img2img.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_img2img.py @@ -1,5 +1,5 @@ import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -45,7 +45,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -74,10 +74,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -92,15 +92,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -194,9 +194,9 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._get_qwen_prompt_embeds def _get_qwen_prompt_embeds( self, - prompt: Union[str, List[str]] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt: str | list[str] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -265,20 +265,20 @@ def get_timesteps(self, num_inference_steps, strength, device): return timesteps, num_inference_steps - t_start - # Copied fromCopied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline.encode_prompt + # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, max_sequence_length: int = 1024, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -297,13 +297,17 @@ def encode_prompt( prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, device) prompt_embeds = prompt_embeds[:, :max_sequence_length] - prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length] - _, seq_len, _ = prompt_embeds.shape prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) - prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + + if prompt_embeds_mask is not None: + prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length] + prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) + prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + + if prompt_embeds_mask.all(): + prompt_embeds_mask = None return prompt_embeds, prompt_embeds_mask @@ -354,15 +358,6 @@ def check_inputs( f" {negative_prompt_embeds}. Please make sure to only forward one of the two." ) - if prompt_embeds is not None and prompt_embeds_mask is None: - raise ValueError( - "If `prompt_embeds` are provided, `prompt_embeds_mask` also have to be passed. Make sure to generate `prompt_embeds_mask` from the same text encoder that was used to generate `prompt_embeds`." - ) - if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None: - raise ValueError( - "If `negative_prompt_embeds` are provided, `negative_prompt_embeds_mask` also have to be passed. Make sure to generate `negative_prompt_embeds_mask` from the same text encoder that was used to generate `negative_prompt_embeds`." - ) - if max_sequence_length is not None and max_sequence_length > 1024: raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}") @@ -526,42 +521,42 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, true_cfg_scale: float = 4.0, image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, strength: float = 0.6, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, - guidance_scale: Optional[float] = None, + sigmas: list[float] | None = None, + guidance_scale: float | None = None, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is not greater than `1`). - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a @@ -587,7 +582,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -603,7 +598,7 @@ def __call__( enable classifier-free guidance computations). num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -631,7 +626,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -775,11 +770,6 @@ def __call__( if self.attention_kwargs is None: self._attention_kwargs = {} - txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None - negative_txt_seq_lens = ( - negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None - ) - # 6. Denoising loop with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): @@ -797,7 +787,6 @@ def __call__( encoder_hidden_states_mask=prompt_embeds_mask, encoder_hidden_states=prompt_embeds, img_shapes=img_shapes, - txt_seq_lens=txt_seq_lens, attention_kwargs=self.attention_kwargs, return_dict=False, )[0] @@ -811,7 +800,6 @@ def __call__( encoder_hidden_states_mask=negative_prompt_embeds_mask, encoder_hidden_states=negative_prompt_embeds, img_shapes=img_shapes, - txt_seq_lens=negative_txt_seq_lens, attention_kwargs=self.attention_kwargs, return_dict=False, )[0] diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_inpaint.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_inpaint.py index 1915c27eb2bb..5baf5bf5f77d 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_inpaint.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_inpaint.py @@ -1,5 +1,5 @@ import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -48,7 +48,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -77,10 +77,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -95,15 +95,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -204,9 +204,9 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._get_qwen_prompt_embeds def _get_qwen_prompt_embeds( self, - prompt: Union[str, List[str]] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt: str | list[str] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -276,20 +276,20 @@ def get_timesteps(self, num_inference_steps, strength, device): return timesteps, num_inference_steps - t_start - # Copied fromCopied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline.encode_prompt + # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, max_sequence_length: int = 1024, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -308,13 +308,17 @@ def encode_prompt( prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, device) prompt_embeds = prompt_embeds[:, :max_sequence_length] - prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length] - _, seq_len, _ = prompt_embeds.shape prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) - prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + + if prompt_embeds_mask is not None: + prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length] + prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) + prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + + if prompt_embeds_mask.all(): + prompt_embeds_mask = None return prompt_embeds, prompt_embeds_mask @@ -369,14 +373,6 @@ def check_inputs( f" {negative_prompt_embeds}. Please make sure to only forward one of the two." ) - if prompt_embeds is not None and prompt_embeds_mask is None: - raise ValueError( - "If `prompt_embeds` are provided, `prompt_embeds_mask` also have to be passed. Make sure to generate `prompt_embeds_mask` from the same text encoder that was used to generate `prompt_embeds`." - ) - if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None: - raise ValueError( - "If `negative_prompt_embeds` are provided, `negative_prompt_embeds_mask` also have to be passed. Make sure to generate `negative_prompt_embeds_mask` from the same text encoder that was used to generate `negative_prompt_embeds`." - ) if padding_mask_crop is not None: if not isinstance(image, PIL.Image.Image): raise ValueError( @@ -636,45 +632,45 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, true_cfg_scale: float = 4.0, image: PipelineImageInput = None, mask_image: PipelineImageInput = None, masked_image_latents: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, - padding_mask_crop: Optional[int] = None, + height: int | None = None, + width: int | None = None, + padding_mask_crop: int | None = None, strength: float = 0.6, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, - guidance_scale: Optional[float] = None, + sigmas: list[float] | None = None, + guidance_scale: float | None = None, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is not greater than `1`). - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a @@ -687,14 +683,14 @@ def __call__( setting `true_cfg_scale > 1` and a provided `negative_prompt`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one color channel (L) instead of 3, so the expected shape for pytorch tensor would be `(B, 1, H, W)`, `(B, H, W)`, `(1, H, W)`, `(H, W)`. And for numpy array would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W, 1)`, or `(H, W)`. - mask_image_latent (`torch.Tensor`, `List[torch.Tensor]`): + mask_image_latent (`torch.Tensor`, `list[torch.Tensor]`): `Tensor` representing an image batch to mask `image` generated by VAE. If not provided, the mask latents tensor will be generated by `mask_image`. height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): @@ -717,7 +713,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -733,7 +729,7 @@ def __call__( enable classifier-free guidance computations). num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -761,7 +757,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -944,11 +940,6 @@ def __call__( if self.attention_kwargs is None: self._attention_kwargs = {} - txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None - negative_txt_seq_lens = ( - negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None - ) - # 6. Denoising loop with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): @@ -966,7 +957,6 @@ def __call__( encoder_hidden_states_mask=prompt_embeds_mask, encoder_hidden_states=prompt_embeds, img_shapes=img_shapes, - txt_seq_lens=txt_seq_lens, attention_kwargs=self.attention_kwargs, return_dict=False, )[0] @@ -980,7 +970,6 @@ def __call__( encoder_hidden_states_mask=negative_prompt_embeds_mask, encoder_hidden_states=negative_prompt_embeds, img_shapes=img_shapes, - txt_seq_lens=negative_txt_seq_lens, attention_kwargs=self.attention_kwargs, return_dict=False, )[0] diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_layered.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_layered.py index 7bb12c26baa4..c7a44d880f9b 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_layered.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_layered.py @@ -14,7 +14,7 @@ import inspect import math -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -89,10 +89,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -107,15 +107,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -148,7 +148,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -250,9 +250,9 @@ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor def _get_qwen_prompt_embeds( self, - prompt: Union[str, List[str]] = None, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + prompt: str | list[str] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -291,17 +291,17 @@ def _get_qwen_prompt_embeds( # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, max_sequence_length: int = 1024, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -320,13 +320,17 @@ def encode_prompt( prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, device) prompt_embeds = prompt_embeds[:, :max_sequence_length] - prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length] - _, seq_len, _ = prompt_embeds.shape prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) - prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + + if prompt_embeds_mask is not None: + prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length] + prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) + prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + + if prompt_embeds_mask.all(): + prompt_embeds_mask = None return prompt_embeds, prompt_embeds_mask @@ -528,26 +532,26 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - image: Optional[PipelineImageInput] = None, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + image: PipelineImageInput | None = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, true_cfg_scale: float = 4.0, - layers: Optional[int] = 4, + layers: int | None = 4, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, - guidance_scale: Optional[float] = None, + sigmas: list[float] | None = None, + guidance_scale: float | None = None, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, resolution: int = 640, cfg_normalize: bool = False, @@ -557,16 +561,16 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image latents as `image`, but if passing latents directly it is not encoded again. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is not greater than `1`). @@ -580,7 +584,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -596,7 +600,7 @@ def __call__( enable classifier-free guidance computations). num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -781,10 +785,6 @@ def __call__( if self.attention_kwargs is None: self._attention_kwargs = {} - txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None - negative_txt_seq_lens = ( - negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None - ) is_rgb = torch.tensor([0] * batch_size).to(device=device, dtype=torch.long) # 6. Denoising loop self.scheduler.set_begin_index(0) @@ -809,7 +809,6 @@ def __call__( encoder_hidden_states_mask=prompt_embeds_mask, encoder_hidden_states=prompt_embeds, img_shapes=img_shapes, - txt_seq_lens=txt_seq_lens, attention_kwargs=self.attention_kwargs, additional_t_cond=is_rgb, return_dict=False, @@ -825,7 +824,6 @@ def __call__( encoder_hidden_states_mask=negative_prompt_embeds_mask, encoder_hidden_states=negative_prompt_embeds, img_shapes=img_shapes, - txt_seq_lens=negative_txt_seq_lens, attention_kwargs=self.attention_kwargs, additional_t_cond=is_rgb, return_dict=False, @@ -885,7 +883,7 @@ def __call__( latents = latents[:, :, 1:] # remove the first frame as it is the orgin input - latents = latents.permute(0, 2, 1, 3, 4).view(-1, c, 1, h, w) + latents = latents.permute(0, 2, 1, 3, 4).reshape(-1, c, 1, h, w) image = self.vae.decode(latents, return_dict=False)[0] # (b f) c 1 h w diff --git a/src/diffusers/pipelines/sana/pipeline_output.py b/src/diffusers/pipelines/sana/pipeline_output.py index f8ac12951644..b9d095906e06 100644 --- a/src/diffusers/pipelines/sana/pipeline_output.py +++ b/src/diffusers/pipelines/sana/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -13,9 +12,9 @@ class SanaPipelineOutput(BaseOutput): Output class for Sana pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image] | np.ndarray diff --git a/src/diffusers/pipelines/sana/pipeline_sana.py b/src/diffusers/pipelines/sana/pipeline_sana.py index 2beff802c6e0..17e0be9ba7aa 100644 --- a/src/diffusers/pipelines/sana/pipeline_sana.py +++ b/src/diffusers/pipelines/sana/pipeline_sana.py @@ -17,7 +17,7 @@ import re import urllib.parse as ul import warnings -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast @@ -130,10 +130,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -148,15 +148,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -201,7 +201,7 @@ class SanaPipeline(DiffusionPipeline, SanaLoraLoaderMixin): def __init__( self, - tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast], + tokenizer: GemmaTokenizer | GemmaTokenizerFast, text_encoder: Gemma2PreTrainedModel, vae: AutoencoderDC, transformer: SanaTransformer2DModel, @@ -275,18 +275,18 @@ def disable_vae_tiling(self): def _get_gemma_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], device: torch.device, dtype: torch.dtype, clean_caption: bool = False, max_sequence_length: int = 300, - complex_human_instruction: Optional[List[str]] = None, + complex_human_instruction: list[str] | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`, *optional*): torch device to place the resulting embeddings on @@ -333,27 +333,27 @@ def _get_gemma_prompt_embeds( def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], do_classifier_free_guidance: bool = True, negative_prompt: str = "", num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, + device: torch.device | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, clean_caption: bool = False, max_sequence_length: int = 300, - complex_human_instruction: Optional[List[str]] = None, - lora_scale: Optional[float] = None, + complex_human_instruction: list[str] | None = None, + lora_scale: float | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For PixArt-Alpha, this should be "". @@ -728,31 +728,31 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, negative_prompt: str = "", num_inference_steps: int = 20, - timesteps: List[int] = None, - sigmas: List[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, guidance_scale: float = 4.5, - num_images_per_prompt: Optional[int] = 1, + num_images_per_prompt: int | None = 1, height: int = 1024, width: int = 1024, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, clean_caption: bool = False, use_resolution_binning: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 300, - complex_human_instruction: List[str] = [ + complex_human_instruction: list[str] = [ "Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:", "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.", "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.", @@ -762,26 +762,26 @@ def __call__( "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", "User Prompt: ", ], - ) -> Union[SanaPipelineOutput, Tuple]: + ) -> SanaPipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_inference_steps (`int`, *optional*, defaults to 20): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -800,7 +800,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -838,13 +838,13 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. max_sequence_length (`int` defaults to `300`): Maximum sequence length to use with the `prompt`. - complex_human_instruction (`List[str]`, *optional*): + complex_human_instruction (`list[str]`, *optional*): Instructions for complex human attention: https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55. @@ -927,8 +927,12 @@ def __call__( prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) # 4. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) # 5. Prepare latents. diff --git a/src/diffusers/pipelines/sana/pipeline_sana_controlnet.py b/src/diffusers/pipelines/sana/pipeline_sana_controlnet.py index 55ed7b84ebdf..d976c7035d9d 100644 --- a/src/diffusers/pipelines/sana/pipeline_sana_controlnet.py +++ b/src/diffusers/pipelines/sana/pipeline_sana_controlnet.py @@ -17,7 +17,7 @@ import re import urllib.parse as ul import warnings -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast @@ -137,10 +137,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -155,15 +155,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -208,7 +208,7 @@ class SanaControlNetPipeline(DiffusionPipeline, SanaLoraLoaderMixin): def __init__( self, - tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast], + tokenizer: GemmaTokenizer | GemmaTokenizerFast, text_encoder: Gemma2PreTrainedModel, vae: AutoencoderDC, transformer: SanaTransformer2DModel, @@ -289,18 +289,18 @@ def disable_vae_tiling(self): # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._get_gemma_prompt_embeds def _get_gemma_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], device: torch.device, dtype: torch.dtype, clean_caption: bool = False, max_sequence_length: int = 300, - complex_human_instruction: Optional[List[str]] = None, + complex_human_instruction: list[str] | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`, *optional*): torch device to place the resulting embeddings on @@ -348,27 +348,27 @@ def _get_gemma_prompt_embeds( # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], do_classifier_free_guidance: bool = True, negative_prompt: str = "", num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, + device: torch.device | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, clean_caption: bool = False, max_sequence_length: int = 300, - complex_human_instruction: Optional[List[str]] = None, - lora_scale: Optional[float] = None, + complex_human_instruction: list[str] | None = None, + lora_scale: float | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For PixArt-Alpha, this should be "". @@ -777,33 +777,33 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, negative_prompt: str = "", num_inference_steps: int = 20, - timesteps: List[int] = None, - sigmas: List[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, guidance_scale: float = 4.5, control_image: PipelineImageInput = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, - num_images_per_prompt: Optional[int] = 1, + controlnet_conditioning_scale: float | list[float] = 1.0, + num_images_per_prompt: int | None = 1, height: int = 1024, width: int = 1024, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, clean_caption: bool = False, use_resolution_binning: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 300, - complex_human_instruction: List[str] = [ + complex_human_instruction: list[str] = [ "Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:", "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.", "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.", @@ -813,26 +813,26 @@ def __call__( "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", "User Prompt: ", ], - ) -> Union[SanaPipelineOutput, Tuple]: + ) -> SanaPipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_inference_steps (`int`, *optional*, defaults to 20): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -842,15 +842,15 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,: - `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`): + control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, `list[np.ndarray]`,: + `list[list[torch.Tensor]]`, `list[list[np.ndarray]]` or `list[list[PIL.Image.Image]]`): The ControlNet input condition to provide guidance to the `unet` for generation. If the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or width are passed, `image` is resized accordingly. If multiple ControlNets are specified in `init`, images must be passed as a list such that each element of the list can be correctly batched for input to a single ControlNet. - controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + controlnet_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set the corresponding scale as a list. @@ -863,7 +863,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -901,13 +901,13 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. max_sequence_length (`int` defaults to `300`): Maximum sequence length to use with the `prompt`. - complex_human_instruction (`List[str]`, *optional*): + complex_human_instruction (`list[str]`, *optional*): Instructions for complex human attention: https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55. @@ -1010,8 +1010,12 @@ def __call__( raise ValueError("`controlnet` must be of type `SanaControlNetModel`.") # 5. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) # 6. Prepare latents. diff --git a/src/diffusers/pipelines/sana/pipeline_sana_sprint.py b/src/diffusers/pipelines/sana/pipeline_sana_sprint.py index 04f45f817efb..c85f05275f51 100644 --- a/src/diffusers/pipelines/sana/pipeline_sana_sprint.py +++ b/src/diffusers/pipelines/sana/pipeline_sana_sprint.py @@ -17,7 +17,7 @@ import re import urllib.parse as ul import warnings -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast @@ -81,10 +81,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -99,15 +99,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -152,7 +152,7 @@ class SanaSprintPipeline(DiffusionPipeline, SanaLoraLoaderMixin): def __init__( self, - tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast], + tokenizer: GemmaTokenizer | GemmaTokenizerFast, text_encoder: Gemma2PreTrainedModel, vae: AutoencoderDC, transformer: SanaTransformer2DModel, @@ -227,18 +227,18 @@ def disable_vae_tiling(self): # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._get_gemma_prompt_embeds def _get_gemma_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], device: torch.device, dtype: torch.dtype, clean_caption: bool = False, max_sequence_length: int = 300, - complex_human_instruction: Optional[List[str]] = None, + complex_human_instruction: list[str] | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`, *optional*): torch device to place the resulting embeddings on @@ -285,21 +285,21 @@ def _get_gemma_prompt_embeds( def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, + device: torch.device | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, clean_caption: bool = False, max_sequence_length: int = 300, - complex_human_instruction: Optional[List[str]] = None, - lora_scale: Optional[float] = None, + complex_human_instruction: list[str] | None = None, + lora_scale: float | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -616,29 +616,29 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_inference_steps: int = 2, - timesteps: List[int] = None, + timesteps: list[int] = None, max_timesteps: float = 1.57080, intermediate_timesteps: float = 1.3, guidance_scale: float = 4.5, - num_images_per_prompt: Optional[int] = 1, + num_images_per_prompt: int | None = 1, height: int = 1024, width: int = 1024, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, clean_caption: bool = False, use_resolution_binning: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 300, - complex_human_instruction: List[str] = [ + complex_human_instruction: list[str] = [ "Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:", "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.", "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.", @@ -648,12 +648,12 @@ def __call__( "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", "User Prompt: ", ], - ) -> Union[SanaPipelineOutput, Tuple]: + ) -> SanaPipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. num_inference_steps (`int`, *optional*, defaults to 20): @@ -663,7 +663,7 @@ def __call__( The maximum timestep value used in the SCM scheduler. intermediate_timesteps (`float`, *optional*, defaults to 1.3): The intermediate timestep value used in SCM scheduler (only used when num_inference_steps=2). - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. @@ -682,7 +682,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -715,13 +715,13 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. max_sequence_length (`int` defaults to `300`): Maximum sequence length to use with the `prompt`. - complex_human_instruction (`List[str]`, *optional*): + complex_human_instruction (`list[str]`, *optional*): Instructions for complex human attention: https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55. @@ -790,10 +790,14 @@ def __call__( ) # 4. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, - device, + timestep_device, timesteps, sigmas=None, max_timesteps=max_timesteps, diff --git a/src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py b/src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py index 8899ed84c4e5..a17c494e88eb 100644 --- a/src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py +++ b/src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py @@ -17,7 +17,7 @@ import re import urllib.parse as ul import warnings -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch import torch.nn.functional as F @@ -87,10 +87,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -105,15 +105,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -159,7 +159,7 @@ class SanaSprintImg2ImgPipeline(DiffusionPipeline, SanaLoraLoaderMixin): def __init__( self, - tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast], + tokenizer: GemmaTokenizer | GemmaTokenizerFast, text_encoder: Gemma2PreTrainedModel, vae: AutoencoderDC, transformer: SanaTransformer2DModel, @@ -237,18 +237,18 @@ def disable_vae_tiling(self): # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._get_gemma_prompt_embeds def _get_gemma_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], device: torch.device, dtype: torch.dtype, clean_caption: bool = False, max_sequence_length: int = 300, - complex_human_instruction: Optional[List[str]] = None, + complex_human_instruction: list[str] | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`, *optional*): torch device to place the resulting embeddings on @@ -296,21 +296,21 @@ def _get_gemma_prompt_embeds( # Copied from diffusers.pipelines.sana.pipeline_sana_sprint.SanaSprintPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, + device: torch.device | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, clean_caption: bool = False, max_sequence_length: int = 300, - complex_human_instruction: Optional[List[str]] = None, - lora_scale: Optional[float] = None, + complex_human_instruction: list[str] | None = None, + lora_scale: float | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -687,31 +687,31 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_inference_steps: int = 2, - timesteps: List[int] = None, + timesteps: list[int] = None, max_timesteps: float = 1.57080, intermediate_timesteps: float = 1.3, guidance_scale: float = 4.5, image: PipelineImageInput = None, strength: float = 0.6, - num_images_per_prompt: Optional[int] = 1, + num_images_per_prompt: int | None = 1, height: int = 1024, width: int = 1024, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, clean_caption: bool = False, use_resolution_binning: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 300, - complex_human_instruction: List[str] = [ + complex_human_instruction: list[str] = [ "Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:", "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.", "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.", @@ -721,12 +721,12 @@ def __call__( "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", "User Prompt: ", ], - ) -> Union[SanaPipelineOutput, Tuple]: + ) -> SanaPipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. num_inference_steps (`int`, *optional*, defaults to 20): @@ -736,7 +736,7 @@ def __call__( The maximum timestep value used in the SCM scheduler. intermediate_timesteps (`float`, *optional*, defaults to 1.3): The intermediate timestep value used in SCM scheduler (only used when num_inference_steps=2). - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. @@ -755,7 +755,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -788,13 +788,13 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. max_sequence_length (`int` defaults to `300`): Maximum sequence length to use with the `prompt`. - complex_human_instruction (`List[str]`, *optional*): + complex_human_instruction (`list[str]`, *optional*): Instructions for complex human attention: https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55. diff --git a/src/diffusers/pipelines/sana_video/pipeline_output.py b/src/diffusers/pipelines/sana_video/pipeline_output.py index 4d37923889eb..a625ad3e1d34 100644 --- a/src/diffusers/pipelines/sana_video/pipeline_output.py +++ b/src/diffusers/pipelines/sana_video/pipeline_output.py @@ -11,7 +11,7 @@ class SanaVideoPipelineOutput(BaseOutput): Output class for Sana-Video pipelines. Args: - frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape `(batch_size, num_frames, channels, height, width)`. diff --git a/src/diffusers/pipelines/sana_video/pipeline_sana_video.py b/src/diffusers/pipelines/sana_video/pipeline_sana_video.py index a786275e45a9..8b44dfc1143c 100644 --- a/src/diffusers/pipelines/sana_video/pipeline_sana_video.py +++ b/src/diffusers/pipelines/sana_video/pipeline_sana_video.py @@ -17,7 +17,7 @@ import re import urllib.parse as ul import warnings -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast @@ -126,10 +126,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -144,15 +144,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -211,9 +211,9 @@ class SanaVideoPipeline(DiffusionPipeline, SanaLoraLoaderMixin): def __init__( self, - tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast], + tokenizer: GemmaTokenizer | GemmaTokenizerFast, text_encoder: Gemma2PreTrainedModel, - vae: Union[AutoencoderDC, AutoencoderKLWan], + vae: AutoencoderDC | AutoencoderKLWan, transformer: SanaVideoTransformer3DModel, scheduler: DPMSolverMultistepScheduler, ): @@ -233,18 +233,18 @@ def __init__( # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._get_gemma_prompt_embeds def _get_gemma_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], device: torch.device, dtype: torch.dtype, clean_caption: bool = False, max_sequence_length: int = 300, - complex_human_instruction: Optional[List[str]] = None, + complex_human_instruction: list[str] | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`, *optional*): torch device to place the resulting embeddings on @@ -291,27 +291,27 @@ def _get_gemma_prompt_embeds( def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], do_classifier_free_guidance: bool = True, negative_prompt: str = "", num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, + device: torch.device | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, clean_caption: bool = False, max_sequence_length: int = 300, - complex_human_instruction: Optional[List[str]] = None, - lora_scale: Optional[float] = None, + complex_human_instruction: list[str] | None = None, + lora_scale: float | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide the video generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For PixArt-Alpha, this should be "". @@ -650,10 +650,10 @@ def prepare_latents( height: int = 480, width: int = 832, num_frames: int = 81, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: if latents is not None: return latents.to(device=device, dtype=dtype) @@ -702,32 +702,32 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, negative_prompt: str = "", num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, guidance_scale: float = 6.0, - num_videos_per_prompt: Optional[int] = 1, + num_videos_per_prompt: int | None = 1, height: int = 480, width: int = 832, frames: int = 81, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, clean_caption: bool = False, use_resolution_binning: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 300, - complex_human_instruction: List[str] = [ + complex_human_instruction: list[str] = [ "Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for video generation. Evaluate the level of detail in the user prompt:", "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, motion, and temporal relationships to create vivid and dynamic scenes.", "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.", @@ -737,26 +737,26 @@ def __call__( "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", "User Prompt: ", ], - ) -> Union[SanaVideoPipelineOutput, Tuple]: + ) -> SanaVideoPipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the video generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality video at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -777,7 +777,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -820,7 +820,7 @@ def __call__( `._callback_tensor_inputs` attribute of your pipeline class. max_sequence_length (`int` defaults to `300`): Maximum sequence length to use with the `prompt`. - complex_human_instruction (`List[str]`, *optional*): + complex_human_instruction (`list[str]`, *optional*): Instructions for complex human attention: https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55. diff --git a/src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py b/src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py index e87880b64cee..b90d7c6f5a60 100644 --- a/src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py +++ b/src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py @@ -17,7 +17,7 @@ import re import urllib.parse as ul import warnings -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import PIL import torch @@ -102,10 +102,10 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -120,15 +120,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -161,7 +161,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -201,9 +201,9 @@ class SanaImageToVideoPipeline(DiffusionPipeline, SanaLoraLoaderMixin): def __init__( self, - tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast], + tokenizer: GemmaTokenizer | GemmaTokenizerFast, text_encoder: Gemma2PreTrainedModel, - vae: Union[AutoencoderDC, AutoencoderKLWan], + vae: AutoencoderDC | AutoencoderKLWan, transformer: SanaVideoTransformer3DModel, scheduler: FlowMatchEulerDiscreteScheduler, ): @@ -230,18 +230,18 @@ def __init__( # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._get_gemma_prompt_embeds def _get_gemma_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], device: torch.device, dtype: torch.dtype, clean_caption: bool = False, max_sequence_length: int = 300, - complex_human_instruction: Optional[List[str]] = None, + complex_human_instruction: list[str] | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`, *optional*): torch device to place the resulting embeddings on @@ -289,27 +289,27 @@ def _get_gemma_prompt_embeds( # Copied from diffusers.pipelines.sana_video.pipeline_sana_video.SanaVideoPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], + prompt: str | list[str], do_classifier_free_guidance: bool = True, negative_prompt: str = "", num_videos_per_prompt: int = 1, - device: Optional[torch.device] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, + device: torch.device | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, clean_caption: bool = False, max_sequence_length: int = 300, - complex_human_instruction: Optional[List[str]] = None, - lora_scale: Optional[float] = None, + complex_human_instruction: list[str] | None = None, + lora_scale: float | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide the video generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For PixArt-Alpha, this should be "". @@ -653,10 +653,10 @@ def prepare_latents( height: int = 480, width: int = 832, num_frames: int = 81, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 shape = ( @@ -726,32 +726,32 @@ def interrupt(self): def __call__( self, image: PipelineImageInput, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, negative_prompt: str = "", num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, guidance_scale: float = 6.0, - num_videos_per_prompt: Optional[int] = 1, + num_videos_per_prompt: int | None = 1, height: int = 480, width: int = 832, frames: int = 81, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_attention_mask: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_attention_mask: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, clean_caption: bool = False, use_resolution_binning: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 300, - complex_human_instruction: List[str] = [ + complex_human_instruction: list[str] = [ "Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for video generation. Evaluate the level of detail in the user prompt:", "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, motion, and temporal relationships to create vivid and dynamic scenes.", "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.", @@ -761,7 +761,7 @@ def __call__( "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", "User Prompt: ", ], - ) -> Union[SanaVideoPipelineOutput, Tuple]: + ) -> SanaVideoPipelineOutput | tuple: """ Function invoked when calling the pipeline for generation. @@ -769,21 +769,21 @@ def __call__( image (`PipelineImageInput`): The input image to condition the video generation on. The first frame of the generated video will be conditioned on this image. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the video generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality video at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -804,7 +804,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -847,7 +847,7 @@ def __call__( `._callback_tensor_inputs` attribute of your pipeline class. max_sequence_length (`int` defaults to `300`): Maximum sequence length to use with the `prompt`. - complex_human_instruction (`List[str]`, *optional*): + complex_human_instruction (`list[str]`, *optional*): Instructions for complex human attention: https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55. diff --git a/src/diffusers/pipelines/semantic_stable_diffusion/pipeline_output.py b/src/diffusers/pipelines/semantic_stable_diffusion/pipeline_output.py index 349912993981..8e5429ce2a8d 100644 --- a/src/diffusers/pipelines/semantic_stable_diffusion/pipeline_output.py +++ b/src/diffusers/pipelines/semantic_stable_diffusion/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Optional, Union import numpy as np import PIL.Image @@ -13,13 +12,13 @@ class SemanticStableDiffusionPipelineOutput(BaseOutput): Output class for Stable Diffusion pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, num_channels)`. - nsfw_content_detected (`List[bool]`) - List indicating whether the corresponding generated image contains “not-safe-for-work” (nsfw) content or + nsfw_content_detected (`list[bool]`) + list indicating whether the corresponding generated image contains “not-safe-for-work” (nsfw) content or `None` if safety checking could not be performed. """ - images: Union[List[PIL.Image.Image], np.ndarray] - nsfw_content_detected: Optional[List[bool]] + images: list[PIL.Image.Image] | np.ndarray + nsfw_content_detected: list[bool] | None diff --git a/src/diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py b/src/diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py index 49b09e205cc5..05d28896e117 100644 --- a/src/diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py +++ b/src/diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py @@ -1,6 +1,6 @@ import inspect from itertools import repeat -from typing import Callable, List, Optional, Union +from typing import Callable import torch from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer @@ -144,7 +144,6 @@ def prepare_extra_step_kwargs(self, generator, eta): extra_step_kwargs["generator"] = generator return extra_step_kwargs - # Copied from diffusers.pipelines.stable_diffusion_k_diffusion.pipeline_stable_diffusion_k_diffusion.StableDiffusionKDiffusionPipeline.check_inputs def check_inputs( self, prompt, @@ -223,37 +222,37 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype @torch.no_grad() def __call__( self, - prompt: Union[str, List[str]], - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str], + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, + negative_prompt: str | list[str] | None = None, num_images_per_prompt: int = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - editing_prompt: Optional[Union[str, List[str]]] = None, - editing_prompt_embeddings: Optional[torch.Tensor] = None, - reverse_editing_direction: Optional[Union[bool, List[bool]]] = False, - edit_guidance_scale: Optional[Union[float, List[float]]] = 5, - edit_warmup_steps: Optional[Union[int, List[int]]] = 10, - edit_cooldown_steps: Optional[Union[int, List[int]]] = None, - edit_threshold: Optional[Union[float, List[float]]] = 0.9, - edit_momentum_scale: Optional[float] = 0.1, - edit_mom_beta: Optional[float] = 0.4, - edit_weights: Optional[List[float]] = None, - sem_guidance: Optional[List[torch.Tensor]] = None, + editing_prompt: str | list[str] | None = None, + editing_prompt_embeddings: torch.Tensor | None = None, + reverse_editing_direction: bool | list[bool] | None = False, + edit_guidance_scale: float | list[float] | None = 5, + edit_warmup_steps: int | list[int] | None = 10, + edit_cooldown_steps: int | list[int] | None = None, + edit_threshold: float | list[float] | None = 0.9, + edit_momentum_scale: float | None = 0.1, + edit_mom_beta: float | None = 0.4, + edit_weights: list[float] | None = None, + sem_guidance: list[torch.Tensor] | None = None, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide image generation. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -265,7 +264,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -273,7 +272,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -291,24 +290,24 @@ def __call__( callback_steps (`int`, *optional*, defaults to 1): The frequency at which the `callback` function is called. If not specified, the callback is called at every step. - editing_prompt (`str` or `List[str]`, *optional*): + editing_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to use for semantic guidance. Semantic guidance is disabled by setting `editing_prompt = None`. Guidance direction of prompt should be specified via `reverse_editing_direction`. editing_prompt_embeddings (`torch.Tensor`, *optional*): Pre-computed embeddings to use for semantic guidance. Guidance direction of embedding should be specified via `reverse_editing_direction`. - reverse_editing_direction (`bool` or `List[bool]`, *optional*, defaults to `False`): + reverse_editing_direction (`bool` or `list[bool]`, *optional*, defaults to `False`): Whether the corresponding prompt in `editing_prompt` should be increased or decreased. - edit_guidance_scale (`float` or `List[float]`, *optional*, defaults to 5): + edit_guidance_scale (`float` or `list[float]`, *optional*, defaults to 5): Guidance scale for semantic guidance. If provided as a list, values should correspond to `editing_prompt`. - edit_warmup_steps (`float` or `List[float]`, *optional*, defaults to 10): + edit_warmup_steps (`float` or `list[float]`, *optional*, defaults to 10): Number of diffusion steps (for each prompt) for which semantic guidance is not applied. Momentum is calculated for those steps and applied once all warmup periods are over. - edit_cooldown_steps (`float` or `List[float]`, *optional*, defaults to `None`): + edit_cooldown_steps (`float` or `list[float]`, *optional*, defaults to `None`): Number of diffusion steps (for each prompt) after which semantic guidance is longer applied. - edit_threshold (`float` or `List[float]`, *optional*, defaults to 0.9): + edit_threshold (`float` or `list[float]`, *optional*, defaults to 0.9): Threshold of semantic guidance. edit_momentum_scale (`float`, *optional*, defaults to 0.1): Scale of the momentum to be added to the semantic guidance at each diffusion step. If set to 0.0, @@ -318,11 +317,11 @@ def __call__( Defines how semantic guidance momentum builds up. `edit_mom_beta` indicates how much of the previous momentum is kept. Momentum is already built up during warmup (for diffusion steps smaller than `edit_warmup_steps`). - edit_weights (`List[float]`, *optional*, defaults to `None`): + edit_weights (`list[float]`, *optional*, defaults to `None`): Indicates how much each individual concept should influence the overall guidance. If no weights are provided all concepts are applied equally. - sem_guidance (`List[torch.Tensor]`, *optional*): - List of pre-generated guidance vectors to be applied at generation. Length of the list has to + sem_guidance (`list[torch.Tensor]`, *optional*): + list of pre-generated guidance vectors to be applied at generation. Length of the list has to correspond to `num_inference_steps`. Examples: @@ -458,7 +457,7 @@ def __call__( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): diff --git a/src/diffusers/pipelines/shap_e/camera.py b/src/diffusers/pipelines/shap_e/camera.py index 31e1759d6154..81807b6ff35a 100644 --- a/src/diffusers/pipelines/shap_e/camera.py +++ b/src/diffusers/pipelines/shap_e/camera.py @@ -13,7 +13,6 @@ # limitations under the License. from dataclasses import dataclass -from typing import Tuple import numpy as np import torch @@ -33,7 +32,7 @@ class DifferentiableProjectiveCamera: height: int x_fov: float y_fov: float - shape: Tuple[int] + shape: tuple[int] def __post_init__(self): assert self.x.shape[0] == self.y.shape[0] == self.z.shape[0] == self.origin.shape[0] diff --git a/src/diffusers/pipelines/shap_e/pipeline_shap_e.py b/src/diffusers/pipelines/shap_e/pipeline_shap_e.py index 49ddfd1196bf..44967dfb3349 100644 --- a/src/diffusers/pipelines/shap_e/pipeline_shap_e.py +++ b/src/diffusers/pipelines/shap_e/pipeline_shap_e.py @@ -14,7 +14,6 @@ import math from dataclasses import dataclass -from typing import List, Optional, Union import numpy as np import PIL.Image @@ -82,7 +81,7 @@ class ShapEPipelineOutput(BaseOutput): A list of images for 3D rendering. """ - images: Union[List[List[PIL.Image.Image]], List[List[np.ndarray]]] + images: list[list[PIL.Image.Image]] | list[list[np.ndarray]] class ShapEPipeline(DiffusionPipeline): @@ -195,25 +194,25 @@ def __call__( prompt: str, num_images_per_prompt: int = 1, num_inference_steps: int = 25, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, guidance_scale: float = 4.0, frame_size: int = 64, - output_type: Optional[str] = "pil", # pil, np, latent, mesh + output_type: str | None = "pil", # pil, np, latent, mesh return_dict: bool = True, ): """ The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. num_inference_steps (`int`, *optional*, defaults to 25): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py b/src/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py index 55d8b85822c4..964db30e7de2 100644 --- a/src/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py +++ b/src/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py @@ -13,7 +13,6 @@ # limitations under the License. from dataclasses import dataclass -from typing import List, Optional, Union import numpy as np import PIL.Image @@ -83,7 +82,7 @@ class ShapEPipelineOutput(BaseOutput): A list of images for 3D rendering. """ - images: Union[PIL.Image.Image, np.ndarray] + images: PIL.Image.Image | np.ndarray class ShapEImg2ImgPipeline(DiffusionPipeline): @@ -147,7 +146,7 @@ def _encode_image( num_images_per_prompt, do_classifier_free_guidance, ): - if isinstance(image, List) and isinstance(image[0], torch.Tensor): + if isinstance(image, list) and isinstance(image[0], torch.Tensor): image = torch.cat(image, axis=0) if image[0].ndim == 4 else torch.stack(image, axis=0) if not isinstance(image, torch.Tensor): @@ -174,21 +173,21 @@ def _encode_image( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - image: Union[PIL.Image.Image, List[PIL.Image.Image]], + image: PIL.Image.Image | list[PIL.Image.Image], num_images_per_prompt: int = 1, num_inference_steps: int = 25, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, guidance_scale: float = 4.0, frame_size: int = 64, - output_type: Optional[str] = "pil", # pil, np, latent, mesh + output_type: str | None = "pil", # pil, np, latent, mesh return_dict: bool = True, ): """ The call function to the pipeline for generation. Args: - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image` or tensor representing an image batch to be used as the starting point. Can also accept image latents as image, but if passing latents directly it is not encoded again. num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -196,7 +195,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 25): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -231,7 +230,7 @@ def __call__( batch_size = len(image) else: raise ValueError( - f"`image` has to be of type `PIL.Image.Image`, `torch.Tensor`, `List[PIL.Image.Image]` or `List[torch.Tensor]` but is {type(image)}" + f"`image` has to be of type `PIL.Image.Image`, `torch.Tensor`, `list[PIL.Image.Image]` or `list[torch.Tensor]` but is {type(image)}" ) device = self._execution_device diff --git a/src/diffusers/pipelines/shap_e/renderer.py b/src/diffusers/pipelines/shap_e/renderer.py index d1d05c894595..0c2058c887fc 100644 --- a/src/diffusers/pipelines/shap_e/renderer.py +++ b/src/diffusers/pipelines/shap_e/renderer.py @@ -14,7 +14,6 @@ import math from dataclasses import dataclass -from typing import Dict, Optional, Tuple import numpy as np import torch @@ -131,7 +130,7 @@ def _convert_srgb_to_linear(u: torch.Tensor): def _create_flat_edge_indices( flat_cube_indices: torch.Tensor, - grid_size: Tuple[int, int, int], + grid_size: tuple[int, int, int], ): num_xs = (grid_size[0] - 1) * grid_size[1] * grid_size[2] y_offset = num_xs @@ -301,7 +300,7 @@ def intersect( self, origin: torch.Tensor, direction: torch.Tensor, - t0_lower: Optional[torch.Tensor] = None, + t0_lower: torch.Tensor | None = None, epsilon=1e-6, ): """ @@ -479,7 +478,7 @@ class MeshDecoderOutput(BaseOutput): verts: torch.Tensor faces: torch.Tensor - vertex_channels: Dict[str, torch.Tensor] + vertex_channels: dict[str, torch.Tensor] class MeshDecoder(nn.Module): @@ -742,13 +741,13 @@ class ShapEParamsProjModel(ModelMixin, ConfigMixin): def __init__( self, *, - param_names: Tuple[str, ...] = ( + param_names: tuple[str] = ( "nerstf.mlp.0.weight", "nerstf.mlp.1.weight", "nerstf.mlp.2.weight", "nerstf.mlp.3.weight", ), - param_shapes: Tuple[Tuple[int]] = ( + param_shapes: tuple[tuple[int]] = ( (256, 93), (256, 256), (256, 256), @@ -786,13 +785,13 @@ class ShapERenderer(ModelMixin, ConfigMixin): def __init__( self, *, - param_names: Tuple[str, ...] = ( + param_names: tuple[str] = ( "nerstf.mlp.0.weight", "nerstf.mlp.1.weight", "nerstf.mlp.2.weight", "nerstf.mlp.3.weight", ), - param_shapes: Tuple[Tuple[int, int], ...] = ( + param_shapes: tuple[tuple[int]] = ( (256, 93), (256, 256), (256, 256), @@ -804,7 +803,7 @@ def __init__( n_hidden_layers: int = 6, act_fn: str = "swish", insert_direction_at: int = 4, - background: Tuple[float, ...] = ( + background: tuple[float] = ( 255.0, 255.0, 255.0, @@ -953,7 +952,7 @@ def decode_to_mesh( device, grid_size: int = 128, query_batch_size: int = 4096, - texture_channels: Tuple = ("R", "G", "B"), + texture_channels: tuple = ("R", "G", "B"), ): # 1. project the parameters from the generated latents projected_params = self.params_proj(latents) diff --git a/src/diffusers/pipelines/skyreels_v2/pipeline_output.py b/src/diffusers/pipelines/skyreels_v2/pipeline_output.py index 7a170d24c39a..dac2316362ec 100644 --- a/src/diffusers/pipelines/skyreels_v2/pipeline_output.py +++ b/src/diffusers/pipelines/skyreels_v2/pipeline_output.py @@ -11,8 +11,8 @@ class SkyReelsV2PipelineOutput(BaseOutput): Output class for SkyReelsV2 pipelines. Args: - frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): - List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): + list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape `(batch_size, num_frames, channels, height, width)`. """ diff --git a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2.py b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2.py index d6cd7d7feceb..c92608fad3b6 100644 --- a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2.py +++ b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2.py @@ -13,11 +13,11 @@ # limitations under the License. import html -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import regex as re import torch -from transformers import AutoTokenizer, UMT5EncoderModel +from transformers import AutoTokenizer, T5EncoderModel, UMT5EncoderModel from ...callbacks import MultiPipelineCallbacks, PipelineCallback from ...loaders import SkyReelsV2LoraLoaderMixin @@ -132,7 +132,7 @@ class SkyReelsV2Pipeline(DiffusionPipeline, SkyReelsV2LoraLoaderMixin): def __init__( self, tokenizer: AutoTokenizer, - text_encoder: UMT5EncoderModel, + text_encoder: T5EncoderModel | UMT5EncoderModel, transformer: SkyReelsV2Transformer3DModel, vae: AutoencoderKLWan, scheduler: UniPCMultistepScheduler, @@ -154,11 +154,11 @@ def __init__( # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -196,23 +196,23 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -324,10 +324,10 @@ def prepare_latents( height: int = 480, width: int = 832, num_frames: int = 81, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: if latents is not None: return latents.to(device=device, dtype=dtype) @@ -377,32 +377,30 @@ def attention_kwargs(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, height: int = 544, width: int = 960, num_frames: int = 97, num_inference_steps: int = 50, guidance_scale: float = 6.0, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "np", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "np", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. height (`int`, defaults to `544`): @@ -422,7 +420,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -445,7 +443,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -545,22 +543,24 @@ def __call__( latent_model_input = latents.to(transformer_dtype) timestep = t.expand(latents.shape[0]) - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - - if self.do_classifier_free_guidance: - noise_uncond = self.transformer( + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( hidden_states=latent_model_input, timestep=timestep, - encoder_hidden_states=negative_prompt_embeds, + encoder_hidden_states=prompt_embeds, attention_kwargs=attention_kwargs, return_dict=False, )[0] + + if self.do_classifier_free_guidance: + with self.transformer.cache_context("uncond"): + noise_uncond = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=negative_prompt_embeds, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond) # compute the previous noisy sample x_t -> x_t-1 diff --git a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing.py b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing.py index 089f92632d38..8751240a1af9 100644 --- a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing.py +++ b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing.py @@ -16,11 +16,10 @@ import math import re from copy import deepcopy -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable -import ftfy import torch -from transformers import AutoTokenizer, UMT5EncoderModel +from transformers import AutoTokenizer, T5EncoderModel, UMT5EncoderModel from ...callbacks import MultiPipelineCallbacks, PipelineCallback from ...loaders import SkyReelsV2LoraLoaderMixin @@ -114,7 +113,7 @@ def prompt_clean(text): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -154,7 +153,7 @@ class SkyReelsV2DiffusionForcingPipeline(DiffusionPipeline, SkyReelsV2LoraLoader def __init__( self, tokenizer: AutoTokenizer, - text_encoder: UMT5EncoderModel, + text_encoder: T5EncoderModel | UMT5EncoderModel, transformer: SkyReelsV2Transformer3DModel, vae: AutoencoderKLWan, scheduler: UniPCMultistepScheduler, @@ -176,11 +175,11 @@ def __init__( # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -218,23 +217,23 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -354,15 +353,15 @@ def prepare_latents( height: int = 480, width: int = 832, num_frames: int = 97, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - base_latent_num_frames: Optional[int] = None, - video_latents: Optional[torch.Tensor] = None, - causal_block_size: Optional[int] = None, - overlap_history_latent_frames: Optional[int] = None, - long_video_iter: Optional[int] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + base_latent_num_frames: int | None = None, + video_latents: torch.Tensor | None = None, + causal_block_size: int | None = None, + overlap_history_latent_frames: int | None = None, + long_video_iter: int | None = None, ) -> torch.Tensor: if latents is not None: return latents.to(device=device, dtype=dtype) @@ -462,7 +461,7 @@ def generate_timestep_matrix( num_latent_frames] - step_update_mask (torch.Tensor): Boolean mask indicating which frames to update Shape: [num_iterations, num_latent_frames] - - valid_interval (list[tuple]): List of (start, end) intervals for each iteration + - valid_interval (list[tuple]): list of (start, end) intervals for each iteration Raises: ValueError: If ar_step is too small for the given configuration @@ -599,41 +598,39 @@ def attention_kwargs(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]], - negative_prompt: Union[str, List[str]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] = None, height: int = 544, width: int = 960, num_frames: int = 97, num_inference_steps: int = 50, guidance_scale: float = 6.0, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "np", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "np", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, - overlap_history: Optional[int] = None, + overlap_history: int | None = None, addnoise_condition: float = 0, base_num_frames: int = 97, ar_step: int = 0, - causal_block_size: Optional[int] = None, + causal_block_size: int | None = None, fps: int = 24, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -654,7 +651,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. (**6.0 for T2V**, **5.0 for I2V**) num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -680,7 +677,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -887,25 +884,28 @@ def __call__( ) timestep[:, valid_interval_start:prefix_video_latents_frames] = addnoise_condition - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - enable_diffusion_forcing=True, - fps=fps_embeds, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - if self.do_classifier_free_guidance: - noise_uncond = self.transformer( + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( hidden_states=latent_model_input, timestep=timestep, - encoder_hidden_states=negative_prompt_embeds, + encoder_hidden_states=prompt_embeds, enable_diffusion_forcing=True, fps=fps_embeds, attention_kwargs=attention_kwargs, return_dict=False, )[0] + + if self.do_classifier_free_guidance: + with self.transformer.cache_context("uncond"): + noise_uncond = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=negative_prompt_embeds, + enable_diffusion_forcing=True, + fps=fps_embeds, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond) update_mask_i = step_update_mask[i] diff --git a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_i2v.py b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_i2v.py index 2951a9447386..a8f1b3a84a4a 100644 --- a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_i2v.py +++ b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_i2v.py @@ -16,12 +16,11 @@ import math import re from copy import deepcopy -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable -import ftfy import PIL import torch -from transformers import AutoTokenizer, UMT5EncoderModel +from transformers import AutoTokenizer, T5EncoderModel, UMT5EncoderModel from diffusers.image_processor import PipelineImageInput from diffusers.utils.torch_utils import randn_tensor @@ -119,7 +118,7 @@ def prompt_clean(text): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -159,7 +158,7 @@ class SkyReelsV2DiffusionForcingImageToVideoPipeline(DiffusionPipeline, SkyReels def __init__( self, tokenizer: AutoTokenizer, - text_encoder: UMT5EncoderModel, + text_encoder: T5EncoderModel | UMT5EncoderModel, transformer: SkyReelsV2Transformer3DModel, vae: AutoencoderKLWan, scheduler: UniPCMultistepScheduler, @@ -181,11 +180,11 @@ def __init__( # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -223,23 +222,23 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -367,23 +366,23 @@ def check_inputs( def prepare_latents( self, - image: Optional[PipelineImageInput], + image: PipelineImageInput | None, batch_size: int, num_channels_latents: int = 16, height: int = 480, width: int = 832, num_frames: int = 97, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - last_image: Optional[torch.Tensor] = None, - video_latents: Optional[torch.Tensor] = None, - base_latent_num_frames: Optional[int] = None, - causal_block_size: Optional[int] = None, - overlap_history_latent_frames: Optional[int] = None, - long_video_iter: Optional[int] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + last_image: torch.Tensor | None = None, + video_latents: torch.Tensor | None = None, + base_latent_num_frames: int | None = None, + causal_block_size: int | None = None, + overlap_history_latent_frames: int | None = None, + long_video_iter: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 latent_height = height // self.vae_scale_factor_spatial latent_width = width // self.vae_scale_factor_spatial @@ -508,7 +507,7 @@ def generate_timestep_matrix( num_latent_frames] - step_update_mask (torch.Tensor): Boolean mask indicating which frames to update Shape: [num_iterations, num_latent_frames] - - valid_interval (list[tuple]): List of (start, end) intervals for each iteration + - valid_interval (list[tuple]): list of (start, end) intervals for each iteration Raises: ValueError: If ar_step is too small for the given configuration @@ -646,33 +645,31 @@ def attention_kwargs(self): def __call__( self, image: PipelineImageInput, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, height: int = 544, width: int = 960, num_frames: int = 97, num_inference_steps: int = 50, guidance_scale: float = 5.0, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - image_embeds: Optional[torch.Tensor] = None, - last_image: Optional[torch.Tensor] = None, - output_type: Optional[str] = "np", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + image_embeds: torch.Tensor | None = None, + last_image: torch.Tensor | None = None, + output_type: str | None = "np", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, - overlap_history: Optional[int] = None, + overlap_history: int | None = None, addnoise_condition: float = 0, base_num_frames: int = 97, ar_step: int = 0, - causal_block_size: Optional[int] = None, + causal_block_size: int | None = None, fps: int = 24, ): r""" @@ -681,10 +678,10 @@ def __call__( Args: image (`PipelineImageInput`): The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -705,7 +702,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. (**6.0 for T2V**, **5.0 for I2V**) num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -737,7 +734,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -966,25 +963,28 @@ def __call__( ) timestep[:, valid_interval_start:prefix_video_latents_frames] = addnoise_condition - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - enable_diffusion_forcing=True, - fps=fps_embeds, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - if self.do_classifier_free_guidance: - noise_uncond = self.transformer( + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( hidden_states=latent_model_input, timestep=timestep, - encoder_hidden_states=negative_prompt_embeds, + encoder_hidden_states=prompt_embeds, enable_diffusion_forcing=True, fps=fps_embeds, attention_kwargs=attention_kwargs, return_dict=False, )[0] + + if self.do_classifier_free_guidance: + with self.transformer.cache_context("uncond"): + noise_uncond = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=negative_prompt_embeds, + enable_diffusion_forcing=True, + fps=fps_embeds, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond) update_mask_i = step_update_mask[i] diff --git a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_v2v.py b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_v2v.py index 6fedfc795a40..924acb850d09 100644 --- a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_v2v.py +++ b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_v2v.py @@ -17,12 +17,11 @@ import math import re from copy import deepcopy -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable -import ftfy import torch from PIL import Image -from transformers import AutoTokenizer, UMT5EncoderModel +from transformers import AutoTokenizer, T5EncoderModel, UMT5EncoderModel from ...callbacks import MultiPipelineCallbacks, PipelineCallback from ...loaders import SkyReelsV2LoraLoaderMixin @@ -116,10 +115,10 @@ def prompt_clean(text): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -134,15 +133,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -175,7 +174,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -215,7 +214,7 @@ class SkyReelsV2DiffusionForcingVideoToVideoPipeline(DiffusionPipeline, SkyReels def __init__( self, tokenizer: AutoTokenizer, - text_encoder: UMT5EncoderModel, + text_encoder: T5EncoderModel | UMT5EncoderModel, transformer: SkyReelsV2Transformer3DModel, vae: AutoencoderKLWan, scheduler: UniPCMultistepScheduler, @@ -237,11 +236,11 @@ def __init__( # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -279,23 +278,23 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -421,16 +420,16 @@ def prepare_latents( height: int = 480, width: int = 832, num_frames: int = 97, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - video_latents: Optional[torch.Tensor] = None, - base_latent_num_frames: Optional[int] = None, - overlap_history: Optional[int] = None, - causal_block_size: Optional[int] = None, - overlap_history_latent_frames: Optional[int] = None, - long_video_iter: Optional[int] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + video_latents: torch.Tensor | None = None, + base_latent_num_frames: int | None = None, + overlap_history: int | None = None, + causal_block_size: int | None = None, + overlap_history_latent_frames: int | None = None, + long_video_iter: int | None = None, ) -> torch.Tensor: if latents is not None: return latents.to(device=device, dtype=dtype) @@ -546,7 +545,7 @@ def generate_timestep_matrix( num_latent_frames] - step_update_mask (torch.Tensor): Boolean mask indicating which frames to update Shape: [num_iterations, num_latent_frames] - - valid_interval (list[tuple]): List of (start, end) intervals for each iteration + - valid_interval (list[tuple]): list of (start, end) intervals for each iteration Raises: ValueError: If ar_step is too small for the given configuration @@ -683,44 +682,42 @@ def attention_kwargs(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - video: List[Image.Image], - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + video: list[Image.Image], + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, height: int = 544, width: int = 960, num_frames: int = 120, num_inference_steps: int = 50, guidance_scale: float = 6.0, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "np", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "np", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, - overlap_history: Optional[int] = None, + overlap_history: int | None = None, addnoise_condition: float = 0, base_num_frames: int = 97, ar_step: int = 0, - causal_block_size: Optional[int] = None, + causal_block_size: int | None = None, fps: int = 24, ): r""" The call function to the pipeline for generation. Args: - video (`List[Image.Image]`): + video (`list[Image.Image]`): The video to guide the video generation. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the video generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -741,7 +738,7 @@ def __call__( usually at the expense of lower image quality. (**6.0 for T2V**, **5.0 for I2V**) num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -767,7 +764,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -974,25 +971,28 @@ def __call__( ) timestep[:, valid_interval_start:prefix_video_latents_frames] = addnoise_condition - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - enable_diffusion_forcing=True, - fps=fps_embeds, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - if self.do_classifier_free_guidance: - noise_uncond = self.transformer( + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( hidden_states=latent_model_input, timestep=timestep, - encoder_hidden_states=negative_prompt_embeds, + encoder_hidden_states=prompt_embeds, enable_diffusion_forcing=True, fps=fps_embeds, attention_kwargs=attention_kwargs, return_dict=False, )[0] + + if self.do_classifier_free_guidance: + with self.transformer.cache_context("uncond"): + noise_uncond = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=negative_prompt_embeds, + enable_diffusion_forcing=True, + fps=fps_embeds, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond) update_mask_i = step_update_mask[i] diff --git a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_i2v.py b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_i2v.py index d61b687eadc3..7c24b898e0bb 100644 --- a/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_i2v.py +++ b/src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_i2v.py @@ -13,12 +13,12 @@ # limitations under the License. import html -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import PIL import regex as re import torch -from transformers import AutoTokenizer, CLIPProcessor, CLIPVisionModelWithProjection, UMT5EncoderModel +from transformers import AutoTokenizer, CLIPProcessor, CLIPVisionModelWithProjection, T5EncoderModel, UMT5EncoderModel from ...callbacks import MultiPipelineCallbacks, PipelineCallback from ...image_processor import PipelineImageInput @@ -112,7 +112,7 @@ def prompt_clean(text): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -157,7 +157,7 @@ class SkyReelsV2ImageToVideoPipeline(DiffusionPipeline, SkyReelsV2LoraLoaderMixi def __init__( self, tokenizer: AutoTokenizer, - text_encoder: UMT5EncoderModel, + text_encoder: T5EncoderModel | UMT5EncoderModel, image_encoder: CLIPVisionModelWithProjection, image_processor: CLIPProcessor, transformer: SkyReelsV2Transformer3DModel, @@ -184,11 +184,11 @@ def __init__( # Copied from diffusers.pipelines.wan.pipeline_wan_i2v.WanImageToVideoPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -227,7 +227,7 @@ def _get_t5_prompt_embeds( def encode_image( self, image: PipelineImageInput, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): device = device or self._execution_device image = self.image_processor(images=image, return_tensors="pt").to(device) @@ -237,23 +237,23 @@ def encode_image( # Copied from diffusers.pipelines.wan.pipeline_wan_i2v.WanImageToVideoPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -378,12 +378,12 @@ def prepare_latents( height: int = 480, width: int = 832, num_frames: int = 81, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - last_image: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + last_image: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 latent_height = height // self.vae_scale_factor_spatial latent_width = width // self.vae_scale_factor_spatial @@ -478,27 +478,25 @@ def attention_kwargs(self): def __call__( self, image: PipelineImageInput, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, height: int = 544, width: int = 960, num_frames: int = 97, num_inference_steps: int = 50, guidance_scale: float = 5.0, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - image_embeds: Optional[torch.Tensor] = None, - last_image: Optional[torch.Tensor] = None, - output_type: Optional[str] = "np", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + image_embeds: torch.Tensor | None = None, + last_image: torch.Tensor | None = None, + output_type: str | None = "np", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" @@ -507,10 +505,10 @@ def __call__( Args: image (`PipelineImageInput`): The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -531,7 +529,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -560,7 +558,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -678,24 +676,26 @@ def __call__( latent_model_input = torch.cat([latents, condition], dim=1).to(transformer_dtype) timestep = t.expand(latents.shape[0]) - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - encoder_hidden_states_image=image_embeds, - attention_kwargs=attention_kwargs, - return_dict=False, - )[0] - - if self.do_classifier_free_guidance: - noise_uncond = self.transformer( + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( hidden_states=latent_model_input, timestep=timestep, - encoder_hidden_states=negative_prompt_embeds, + encoder_hidden_states=prompt_embeds, encoder_hidden_states_image=image_embeds, attention_kwargs=attention_kwargs, return_dict=False, )[0] + + if self.do_classifier_free_guidance: + with self.transformer.cache_context("uncond"): + noise_uncond = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=negative_prompt_embeds, + encoder_hidden_states_image=image_embeds, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond) # compute the previous noisy sample x_t -> x_t-1 diff --git a/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py b/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py index 07b382dfc49f..d40269411bc0 100644 --- a/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py +++ b/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py @@ -14,7 +14,6 @@ from dataclasses import dataclass from math import pi -from typing import Optional import torch import torch.nn as nn @@ -57,9 +56,9 @@ class StableAudioProjectionModelOutput(BaseOutput): Sequence of hidden-states obtained by linearly projecting the audio end hidden states. """ - text_hidden_states: Optional[torch.Tensor] = None - seconds_start_hidden_states: Optional[torch.Tensor] = None - seconds_end_hidden_states: Optional[torch.Tensor] = None + text_hidden_states: torch.Tensor | None = None + seconds_start_hidden_states: torch.Tensor | None = None + seconds_end_hidden_states: torch.Tensor | None = None class StableAudioNumberConditioner(nn.Module): @@ -82,7 +81,7 @@ def __init__( number_embedding_dim, min_value, max_value, - internal_dim: Optional[int] = 256, + internal_dim: int | None = 256, ): super().__init__() self.time_positional_embedding = nn.Sequential( @@ -138,9 +137,9 @@ def __init__(self, text_encoder_dim, conditioning_dim, min_value, max_value): def forward( self, - text_hidden_states: Optional[torch.Tensor] = None, - start_seconds: Optional[torch.Tensor] = None, - end_seconds: Optional[torch.Tensor] = None, + text_hidden_states: torch.Tensor | None = None, + start_seconds: torch.Tensor | None = None, + end_seconds: torch.Tensor | None = None, ): text_hidden_states = ( text_hidden_states if text_hidden_states is None else self.text_projection(text_hidden_states) diff --git a/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py b/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py index b7faf097ab0d..351c8b65de0e 100644 --- a/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py +++ b/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Callable, List, Optional, Union +from typing import Callable import torch from transformers import ( @@ -108,7 +108,7 @@ def __init__( vae: AutoencoderOobleck, text_encoder: T5EncoderModel, projection_model: StableAudioProjectionModel, - tokenizer: Union[T5Tokenizer, T5TokenizerFast], + tokenizer: T5Tokenizer | T5TokenizerFast, transformer: StableAudioDiTModel, scheduler: EDMDPMSolverMultistepScheduler, ): @@ -158,10 +158,10 @@ def encode_prompt( device, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.LongTensor] = None, - negative_attention_mask: Optional[torch.LongTensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + attention_mask: torch.LongTensor | None = None, + negative_attention_mask: torch.LongTensor | None = None, ): if prompt is not None and isinstance(prompt, str): batch_size = 1 @@ -206,7 +206,7 @@ def encode_prompt( prompt_embeds = prompt_embeds[0] if do_classifier_free_guidance and negative_prompt is not None: - uncond_tokens: List[str] + uncond_tokens: list[str] if type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -491,32 +491,32 @@ def prepare_latents( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - audio_end_in_s: Optional[float] = None, - audio_start_in_s: Optional[float] = 0.0, + prompt: str | list[str] = None, + audio_end_in_s: float | None = None, + audio_start_in_s: float | None = 0.0, num_inference_steps: int = 100, guidance_scale: float = 7.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_waveforms_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_waveforms_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - initial_audio_waveforms: Optional[torch.Tensor] = None, - initial_audio_sampling_rate: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.LongTensor] = None, - negative_attention_mask: Optional[torch.LongTensor] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + initial_audio_waveforms: torch.Tensor | None = None, + initial_audio_sampling_rate: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + attention_mask: torch.LongTensor | None = None, + negative_attention_mask: torch.LongTensor | None = None, return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: Optional[int] = 1, - output_type: Optional[str] = "pt", + callback: Callable[[int, int, torch.Tensor], None] | None = None, + callback_steps: int | None = 1, + output_type: str | None = "pt", ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide audio generation. If not defined, you need to pass `prompt_embeds`. audio_end_in_s (`float`, *optional*, defaults to 47.55): Audio end index in seconds. @@ -528,7 +528,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.0): A higher guidance scale value encourages the model to generate audio that is closely linked to the text `prompt` at the expense of lower sound quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in audio generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_waveforms_per_prompt (`int`, *optional*, defaults to 1): @@ -536,7 +536,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py index a6a60ad94be6..ef40078bfbb9 100644 --- a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py +++ b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import torch from transformers import CLIPTextModelWithProjection, CLIPTokenizer @@ -139,10 +139,10 @@ def encode_prompt( do_classifier_free_guidance, prompt=None, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_pooled: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_pooled: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_pooled: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_pooled: torch.Tensor | None = None, ): if prompt_embeds is None: # get prompt text embeddings @@ -184,7 +184,7 @@ def encode_prompt( prompt_embeds_pooled = prompt_embeds_pooled.repeat_interleave(num_images_per_prompt, dim=0) if negative_prompt_embeds is None and do_classifier_free_guidance: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -305,30 +305,30 @@ def get_timestep_ratio_conditioning(self, t, alphas_cumprod): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - image_embeddings: Union[torch.Tensor, List[torch.Tensor]], - prompt: Union[str, List[str]] = None, + image_embeddings: torch.Tensor | list[torch.Tensor], + prompt: str | list[str] = None, num_inference_steps: int = 10, guidance_scale: float = 0.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_pooled: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_pooled: Optional[torch.Tensor] = None, + negative_prompt: str | list[str] | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_pooled: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_pooled: torch.Tensor | None = None, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], ): """ Function invoked when calling the pipeline for generation. Args: - image_embedding (`torch.Tensor` or `List[torch.Tensor]`): + image_embedding (`torch.Tensor` or `list[torch.Tensor]`): Image Embeddings either extracted from an image or generated by a Prior Model. - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. num_inference_steps (`int`, *optional*, defaults to 12): The number of denoising steps. More denoising steps usually lead to a higher quality image at the @@ -339,7 +339,7 @@ def __call__( equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `decoder_guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `decoder_guidance_scale` is less than `1`). prompt_embeds (`torch.Tensor`, *optional*): @@ -358,7 +358,7 @@ def __call__( input argument. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -375,7 +375,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py index 838b93faaa0c..50e6c02b6017 100644 --- a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py +++ b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import PIL import torch @@ -90,8 +90,8 @@ def __init__( prior_text_encoder: CLIPTextModelWithProjection, prior_tokenizer: CLIPTokenizer, prior_scheduler: DDPMWuerstchenScheduler, - prior_feature_extractor: Optional[CLIPImageProcessor] = None, - prior_image_encoder: Optional[CLIPVisionModelWithProjection] = None, + prior_feature_extractor: CLIPImageProcessor | None = None, + prior_image_encoder: CLIPVisionModelWithProjection | None = None, ): super().__init__() @@ -124,10 +124,10 @@ def __init__( vqgan=vqgan, ) - def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None): + def enable_xformers_memory_efficient_attention(self, attention_op: Callable | None = None): self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op) - def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None): + def enable_model_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None): r""" Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward` @@ -137,7 +137,7 @@ def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[t self.prior_pipe.enable_model_cpu_offload(gpu_id=gpu_id, device=device) self.decoder_pipe.enable_model_cpu_offload(gpu_id=gpu_id, device=device) - def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None): + def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None): r""" Offloads all models (`unet`, `text_encoder`, `vae`, and `safety checker` state dicts) to CPU using 🤗 Accelerate, significantly reducing memory usage. Models are moved to a `torch.device('meta')` and loaded on a @@ -159,38 +159,38 @@ def set_progress_bar_config(self, **kwargs): @replace_example_docstring(TEXT2IMAGE_EXAMPLE_DOC_STRING) def __call__( self, - prompt: Optional[Union[str, List[str]]] = None, - images: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]] = None, + prompt: str | list[str] | None = None, + images: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image] = None, height: int = 512, width: int = 512, prior_num_inference_steps: int = 60, prior_guidance_scale: float = 4.0, num_inference_steps: int = 12, decoder_guidance_scale: float = 0.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_pooled: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_pooled: Optional[torch.Tensor] = None, + negative_prompt: str | list[str] | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_pooled: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_pooled: torch.Tensor | None = None, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - prior_callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - prior_callback_on_step_end_tensor_inputs: List[str] = ["latents"], - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + prior_callback_on_step_end: Callable[[int, int], None] | None = None, + prior_callback_on_step_end_tensor_inputs: list[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], ): """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation for the prior and decoder. - images (`torch.Tensor`, `PIL.Image.Image`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, *optional*): + images (`torch.Tensor`, `PIL.Image.Image`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, *optional*): The images to guide the image generation for the prior. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). prompt_embeds (`torch.Tensor`, *optional*): @@ -219,7 +219,7 @@ def __call__( equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `prior_guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - prior_num_inference_steps (`Union[int, Dict[float, int]]`, *optional*, defaults to 60): + prior_num_inference_steps (`int | dict[float, int]`, *optional*, defaults to 60): The number of prior denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. For more specific timestep spacing, you can pass customized `prior_timesteps` @@ -233,7 +233,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -249,7 +249,7 @@ def __call__( A function that calls at the end of each denoising steps during the inference. The function is called with the following arguments: `prior_callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. - prior_callback_on_step_end_tensor_inputs (`List`, *optional*): + prior_callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `prior_callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -258,7 +258,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py index 29ad8b5429d7..0c5ea9ed61b4 100644 --- a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py +++ b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py @@ -14,7 +14,7 @@ from dataclasses import dataclass from math import ceil -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import numpy as np import PIL @@ -70,11 +70,11 @@ class StableCascadePriorPipelineOutput(BaseOutput): Text embeddings for the negative prompt. """ - image_embeddings: Union[torch.Tensor, np.ndarray] - prompt_embeds: Union[torch.Tensor, np.ndarray] - prompt_embeds_pooled: Union[torch.Tensor, np.ndarray] - negative_prompt_embeds: Union[torch.Tensor, np.ndarray] - negative_prompt_embeds_pooled: Union[torch.Tensor, np.ndarray] + image_embeddings: torch.Tensor | np.ndarray + prompt_embeds: torch.Tensor | np.ndarray + prompt_embeds_pooled: torch.Tensor | np.ndarray + negative_prompt_embeds: torch.Tensor | np.ndarray + negative_prompt_embeds_pooled: torch.Tensor | np.ndarray class StableCascadePriorPipeline(DeprecatedPipelineMixin, DiffusionPipeline): @@ -118,8 +118,8 @@ def __init__( prior: StableCascadeUNet, scheduler: DDPMWuerstchenScheduler, resolution_multiple: float = 42.67, - feature_extractor: Optional[CLIPImageProcessor] = None, - image_encoder: Optional[CLIPVisionModelWithProjection] = None, + feature_extractor: CLIPImageProcessor | None = None, + image_encoder: CLIPVisionModelWithProjection | None = None, ) -> None: super().__init__() self.register_modules( @@ -160,10 +160,10 @@ def encode_prompt( do_classifier_free_guidance, prompt=None, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_pooled: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_pooled: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_pooled: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_pooled: torch.Tensor | None = None, ): if prompt_embeds is None: # get prompt text embeddings @@ -205,7 +205,7 @@ def encode_prompt( prompt_embeds_pooled = prompt_embeds_pooled.repeat_interleave(num_images_per_prompt, dim=0) if negative_prompt_embeds is None and do_classifier_free_guidance: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -376,32 +376,32 @@ def get_timestep_ratio_conditioning(self, t, alphas_cumprod): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Optional[Union[str, List[str]]] = None, - images: Union[torch.Tensor, PIL.Image.Image, List[torch.Tensor], List[PIL.Image.Image]] = None, + prompt: str | list[str] | None = None, + images: torch.Tensor | PIL.Image.Image | list[torch.Tensor] | list[PIL.Image.Image] = None, height: int = 1024, width: int = 1024, num_inference_steps: int = 20, - timesteps: List[float] = None, + timesteps: list[float] = None, guidance_scale: float = 4.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - prompt_embeds_pooled: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds_pooled: Optional[torch.Tensor] = None, - image_embeds: Optional[torch.Tensor] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pt", + negative_prompt: str | list[str] | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_pooled: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_pooled: torch.Tensor | None = None, + image_embeds: torch.Tensor | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pt", return_dict: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], ): """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. height (`int`, *optional*, defaults to 1024): The height in pixels of the generated image. @@ -416,7 +416,7 @@ def __call__( equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `decoder_guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `decoder_guidance_scale` is less than `1`). prompt_embeds (`torch.Tensor`, *optional*): @@ -438,7 +438,7 @@ def __call__( not provided, image embeddings will be generated from `image` input argument if existing. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -455,7 +455,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/stable_diffusion/__init__.py b/src/diffusers/pipelines/stable_diffusion/__init__.py index b05a3ce2a7c9..c2eebf586ef8 100644 --- a/src/diffusers/pipelines/stable_diffusion/__init__.py +++ b/src/diffusers/pipelines/stable_diffusion/__init__.py @@ -6,8 +6,6 @@ _LazyModule, get_objects_from_module, is_flax_available, - is_k_diffusion_available, - is_k_diffusion_version, is_onnx_available, is_torch_available, is_transformers_available, diff --git a/src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py b/src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py index 6c0221d2092a..8fa4bd5941d2 100644 --- a/src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py +++ b/src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py @@ -17,7 +17,6 @@ import re from contextlib import nullcontext from io import BytesIO -from typing import Dict, Optional, Union import requests import torch @@ -1049,7 +1048,7 @@ def stable_unclip_image_encoder(original_config, local_files_only=False): def stable_unclip_image_noising_components( - original_config, clip_stats_path: Optional[str] = None, device: Optional[str] = None + original_config, clip_stats_path: str | None = None, device: str | None = None ): """ Returns the noising components for the img2img and txt2img unclip pipelines. @@ -1144,25 +1143,25 @@ def convert_controlnet_checkpoint( def download_from_original_stable_diffusion_ckpt( - checkpoint_path_or_dict: Union[str, Dict[str, torch.Tensor]], + checkpoint_path_or_dict: str | dict[str, torch.Tensor], original_config_file: str = None, - image_size: Optional[int] = None, + image_size: int | None = None, prediction_type: str = None, model_type: str = None, extract_ema: bool = False, scheduler_type: str = "pndm", - num_in_channels: Optional[int] = None, - upcast_attention: Optional[bool] = None, + num_in_channels: int | None = None, + upcast_attention: bool | None = None, device: str = None, from_safetensors: bool = False, - stable_unclip: Optional[str] = None, - stable_unclip_prior: Optional[str] = None, - clip_stats_path: Optional[str] = None, - controlnet: Optional[bool] = None, - adapter: Optional[bool] = None, + stable_unclip: str | None = None, + stable_unclip_prior: str | None = None, + clip_stats_path: str | None = None, + controlnet: bool | None = None, + adapter: bool | None = None, load_safety_checker: bool = True, - safety_checker: Optional[StableDiffusionSafetyChecker] = None, - feature_extractor: Optional[AutoFeatureExtractor] = None, + safety_checker: StableDiffusionSafetyChecker | None = None, + feature_extractor: AutoFeatureExtractor | None = None, pipeline_class: DiffusionPipeline = None, local_files_only=False, vae_path=None, @@ -1237,7 +1236,7 @@ def download_from_original_stable_diffusion_ckpt( [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer) to use. If this parameter is `None`, the function will load a new instance of [CLIPTokenizer] by itself, if needed. - config_files (`Dict[str, str]`, *optional*, defaults to `None`): + config_files (`dict[str, str]`, *optional*, defaults to `None`): A dictionary mapping from config file names to their contents. If this parameter is `None`, the function will load the config files by itself, if needed. Valid keys are: - `v1`: Config file for Stable Diffusion v1 @@ -1827,12 +1826,12 @@ def download_controlnet_from_original_ckpt( original_config_file: str, image_size: int = 512, extract_ema: bool = False, - num_in_channels: Optional[int] = None, - upcast_attention: Optional[bool] = None, + num_in_channels: int | None = None, + upcast_attention: bool | None = None, device: str = None, from_safetensors: bool = False, - use_linear_projection: Optional[bool] = None, - cross_attention_dim: Optional[bool] = None, + use_linear_projection: bool | None = None, + cross_attention_dim: bool | None = None, ) -> DiffusionPipeline: if from_safetensors: from safetensors import safe_open diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py index 6befe77aa4b1..c15163917a7c 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py @@ -14,7 +14,6 @@ import warnings from functools import partial -from typing import Dict, List, Optional, Union import jax import jax.numpy as jnp @@ -112,9 +111,7 @@ def __init__( text_encoder: FlaxCLIPTextModel, tokenizer: CLIPTokenizer, unet: FlaxUNet2DConditionModel, - scheduler: Union[ - FlaxDDIMScheduler, FlaxPNDMScheduler, FlaxLMSDiscreteScheduler, FlaxDPMSolverMultistepScheduler - ], + scheduler: FlaxDDIMScheduler | FlaxPNDMScheduler | FlaxLMSDiscreteScheduler | FlaxDPMSolverMultistepScheduler, safety_checker: FlaxStableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, dtype: jnp.dtype = jnp.float32, @@ -168,7 +165,7 @@ def __init__( ) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - def prepare_inputs(self, prompt: Union[str, List[str]]): + def prepare_inputs(self, prompt: str | list[str]): if not isinstance(prompt, (str, list)): raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") @@ -218,14 +215,14 @@ def _run_safety_checker(self, images, safety_model_params, jit=False): def _generate( self, prompt_ids: jnp.array, - params: Union[Dict, FrozenDict], + params: dict | FrozenDict, prng_seed: jax.Array, num_inference_steps: int, height: int, width: int, guidance_scale: float, - latents: Optional[jnp.ndarray] = None, - neg_prompt_ids: Optional[jnp.ndarray] = None, + latents: jnp.ndarray | None = None, + neg_prompt_ids: jnp.ndarray | None = None, ): if height % 8 != 0 or width % 8 != 0: raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") @@ -315,12 +312,12 @@ def loop_body(step, args): def __call__( self, prompt_ids: jnp.array, - params: Union[Dict, FrozenDict], + params: dict | FrozenDict, prng_seed: jax.Array, num_inference_steps: int = 50, - height: Optional[int] = None, - width: Optional[int] = None, - guidance_scale: Union[float, jnp.ndarray] = 7.5, + height: int | None = None, + width: int | None = None, + guidance_scale: float | jnp.ndarray = 7.5, latents: jnp.ndarray = None, neg_prompt_ids: jnp.ndarray = None, return_dict: bool = True, @@ -330,7 +327,7 @@ def __call__( The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py index 81656beba7e1..e1c004ceb44c 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py @@ -14,7 +14,6 @@ import warnings from functools import partial -from typing import Dict, List, Optional, Union import jax import jax.numpy as jnp @@ -136,9 +135,7 @@ def __init__( text_encoder: FlaxCLIPTextModel, tokenizer: CLIPTokenizer, unet: FlaxUNet2DConditionModel, - scheduler: Union[ - FlaxDDIMScheduler, FlaxPNDMScheduler, FlaxLMSDiscreteScheduler, FlaxDPMSolverMultistepScheduler - ], + scheduler: FlaxDDIMScheduler | FlaxPNDMScheduler | FlaxLMSDiscreteScheduler | FlaxDPMSolverMultistepScheduler, safety_checker: FlaxStableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, dtype: jnp.dtype = jnp.float32, @@ -167,7 +164,7 @@ def __init__( ) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - def prepare_inputs(self, prompt: Union[str, List[str]], image: Union[Image.Image, List[Image.Image]]): + def prepare_inputs(self, prompt: str | list[str], image: Image.Image | list[Image.Image]): if not isinstance(prompt, (str, list)): raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") @@ -234,15 +231,15 @@ def _generate( self, prompt_ids: jnp.ndarray, image: jnp.ndarray, - params: Union[Dict, FrozenDict], + params: dict | FrozenDict, prng_seed: jax.Array, start_timestep: int, num_inference_steps: int, height: int, width: int, guidance_scale: float, - noise: Optional[jnp.ndarray] = None, - neg_prompt_ids: Optional[jnp.ndarray] = None, + noise: jnp.ndarray | None = None, + neg_prompt_ids: jnp.ndarray | None = None, ): if height % 8 != 0 or width % 8 != 0: raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") @@ -339,13 +336,13 @@ def __call__( self, prompt_ids: jnp.ndarray, image: jnp.ndarray, - params: Union[Dict, FrozenDict], + params: dict | FrozenDict, prng_seed: jax.Array, strength: float = 0.8, num_inference_steps: int = 50, - height: Optional[int] = None, - width: Optional[int] = None, - guidance_scale: Union[float, jnp.ndarray] = 7.5, + height: int | None = None, + width: int | None = None, + guidance_scale: float | jnp.ndarray = 7.5, noise: jnp.ndarray = None, neg_prompt_ids: jnp.ndarray = None, return_dict: bool = True, diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py index 5938fe232a71..dbda93deebff 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py @@ -14,7 +14,6 @@ import warnings from functools import partial -from typing import Dict, List, Optional, Union import jax import jax.numpy as jnp @@ -135,9 +134,7 @@ def __init__( text_encoder: FlaxCLIPTextModel, tokenizer: CLIPTokenizer, unet: FlaxUNet2DConditionModel, - scheduler: Union[ - FlaxDDIMScheduler, FlaxPNDMScheduler, FlaxLMSDiscreteScheduler, FlaxDPMSolverMultistepScheduler - ], + scheduler: FlaxDDIMScheduler | FlaxPNDMScheduler | FlaxLMSDiscreteScheduler | FlaxDPMSolverMultistepScheduler, safety_checker: FlaxStableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, dtype: jnp.dtype = jnp.float32, @@ -193,9 +190,9 @@ def __init__( def prepare_inputs( self, - prompt: Union[str, List[str]], - image: Union[Image.Image, List[Image.Image]], - mask: Union[Image.Image, List[Image.Image]], + prompt: str | list[str], + image: Image.Image | list[Image.Image], + mask: Image.Image | list[Image.Image], ): if not isinstance(prompt, (str, list)): raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") @@ -269,14 +266,14 @@ def _generate( prompt_ids: jnp.ndarray, mask: jnp.ndarray, masked_image: jnp.ndarray, - params: Union[Dict, FrozenDict], + params: dict | FrozenDict, prng_seed: jax.Array, num_inference_steps: int, height: int, width: int, guidance_scale: float, - latents: Optional[jnp.ndarray] = None, - neg_prompt_ids: Optional[jnp.ndarray] = None, + latents: jnp.ndarray | None = None, + neg_prompt_ids: jnp.ndarray | None = None, ): if height % 8 != 0 or width % 8 != 0: raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") @@ -397,12 +394,12 @@ def __call__( prompt_ids: jnp.ndarray, mask: jnp.ndarray, masked_image: jnp.ndarray, - params: Union[Dict, FrozenDict], + params: dict | FrozenDict, prng_seed: jax.Array, num_inference_steps: int = 50, - height: Optional[int] = None, - width: Optional[int] = None, - guidance_scale: Union[float, jnp.ndarray] = 7.5, + height: int | None = None, + width: int | None = None, + guidance_scale: float | jnp.ndarray = 7.5, latents: jnp.ndarray = None, neg_prompt_ids: jnp.ndarray = None, return_dict: bool = True, @@ -412,7 +409,7 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide image generation. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py index 6ebe0986a1ab..0f66ca909e7d 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Callable, List, Optional, Union +from typing import Callable import numpy as np import torch @@ -36,7 +36,7 @@ class OnnxStableDiffusionPipeline(DiffusionPipeline): text_encoder: OnnxRuntimeModel tokenizer: CLIPTokenizer unet: OnnxRuntimeModel - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler] + scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler safety_checker: OnnxRuntimeModel feature_extractor: CLIPImageProcessor @@ -50,7 +50,7 @@ def __init__( text_encoder: OnnxRuntimeModel, tokenizer: CLIPTokenizer, unet: OnnxRuntimeModel, - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], + scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler, safety_checker: OnnxRuntimeModel, feature_extractor: CLIPImageProcessor, requires_safety_checker: bool = True, @@ -114,24 +114,24 @@ def __init__( def _encode_prompt( self, - prompt: Union[str, List[str]], - num_images_per_prompt: Optional[int], + prompt: str | list[str], + num_images_per_prompt: int | None, do_classifier_free_guidance: bool, - negative_prompt: Optional[str], - prompt_embeds: Optional[np.ndarray] = None, - negative_prompt_embeds: Optional[np.ndarray] = None, + negative_prompt: str | None, + prompt_embeds: np.ndarray | None = None, + negative_prompt_embeds: np.ndarray | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): prompt to be encoded num_images_per_prompt (`int`): number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`): + negative_prompt (`str` or `list[str]`): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). prompt_embeds (`np.ndarray`, *optional*): @@ -176,7 +176,7 @@ def _encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -217,13 +217,13 @@ def _encode_prompt( def check_inputs( self, - prompt: Union[str, List[str]], - height: Optional[int], - width: Optional[int], + prompt: str | list[str], + height: int | None, + width: int | None, callback_steps: int, - negative_prompt: Optional[str] = None, - prompt_embeds: Optional[np.ndarray] = None, - negative_prompt_embeds: Optional[np.ndarray] = None, + negative_prompt: str | None = None, + prompt_embeds: np.ndarray | None = None, + negative_prompt_embeds: np.ndarray | None = None, ): if height % 8 != 0 or width % 8 != 0: raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") @@ -264,31 +264,31 @@ def check_inputs( def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = 512, - width: Optional[int] = 512, - num_inference_steps: Optional[int] = 50, - guidance_scale: Optional[float] = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - eta: Optional[float] = 0.0, - generator: Optional[np.random.RandomState] = None, - latents: Optional[np.ndarray] = None, - prompt_embeds: Optional[np.ndarray] = None, - negative_prompt_embeds: Optional[np.ndarray] = None, - output_type: Optional[str] = "pil", + prompt: str | list[str] = None, + height: int | None = 512, + width: int | None = 512, + num_inference_steps: int | None = 50, + guidance_scale: float | None = 7.5, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + eta: float | None = 0.0, + generator: np.random.RandomState | None = None, + latents: np.ndarray | None = None, + prompt_embeds: np.ndarray | None = None, + negative_prompt_embeds: np.ndarray | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, np.ndarray], None]] = None, + callback: Callable[[int, int, np.ndarray], None] | None = None, callback_steps: int = 1, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - image (`PIL.Image.Image` or List[`PIL.Image.Image`] or `torch.Tensor`): + image (`PIL.Image.Image` or list[`PIL.Image.Image`] or `torch.Tensor`): `Image`, or tensor representing an image batch which will be upscaled. * num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the @@ -299,7 +299,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`. instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -470,7 +470,7 @@ def __init__( text_encoder: OnnxRuntimeModel, tokenizer: CLIPTokenizer, unet: OnnxRuntimeModel, - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], + scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler, safety_checker: OnnxRuntimeModel, feature_extractor: CLIPImageProcessor, ): diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_img2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_img2img.py index d63bf3bf4564..2bb3eb6990e1 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_img2img.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_img2img.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Callable, List, Optional, Union +from typing import Callable import numpy as np import PIL.Image @@ -89,7 +89,7 @@ class OnnxStableDiffusionImg2ImgPipeline(DiffusionPipeline): text_encoder: OnnxRuntimeModel tokenizer: CLIPTokenizer unet: OnnxRuntimeModel - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler] + scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler safety_checker: OnnxRuntimeModel feature_extractor: CLIPImageProcessor @@ -103,7 +103,7 @@ def __init__( text_encoder: OnnxRuntimeModel, tokenizer: CLIPTokenizer, unet: OnnxRuntimeModel, - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], + scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler, safety_checker: OnnxRuntimeModel, feature_extractor: CLIPImageProcessor, requires_safety_checker: bool = True, @@ -168,24 +168,24 @@ def __init__( # Copied from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion.OnnxStableDiffusionPipeline._encode_prompt def _encode_prompt( self, - prompt: Union[str, List[str]], - num_images_per_prompt: Optional[int], + prompt: str | list[str], + num_images_per_prompt: int | None, do_classifier_free_guidance: bool, - negative_prompt: Optional[str], - prompt_embeds: Optional[np.ndarray] = None, - negative_prompt_embeds: Optional[np.ndarray] = None, + negative_prompt: str | None, + prompt_embeds: np.ndarray | None = None, + negative_prompt_embeds: np.ndarray | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): prompt to be encoded num_images_per_prompt (`int`): number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`): + negative_prompt (`str` or `list[str]`): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). prompt_embeds (`np.ndarray`, *optional*): @@ -230,7 +230,7 @@ def _encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -271,11 +271,11 @@ def _encode_prompt( def check_inputs( self, - prompt: Union[str, List[str]], + prompt: str | list[str], callback_steps: int, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[np.ndarray] = None, - negative_prompt_embeds: Optional[np.ndarray] = None, + negative_prompt: str | list[str] | None = None, + prompt_embeds: np.ndarray | None = None, + negative_prompt_embeds: np.ndarray | None = None, ): if (callback_steps is None) or ( callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0) @@ -313,27 +313,27 @@ def check_inputs( def __call__( self, - prompt: Union[str, List[str]], - image: Union[np.ndarray, PIL.Image.Image] = None, + prompt: str | list[str], + image: np.ndarray | PIL.Image.Image = None, strength: float = 0.8, - num_inference_steps: Optional[int] = 50, - guidance_scale: Optional[float] = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - eta: Optional[float] = 0.0, - generator: Optional[np.random.RandomState] = None, - prompt_embeds: Optional[np.ndarray] = None, - negative_prompt_embeds: Optional[np.ndarray] = None, - output_type: Optional[str] = "pil", + num_inference_steps: int | None = 50, + guidance_scale: float | None = 7.5, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + eta: float | None = 0.0, + generator: np.random.RandomState | None = None, + prompt_embeds: np.ndarray | None = None, + negative_prompt_embeds: np.ndarray | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, np.ndarray], None]] = None, + callback: Callable[[int, int, np.ndarray], None] | None = None, callback_steps: int = 1, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. image (`np.ndarray` or `PIL.Image.Image`): `Image`, or tensor representing an image batch, that will be used as the starting point for the @@ -353,7 +353,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_inpaint.py b/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_inpaint.py index 158bcabbebfd..7e3416a33832 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_inpaint.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_inpaint.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Callable, List, Optional, Union +from typing import Callable import numpy as np import PIL.Image @@ -87,7 +87,7 @@ class OnnxStableDiffusionInpaintPipeline(DiffusionPipeline): text_encoder: OnnxRuntimeModel tokenizer: CLIPTokenizer unet: OnnxRuntimeModel - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler] + scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler safety_checker: OnnxRuntimeModel feature_extractor: CLIPImageProcessor @@ -101,7 +101,7 @@ def __init__( text_encoder: OnnxRuntimeModel, tokenizer: CLIPTokenizer, unet: OnnxRuntimeModel, - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], + scheduler: DDIMScheduler | PNDMScheduler | LMSDiscreteScheduler, safety_checker: OnnxRuntimeModel, feature_extractor: CLIPImageProcessor, requires_safety_checker: bool = True, @@ -167,24 +167,24 @@ def __init__( # Copied from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion.OnnxStableDiffusionPipeline._encode_prompt def _encode_prompt( self, - prompt: Union[str, List[str]], - num_images_per_prompt: Optional[int], + prompt: str | list[str], + num_images_per_prompt: int | None, do_classifier_free_guidance: bool, - negative_prompt: Optional[str], - prompt_embeds: Optional[np.ndarray] = None, - negative_prompt_embeds: Optional[np.ndarray] = None, + negative_prompt: str | None, + prompt_embeds: np.ndarray | None = None, + negative_prompt_embeds: np.ndarray | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): prompt to be encoded num_images_per_prompt (`int`): number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`): + negative_prompt (`str` or `list[str]`): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). prompt_embeds (`np.ndarray`, *optional*): @@ -229,7 +229,7 @@ def _encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -271,13 +271,13 @@ def _encode_prompt( # Copied from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion.OnnxStableDiffusionPipeline.check_inputs def check_inputs( self, - prompt: Union[str, List[str]], - height: Optional[int], - width: Optional[int], + prompt: str | list[str], + height: int | None, + width: int | None, callback_steps: int, - negative_prompt: Optional[str] = None, - prompt_embeds: Optional[np.ndarray] = None, - negative_prompt_embeds: Optional[np.ndarray] = None, + negative_prompt: str | None = None, + prompt_embeds: np.ndarray | None = None, + negative_prompt_embeds: np.ndarray | None = None, ): if height % 8 != 0 or width % 8 != 0: raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") @@ -319,30 +319,30 @@ def check_inputs( @torch.no_grad() def __call__( self, - prompt: Union[str, List[str]], + prompt: str | list[str], image: PIL.Image.Image, mask_image: PIL.Image.Image, - height: Optional[int] = 512, - width: Optional[int] = 512, + height: int | None = 512, + width: int | None = 512, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[np.random.RandomState] = None, - latents: Optional[np.ndarray] = None, - prompt_embeds: Optional[np.ndarray] = None, - negative_prompt_embeds: Optional[np.ndarray] = None, - output_type: Optional[str] = "pil", + generator: np.random.RandomState | None = None, + latents: np.ndarray | None = None, + prompt_embeds: np.ndarray | None = None, + negative_prompt_embeds: np.ndarray | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, np.ndarray], None]] = None, + callback: Callable[[int, int, np.ndarray], None] | None = None, callback_steps: int = 1, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. image (`PIL.Image.Image`): `Image`, or tensor representing an image batch which will be inpainted, *i.e.* parts of the image will @@ -365,7 +365,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py b/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py index a765163175a2..c721a4f2ffa0 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -74,8 +74,8 @@ def __init__( unet: OnnxRuntimeModel, low_res_scheduler: DDPMScheduler, scheduler: KarrasDiffusionSchedulers, - safety_checker: Optional[OnnxRuntimeModel] = None, - feature_extractor: Optional[CLIPImageProcessor] = None, + safety_checker: OnnxRuntimeModel | None = None, + feature_extractor: CLIPImageProcessor | None = None, max_noise_level: int = 350, num_latent_channels=4, num_unet_input_channels=7, @@ -144,7 +144,7 @@ def __init__( def check_inputs( self, - prompt: Union[str, List[str]], + prompt: str | list[str], image, noise_level, callback_steps, @@ -245,24 +245,24 @@ def decode_latents(self, latents): def _encode_prompt( self, - prompt: Union[str, List[str]], - num_images_per_prompt: Optional[int], + prompt: str | list[str], + num_images_per_prompt: int | None, do_classifier_free_guidance: bool, - negative_prompt: Optional[str], - prompt_embeds: Optional[np.ndarray] = None, - negative_prompt_embeds: Optional[np.ndarray] = None, + negative_prompt: str | None, + prompt_embeds: np.ndarray | None = None, + negative_prompt_embeds: np.ndarray | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): prompt to be encoded num_images_per_prompt (`int`): number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`): + negative_prompt (`str` or `list[str]`): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). prompt_embeds (`np.ndarray`, *optional*): @@ -307,7 +307,7 @@ def _encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -348,28 +348,28 @@ def _encode_prompt( def __call__( self, - prompt: Union[str, List[str]], - image: Union[np.ndarray, PIL.Image.Image, List[PIL.Image.Image]], + prompt: str | list[str], + image: np.ndarray | PIL.Image.Image | list[PIL.Image.Image], num_inference_steps: int = 75, guidance_scale: float = 9.0, noise_level: int = 20, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[np.random.RandomState, List[np.random.RandomState]]] = None, - latents: Optional[np.ndarray] = None, - prompt_embeds: Optional[np.ndarray] = None, - negative_prompt_embeds: Optional[np.ndarray] = None, - output_type: Optional[str] = "pil", + generator: np.random.RandomState | list[np.random.RandomState] | None = None, + latents: np.ndarray | None = None, + prompt_embeds: np.ndarray | None = None, + negative_prompt_embeds: np.ndarray | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, np.ndarray], None]] = None, - callback_steps: Optional[int] = 1, + callback: Callable[[int, int, np.ndarray], None] | None = None, + callback_steps: int | None = 1, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. image (`np.ndarray` or `PIL.Image.Image`): `Image`, or tensor representing an image batch, that will be used as the starting point for the @@ -385,7 +385,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. noise_level (`float`, defaults to 0.2): Deteremines the amount of noise to add to the initial image before performing upscaling. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_output.py b/src/diffusers/pipelines/stable_diffusion/pipeline_output.py index 5fb9b1a1412d..804e5655d369 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_output.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Optional, Union import numpy as np import PIL.Image @@ -13,16 +12,16 @@ class StableDiffusionPipelineOutput(BaseOutput): Output class for Stable Diffusion pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, num_channels)`. - nsfw_content_detected (`List[bool]`) - List indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or + nsfw_content_detected (`list[bool]`) + list indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or `None` if safety checking could not be performed. """ - images: Union[List[PIL.Image.Image], np.ndarray] - nsfw_content_detected: Optional[List[bool]] + images: list[PIL.Image.Image] | np.ndarray + nsfw_content_detected: list[bool] | None if is_flax_available(): @@ -36,10 +35,10 @@ class FlaxStableDiffusionPipelineOutput(BaseOutput): Args: images (`np.ndarray`): Denoised images of array shape of `(batch_size, height, width, num_channels)`. - nsfw_content_detected (`List[bool]`): - List indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content + nsfw_content_detected (`list[bool]`): + list indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or `None` if safety checking could not be performed. """ images: np.ndarray - nsfw_content_detected: List[bool] + nsfw_content_detected: list[bool] diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index cb97f18efeff..42d62f53a20a 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch from packaging import version @@ -94,10 +94,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -112,15 +112,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -304,9 +304,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -336,16 +336,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -353,7 +353,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -452,7 +452,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -779,38 +779,36 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, guidance_rescale: float = 0.0, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -819,18 +817,18 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -838,7 +836,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -852,7 +850,7 @@ def __call__( Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -877,7 +875,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -989,8 +987,11 @@ def __call__( ) # 4. Prepare timesteps + timestep_device = device + if XLA_AVAILABLE: + timestep_device = "cpu" timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) # 5. Prepare latent variables @@ -1093,6 +1094,8 @@ def __call__( do_denormalize = [True] * image.shape[0] else: do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] + if XLA_AVAILABLE: + xm.mark_step() image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) # Offload all models diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py index e957c6661f87..3ec64c30763f 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py @@ -14,7 +14,7 @@ import contextlib import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -53,7 +53,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -175,9 +175,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -208,16 +208,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -225,7 +225,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -324,7 +324,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -634,33 +634,33 @@ def num_timesteps(self): @torch.no_grad() def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, image: PipelineImageInput = None, - depth_map: Optional[torch.Tensor] = None, + depth_map: torch.Tensor | None = None, strength: float = 0.8, - num_inference_steps: Optional[int] = 50, - guidance_scale: Optional[float] = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - eta: Optional[float] = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + num_inference_steps: int | None = 50, + guidance_scale: float | None = 7.5, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + eta: float | None = 0.0, + generator: torch.Generator | list[torch.Generator] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + cross_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image` or tensor representing an image batch to be used as the starting point. Can accept image latents as `image` only if `depth_map` is not `None`. depth_map (`torch.Tensor`, *optional*): @@ -678,7 +678,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -686,7 +686,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. prompt_embeds (`torch.Tensor`, *optional*): @@ -711,7 +711,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py index d47e2f0593dd..5974aea6b448 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Callable, List, Optional, Union +from typing import Callable import PIL.Image import torch @@ -218,7 +218,7 @@ def check_inputs(self, image, height, width, callback_steps): and not isinstance(image, list) ): raise ValueError( - "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is" + "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]` but is" f" {type(image)}" ) @@ -259,25 +259,25 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype @torch.no_grad() def __call__( self, - image: Union[PIL.Image.Image, List[PIL.Image.Image], torch.Tensor], - height: Optional[int] = None, - width: Optional[int] = None, + image: PIL.Image.Image | list[PIL.Image.Image] | torch.Tensor, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, - num_images_per_prompt: Optional[int] = 1, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, ): r""" The call function to the pipeline for generation. Args: - image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.Tensor`): + image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.Tensor`): Image or images to guide image generation. If you provide a tensor, it needs to be compatible with [`CLIPImageProcessor`](https://huggingface.co/lambdalabs/sd-image-variations-diffusers/blob/main/feature_extractor/preprocessor_config.json). height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): @@ -295,7 +295,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py index 95d3ab06f02a..abcd06a2bb3b 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -84,7 +84,7 @@ def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -122,10 +122,10 @@ def preprocess(image): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -140,15 +140,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -329,9 +329,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -362,16 +362,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -379,7 +379,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -478,7 +478,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -859,38 +859,36 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, image: PipelineImageInput = None, strength: float = 0.8, - num_inference_steps: Optional[int] = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - guidance_scale: Optional[float] = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - eta: Optional[float] = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + num_inference_steps: int | None = 50, + timesteps: list[int] = None, + sigmas: list[float] = None, + guidance_scale: float | None = 7.5, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + eta: float | None = 0.0, + generator: torch.Generator | list[torch.Generator] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, clip_skip: int = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a @@ -905,18 +903,18 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. This parameter is modulated by `strength`. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -924,7 +922,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. prompt_embeds (`torch.Tensor`, *optional*): @@ -934,7 +932,7 @@ def __call__( Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -955,7 +953,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1050,8 +1048,12 @@ def __call__( image = self.image_processor.preprocess(image) # 5. set timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py index 148d7386a732..a3e09d1ed1ad 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import PIL.Image import torch @@ -53,7 +53,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -68,10 +68,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -86,15 +86,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -173,7 +173,7 @@ class StableDiffusionInpaintPipeline( def __init__( self, - vae: Union[AutoencoderKL, AsymmetricAutoencoderKL], + vae: AutoencoderKL | AsymmetricAutoencoderKL, text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, unet: UNet2DConditionModel, @@ -283,9 +283,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -316,16 +316,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -333,7 +333,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -432,7 +432,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -880,51 +880,49 @@ def interrupt(self): @torch.no_grad() def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, image: PipelineImageInput = None, mask_image: PipelineImageInput = None, masked_image_latents: torch.Tensor = None, - height: Optional[int] = None, - width: Optional[int] = None, - padding_mask_crop: Optional[int] = None, + height: int | None = None, + width: int | None = None, + padding_mask_crop: int | None = None, strength: float = 1.0, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, clip_skip: int = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be inpainted (which parts of the image to be masked out with `mask_image` and repainted according to `prompt`). For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image latents as `image`, but if passing latents directly it is not encoded again. - mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one @@ -951,18 +949,18 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. This parameter is modulated by `strength`. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -970,7 +968,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -984,7 +982,7 @@ def __call__( Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -1005,7 +1003,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1136,8 +1134,12 @@ def __call__( ) # 4. set timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) timesteps, num_inference_steps = self.get_timesteps( num_inference_steps=num_inference_steps, strength=strength, device=device diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py index 843d25d67c10..c89d593d57be 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -68,7 +68,7 @@ def preprocess(image): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -134,7 +134,7 @@ def __init__( scheduler: KarrasDiffusionSchedulers, safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, - image_encoder: Optional[CLIPVisionModelWithProjection] = None, + image_encoder: CLIPVisionModelWithProjection | None = None, requires_safety_checker: bool = True, ): super().__init__() @@ -172,36 +172,34 @@ def __init__( @torch.no_grad() def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, image: PipelineImageInput = None, num_inference_steps: int = 100, guidance_scale: float = 7.5, image_guidance_scale: float = 1.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + cross_attention_kwargs: dict[str, Any] | None = None, **kwargs, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - image (`torch.Tensor` `np.ndarray`, `PIL.Image.Image`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor` `np.ndarray`, `PIL.Image.Image`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image` or tensor representing an image batch to be repainted according to `prompt`. Can also accept image latents as `image`, but if passing latents directly it is not encoded again. num_inference_steps (`int`, *optional*, defaults to 100): @@ -215,7 +213,7 @@ def __call__( `image_guidance_scale > 1`. Higher image guidance scale encourages generated images that are closely linked to the source `image`, usually at the expense of lower image quality. This pipeline requires a value of at least `1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -248,7 +246,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -497,14 +495,14 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -512,7 +510,7 @@ def _encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_ prompt (`str` or `List[str]`, *optional*): + negative_ prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -579,7 +577,7 @@ def _encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py index 66d5ffa6b849..02dc483c277a 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py @@ -13,7 +13,7 @@ # limitations under the License. import warnings -from typing import Callable, List, Optional, Union +from typing import Callable import numpy as np import PIL.Image @@ -42,7 +42,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -132,10 +132,10 @@ def _encode_prompt( device, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -169,10 +169,10 @@ def encode_prompt( device, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, ): r""" Encodes the prompt into text encoder hidden states. @@ -184,7 +184,7 @@ def encode_prompt( torch device do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`): + negative_prompt (`str` or `list[str]`): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). prompt_embeds (`torch.FloatTensor`, *optional*): @@ -243,7 +243,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance: if negative_prompt_embeds is None or negative_pooled_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -396,29 +396,29 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype @torch.no_grad() def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, image: PipelineImageInput = None, num_inference_steps: int = 75, guidance_scale: float = 9.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + negative_prompt: str | list[str] | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide image upscaling. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image` or tensor representing an image batch to be upscaled. If it's a tensor, it can be either a latent output from a Stable Diffusion model or an image tensor in the range `[-1, 1]`. It is considered a `latent` if `image.shape[1]` is `4`; otherwise, it is considered to be an image representation and @@ -429,13 +429,13 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -459,7 +459,6 @@ def __call__( >>> from diffusers import StableDiffusionLatentUpscalePipeline, StableDiffusionPipeline >>> import torch - >>> pipeline = StableDiffusionPipeline.from_pretrained( ... "CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16 ... ) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py index f13cdc67073f..4befa44550b7 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py @@ -14,7 +14,7 @@ import inspect import warnings -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -123,9 +123,9 @@ def __init__( unet: UNet2DConditionModel, low_res_scheduler: DDPMScheduler, scheduler: KarrasDiffusionSchedulers, - safety_checker: Optional[Any] = None, - feature_extractor: Optional[CLIPImageProcessor] = None, - watermarker: Optional[Any] = None, + safety_checker: Any | None = None, + feature_extractor: CLIPImageProcessor | None = None, + watermarker: Any | None = None, max_noise_level: int = 350, ): super().__init__() @@ -188,9 +188,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -221,16 +221,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -238,7 +238,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -337,7 +337,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -535,32 +535,32 @@ def upcast_vae(self): @torch.no_grad() def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, image: PipelineImageInput = None, num_inference_steps: int = 75, guidance_scale: float = 9.0, noise_level: int = 20, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, clip_skip: int = None, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image` or tensor representing an image batch to be upscaled. num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the @@ -568,7 +568,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -576,7 +576,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py index a134244e3ee4..39857ed230e6 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer @@ -173,8 +173,8 @@ def _encode_prior_prompt( device, num_images_per_prompt, do_classifier_free_guidance, - text_model_output: Optional[Union[CLIPTextModelOutput, Tuple]] = None, - text_attention_mask: Optional[torch.Tensor] = None, + text_model_output: CLIPTextModelOutput | tuple | None = None, + text_attention_mask: torch.Tensor | None = None, ): if text_model_output is None: batch_size = len(prompt) if isinstance(prompt, list) else 1 @@ -268,9 +268,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -301,16 +301,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -318,7 +318,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -417,7 +417,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -600,8 +600,8 @@ def noise_image_embeddings( self, image_embeds: torch.Tensor, noise_level: int, - noise: Optional[torch.Tensor] = None, - generator: Optional[torch.Generator] = None, + noise: torch.Tensor | None = None, + generator: torch.Generator | None = None, ): """ Add noise to the image embeddings. The amount of noise is controlled by a `noise_level` input. A higher @@ -647,35 +647,35 @@ def noise_image_embeddings( def __call__( self, # regular denoising process args - prompt: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 20, guidance_scale: float = 10.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, noise_level: int = 0, # prior args prior_num_inference_steps: int = 25, prior_guidance_scale: float = 4.0, - prior_latents: Optional[torch.Tensor] = None, - clip_skip: Optional[int] = None, + prior_latents: torch.Tensor | None = None, + clip_skip: int | None = None, ): """ The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -687,7 +687,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 10.0): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -695,7 +695,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py index abb4cc3a05d5..bb96e5db0295 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import PIL.Image import torch @@ -175,9 +175,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -263,16 +263,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -280,7 +280,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -379,7 +379,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -547,7 +547,7 @@ def check_inputs( and not isinstance(image, list) ): raise ValueError( - "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is" + "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]` but is" f" {type(image)}" ) @@ -579,8 +579,8 @@ def noise_image_embeddings( self, image_embeds: torch.Tensor, noise_level: int, - noise: Optional[torch.Tensor] = None, - generator: Optional[torch.Generator] = None, + noise: torch.Tensor | None = None, + generator: torch.Generator | None = None, ): """ Add noise to the image embeddings. The amount of noise is controlled by a `noise_level` input. A higher @@ -625,33 +625,33 @@ def noise_image_embeddings( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - image: Union[torch.Tensor, PIL.Image.Image] = None, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + image: torch.Tensor | PIL.Image.Image = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 20, guidance_scale: float = 10, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[torch.Generator] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, noise_level: int = 0, - image_embeds: Optional[torch.Tensor] = None, - clip_skip: Optional[int] = None, + image_embeds: torch.Tensor | None = None, + clip_skip: int | None = None, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, either `prompt_embeds` will be used or prompt is initialized to `""`. image (`torch.Tensor` or `PIL.Image.Image`): @@ -668,7 +668,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 10.0): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -676,7 +676,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/stable_diffusion/safety_checker_flax.py b/src/diffusers/pipelines/stable_diffusion/safety_checker_flax.py index 55d5023b6869..77a31b91d52d 100644 --- a/src/diffusers/pipelines/stable_diffusion/safety_checker_flax.py +++ b/src/diffusers/pipelines/stable_diffusion/safety_checker_flax.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Tuple - import jax import jax.numpy as jnp from flax import linen as nn @@ -76,7 +74,7 @@ class FlaxStableDiffusionSafetyChecker(FlaxPreTrainedModel): def __init__( self, config: CLIPConfig, - input_shape: Optional[Tuple] = None, + input_shape: tuple | None = None, seed: int = 0, dtype: jnp.dtype = jnp.float32, _do_init: bool = True, @@ -87,7 +85,7 @@ def __init__( module = self.module_class(config=config, dtype=dtype, **kwargs) super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init) - def init_weights(self, rng: jax.Array, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict: + def init_weights(self, rng: jax.Array, input_shape: tuple, params: FrozenDict = None) -> FrozenDict: # init input tensor clip_input = jax.random.normal(rng, input_shape) diff --git a/src/diffusers/pipelines/stable_diffusion/stable_unclip_image_normalizer.py b/src/diffusers/pipelines/stable_diffusion/stable_unclip_image_normalizer.py index ffd66792fe46..ba91a0f23923 100644 --- a/src/diffusers/pipelines/stable_diffusion/stable_unclip_image_normalizer.py +++ b/src/diffusers/pipelines/stable_diffusion/stable_unclip_image_normalizer.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Union - import torch from torch import nn @@ -41,8 +39,8 @@ def __init__( def to( self, - torch_device: Optional[Union[str, torch.device]] = None, - torch_dtype: Optional[torch.dtype] = None, + torch_device: str | torch.device | None = None, + torch_dtype: torch.dtype | None = None, ): self.mean = nn.Parameter(self.mean.to(torch_device).to(torch_dtype)) self.std = nn.Parameter(self.std.to(torch_device).to(torch_dtype)) diff --git a/src/diffusers/pipelines/stable_diffusion_3/pipeline_output.py b/src/diffusers/pipelines/stable_diffusion_3/pipeline_output.py index 4655f446102a..43f801d81e2e 100644 --- a/src/diffusers/pipelines/stable_diffusion_3/pipeline_output.py +++ b/src/diffusers/pipelines/stable_diffusion_3/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -13,9 +12,9 @@ class StableDiffusion3PipelineOutput(BaseOutput): Output class for Stable Diffusion pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image] | np.ndarray diff --git a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py index 660d9801df56..7764a79d7faf 100644 --- a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py +++ b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch from transformers import ( @@ -88,10 +88,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -106,15 +106,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -232,11 +232,11 @@ def __init__( def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 256, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -288,10 +288,10 @@ def _get_t5_prompt_embeds( def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - clip_skip: Optional[int] = None, + device: torch.device | None = None, + clip_skip: int | None = None, clip_model_index: int = 0, ): device = device or self._execution_device @@ -343,32 +343,32 @@ def _get_clip_prompt_embeds( def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], - prompt_3: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + prompt_2: str | list[str], + prompt_3: str | list[str], + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt_3: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_3: str | list[str] | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + clip_skip: int | None = None, max_sequence_length: int = 256, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in all text-encoders - prompt_3 (`str` or `List[str]`, *optional*): + prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is used in all text-encoders device: (`torch.device`): @@ -377,14 +377,14 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders. - negative_prompt_3 (`str` or `List[str]`, *optional*): + negative_prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `negative_prompt` is used in all the text-encoders. prompt_embeds (`torch.FloatTensor`, *optional*): @@ -715,9 +715,9 @@ def encode_image(self, image: PipelineImageInput, device: torch.device) -> torch # Adapted from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.prepare_ip_adapter_image_embeds def prepare_ip_adapter_image_embeds( self, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: torch.Tensor | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, ) -> torch.Tensor: @@ -773,50 +773,50 @@ def enable_sequential_cpu_offload(self, *args, **kwargs): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - prompt_3: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + prompt_3: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 28, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 7.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt_3: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_3: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 256, - skip_guidance_layers: List[int] = None, + skip_guidance_layers: list[int] = None, skip_layer_guidance_scale: float = 2.8, skip_layer_guidance_stop: float = 0.2, skip_layer_guidance_start: float = 0.01, - mu: Optional[float] = None, + mu: float | None = None, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead - prompt_3 (`str` or `List[str]`, *optional*): + prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is will be used instead height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): @@ -826,7 +826,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -836,19 +836,19 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used instead - negative_prompt_3 (`str` or `List[str]`, *optional*): + negative_prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `negative_prompt` is used instead num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -890,12 +890,12 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. max_sequence_length (`int` defaults to 256): Maximum sequence length to use with the `prompt`. - skip_guidance_layers (`List[int]`, *optional*): + skip_guidance_layers (`list[int]`, *optional*): A list of integers that specify layers to skip during guidance. If not provided, all layers will be used for guidance. If provided, the guidance will only be applied to the layers specified in the list. Recommended value by StabiltyAI for Stable Diffusion 3.5 Medium is [7, 8, 9]. @@ -1025,10 +1025,14 @@ def __call__( scheduler_kwargs["mu"] = mu elif mu is not None: scheduler_kwargs["mu"] = mu + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, - device, + timestep_device, sigmas=sigmas, **scheduler_kwargs, ) diff --git a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py index 9b11bc8781e7..7951b970cd0c 100644 --- a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py +++ b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import PIL.Image import torch @@ -93,7 +93,7 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -108,10 +108,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -126,15 +126,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -218,8 +218,8 @@ def __init__( tokenizer_2: CLIPTokenizer, text_encoder_3: T5EncoderModel, tokenizer_3: T5TokenizerFast, - image_encoder: Optional[SiglipVisionModel] = None, - feature_extractor: Optional[SiglipImageProcessor] = None, + image_encoder: SiglipVisionModel | None = None, + feature_extractor: SiglipImageProcessor | None = None, ): super().__init__() @@ -256,11 +256,11 @@ def __init__( # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 256, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -313,10 +313,10 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - clip_skip: Optional[int] = None, + device: torch.device | None = None, + clip_skip: int | None = None, clip_model_index: int = 0, ): device = device or self._execution_device @@ -369,32 +369,32 @@ def _get_clip_prompt_embeds( # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], - prompt_3: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + prompt_2: str | list[str], + prompt_3: str | list[str], + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt_3: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_3: str | list[str] | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + clip_skip: int | None = None, max_sequence_length: int = 256, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in all text-encoders - prompt_3 (`str` or `List[str]`, *optional*): + prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is used in all text-encoders device: (`torch.device`): @@ -403,14 +403,14 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders. - negative_prompt_3 (`str` or `List[str]`, *optional*): + negative_prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `negative_prompt` is used in all the text-encoders. prompt_embeds (`torch.FloatTensor`, *optional*): @@ -771,9 +771,9 @@ def encode_image(self, image: PipelineImageInput, device: torch.device) -> torch # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.prepare_ip_adapter_image_embeds def prepare_ip_adapter_image_embeds( self, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: torch.Tensor | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, ) -> torch.Tensor: @@ -830,48 +830,48 @@ def enable_sequential_cpu_offload(self, *args, **kwargs): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - prompt_3: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + prompt_3: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, image: PipelineImageInput = None, strength: float = 0.6, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 7.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt_3: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[torch.Tensor] = None, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_3: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: torch.Tensor | None = None, return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 256, - mu: Optional[float] = None, + mu: float | None = None, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead - prompt_3 (`str` or `List[str]`, *optional*): + prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is will be used instead height (`int`, *optional*, defaults to self.transformer.config.sample_size * self.vae_scale_factor): @@ -881,7 +881,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -891,19 +891,19 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used instead - negative_prompt_3 (`str` or `List[str]`, *optional*): + negative_prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `negative_prompt` is used instead num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -945,7 +945,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1047,8 +1047,12 @@ def __call__( scheduler_kwargs["mu"] = mu elif mu is not None: scheduler_kwargs["mu"] = mu + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, sigmas=sigmas, **scheduler_kwargs + self.scheduler, num_inference_steps, timestep_device, sigmas=sigmas, **scheduler_kwargs ) timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) diff --git a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py index b947cbff0914..d3594b868f89 100644 --- a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py +++ b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch from transformers import ( @@ -92,7 +92,7 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -107,10 +107,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -125,15 +125,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -217,8 +217,8 @@ def __init__( tokenizer_2: CLIPTokenizer, text_encoder_3: T5EncoderModel, tokenizer_3: T5TokenizerFast, - image_encoder: Optional[SiglipVisionModel] = None, - feature_extractor: Optional[SiglipImageProcessor] = None, + image_encoder: SiglipVisionModel | None = None, + feature_extractor: SiglipImageProcessor | None = None, ): super().__init__() @@ -262,11 +262,11 @@ def __init__( # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 256, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -319,10 +319,10 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, - clip_skip: Optional[int] = None, + device: torch.device | None = None, + clip_skip: int | None = None, clip_model_index: int = 0, ): device = device or self._execution_device @@ -375,32 +375,32 @@ def _get_clip_prompt_embeds( # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - prompt_2: Union[str, List[str]], - prompt_3: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + prompt_2: str | list[str], + prompt_3: str | list[str], + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt_3: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_3: str | list[str] | None = None, + prompt_embeds: torch.FloatTensor | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + negative_pooled_prompt_embeds: torch.FloatTensor | None = None, + clip_skip: int | None = None, max_sequence_length: int = 256, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in all text-encoders - prompt_3 (`str` or `List[str]`, *optional*): + prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is used in all text-encoders device: (`torch.device`): @@ -409,14 +409,14 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders. - negative_prompt_3 (`str` or `List[str]`, *optional*): + negative_prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `negative_prompt` is used in all the text-encoders. prompt_embeds (`torch.FloatTensor`, *optional*): @@ -863,9 +863,9 @@ def encode_image(self, image: PipelineImageInput, device: torch.device) -> torch # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.prepare_ip_adapter_image_embeds def prepare_ip_adapter_image_embeds( self, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[torch.Tensor] = None, - device: Optional[torch.device] = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: torch.Tensor | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, ) -> torch.Tensor: @@ -922,67 +922,67 @@ def enable_sequential_cpu_offload(self, *args, **kwargs): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - prompt_3: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + prompt_3: str | list[str] | None = None, image: PipelineImageInput = None, mask_image: PipelineImageInput = None, masked_image_latents: PipelineImageInput = None, height: int = None, width: int = None, - padding_mask_crop: Optional[int] = None, + padding_mask_crop: int | None = None, strength: float = 0.6, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 7.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - negative_prompt_3: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + negative_prompt_3: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 256, - mu: Optional[float] = None, + mu: float | None = None, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead - prompt_3 (`str` or `List[str]`, *optional*): + prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is will be used instead - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image latents as `image`, but if passing latents directly it is not encoded again. - mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to mask `image`. White pixels in the mask are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one color channel (L) instead of 3, so the expected shape for pytorch tensor would be `(B, 1, H, W)`, `(B, H, W)`, `(1, H, W)`, `(H, W)`. And for numpy array would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W, 1)`, or `(H, W)`. - mask_image_latent (`torch.Tensor`, `List[torch.Tensor]`): + mask_image_latent (`torch.Tensor`, `list[torch.Tensor]`): `Tensor` representing an image batch to mask `image` generated by VAE. If not provided, the mask latents tensor will be generated by `mask_image`. height (`int`, *optional*, defaults to self.transformer.config.sample_size * self.vae_scale_factor): @@ -1005,7 +1005,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -1015,19 +1015,19 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used instead - negative_prompt_3 (`str` or `List[str]`, *optional*): + negative_prompt_3 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `negative_prompt` is used instead num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -1069,7 +1069,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1167,8 +1167,13 @@ def __call__( scheduler_kwargs["mu"] = mu elif mu is not None: scheduler_kwargs["mu"] = mu + + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, sigmas=sigmas, **scheduler_kwargs + self.scheduler, num_inference_steps, timestep_device, sigmas=sigmas, **scheduler_kwargs ) timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) # check that number of inference steps is not < 1 - as this doesn't make sense diff --git a/src/diffusers/pipelines/stable_diffusion_attend_and_excite/pipeline_stable_diffusion_attend_and_excite.py b/src/diffusers/pipelines/stable_diffusion_attend_and_excite/pipeline_stable_diffusion_attend_and_excite.py index a1ff99b6aa34..80b0c09bc9a5 100644 --- a/src/diffusers/pipelines/stable_diffusion_attend_and_excite/pipeline_stable_diffusion_attend_and_excite.py +++ b/src/diffusers/pipelines/stable_diffusion_attend_and_excite/pipeline_stable_diffusion_attend_and_excite.py @@ -14,7 +14,7 @@ import inspect import math -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import torch @@ -111,7 +111,7 @@ def get_average_attention(self): average_attention = self.attention_store return average_attention - def aggregate_attention(self, from_where: List[str]) -> torch.Tensor: + def aggregate_attention(self, from_where: list[str]) -> torch.Tensor: """Aggregates the attention across the different layers and heads at the specified resolution.""" out = [] attention_maps = self.get_average_attention() @@ -267,9 +267,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -300,16 +300,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -317,7 +317,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -416,7 +416,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -619,8 +619,8 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype @staticmethod def _compute_max_attention_per_index( attention_maps: torch.Tensor, - indices: List[int], - ) -> List[torch.Tensor]: + indices: list[int], + ) -> list[torch.Tensor]: """Computes the maximum attention value for each of the tokens we wish to alter.""" attention_for_text = attention_maps[:, :, 1:-1] attention_for_text *= 100 @@ -641,7 +641,7 @@ def _compute_max_attention_per_index( def _aggregate_and_get_max_attention_per_token( self, - indices: List[int], + indices: list[int], ): """Aggregates the attention for each token and computes the max activation value for each token to alter.""" attention_maps = self.attention_store.aggregate_attention( @@ -654,7 +654,7 @@ def _aggregate_and_get_max_attention_per_token( return max_attention_per_index @staticmethod - def _compute_loss(max_attention_per_index: List[torch.Tensor]) -> torch.Tensor: + def _compute_loss(max_attention_per_index: list[torch.Tensor]) -> torch.Tensor: """Computes the attend-and-excite loss using the maximum attention value for each token.""" losses = [max(0, 1.0 - curr_max) for curr_max in max_attention_per_index] loss = max(losses) @@ -670,7 +670,7 @@ def _update_latent(latents: torch.Tensor, loss: torch.Tensor, step_size: float) def _perform_iterative_refinement_step( self, latents: torch.Tensor, - indices: List[int], + indices: list[int], loss: torch.Tensor, threshold: float, text_embeddings: torch.Tensor, @@ -740,7 +740,7 @@ def register_attention_control(self): self.unet.set_attn_processor(attn_procs) self.attention_store.num_att_layers = cross_att_count - def get_indices(self, prompt: str) -> Dict[str, int]: + def get_indices(self, prompt: str) -> dict[str, int]: """Utility function to list the indices of the tokens you wish to alte""" ids = self.tokenizer(prompt).input_ids indices = {i: tok for tok, i in zip(self.tokenizer.convert_ids_to_tokens(ids), range(len(ids)))} @@ -750,37 +750,37 @@ def get_indices(self, prompt: str) -> Dict[str, int]: @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]], - token_indices: Union[List[int], List[List[int]]], - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str], + token_indices: list[int] | list[list[int]], + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, + negative_prompt: str | list[str] | None = None, num_images_per_prompt: int = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, max_iter_to_alter: int = 25, thresholds: dict = {0: 0.05, 10: 0.5, 20: 0.8}, scale_factor: int = 20, - attn_res: Optional[Tuple[int]] = (16, 16), - clip_skip: Optional[int] = None, + attn_res: tuple[int] | None = (16, 16), + clip_skip: int | None = None, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - token_indices (`List[int]`): + token_indices (`list[int]`): The token indices to alter with attend-and-excite. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -792,7 +792,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -800,7 +800,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py b/src/diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py index 65c25ffbe492..43bc2eb955c7 100644 --- a/src/diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py +++ b/src/diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py @@ -14,7 +14,7 @@ import inspect from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -63,14 +63,14 @@ class DiffEditInversionPipelineOutput(BaseOutput): Args: latents (`torch.Tensor`) inverted latents tensor - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `num_timesteps * batch_size` or numpy array of shape `(num_timesteps, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `num_timesteps * batch_size` or numpy array of shape `(num_timesteps, batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ latents: torch.Tensor - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image] | np.ndarray EXAMPLE_DOC_STRING = """ @@ -393,9 +393,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -426,16 +426,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -443,7 +443,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -542,7 +542,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -844,23 +844,23 @@ def get_epsilon(self, model_output: torch.Tensor, sample: torch.Tensor, timestep @replace_example_docstring(EXAMPLE_DOC_STRING) def generate_mask( self, - image: Union[torch.Tensor, PIL.Image.Image] = None, - target_prompt: Optional[Union[str, List[str]]] = None, - target_negative_prompt: Optional[Union[str, List[str]]] = None, - target_prompt_embeds: Optional[torch.Tensor] = None, - target_negative_prompt_embeds: Optional[torch.Tensor] = None, - source_prompt: Optional[Union[str, List[str]]] = None, - source_negative_prompt: Optional[Union[str, List[str]]] = None, - source_prompt_embeds: Optional[torch.Tensor] = None, - source_negative_prompt_embeds: Optional[torch.Tensor] = None, - num_maps_per_mask: Optional[int] = 10, - mask_encode_strength: Optional[float] = 0.5, - mask_thresholding_ratio: Optional[float] = 3.0, + image: torch.Tensor | PIL.Image.Image = None, + target_prompt: str | list[str] | None = None, + target_negative_prompt: str | list[str] | None = None, + target_prompt_embeds: torch.Tensor | None = None, + target_negative_prompt_embeds: torch.Tensor | None = None, + source_prompt: str | list[str] | None = None, + source_negative_prompt: str | list[str] | None = None, + source_prompt_embeds: torch.Tensor | None = None, + source_negative_prompt_embeds: torch.Tensor | None = None, + num_maps_per_mask: int | None = 10, + mask_encode_strength: float | None = 0.5, + mask_thresholding_ratio: float | None = 3.0, num_inference_steps: int = 50, guidance_scale: float = 7.5, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - output_type: Optional[str] = "np", - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + output_type: str | None = "np", + cross_attention_kwargs: dict[str, Any] | None = None, ): r""" Generate a latent mask given a mask prompt, a target prompt, and an image. @@ -868,10 +868,10 @@ def generate_mask( Args: image (`PIL.Image.Image`): `Image` or tensor representing an image batch to be used for computing the mask. - target_prompt (`str` or `List[str]`, *optional*): + target_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide semantic mask generation. If not defined, you need to pass `prompt_embeds`. - target_negative_prompt (`str` or `List[str]`, *optional*): + target_negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). target_prompt_embeds (`torch.Tensor`, *optional*): @@ -880,10 +880,10 @@ def generate_mask( target_negative_prompt_embeds (`torch.Tensor`, *optional*): Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. - source_prompt (`str` or `List[str]`, *optional*): + source_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide semantic mask generation using DiffEdit. If not defined, you need to pass `source_prompt_embeds` or `source_image` instead. - source_negative_prompt (`str` or `List[str]`, *optional*): + source_negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide semantic mask generation away from using DiffEdit. If not defined, you need to pass `source_negative_prompt_embeds` or `source_image` instead. source_prompt_embeds (`torch.Tensor`, *optional*): @@ -908,7 +908,7 @@ def generate_mask( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. output_type (`str`, *optional*, defaults to `"pil"`): @@ -921,8 +921,8 @@ def generate_mask( Examples: Returns: - `List[PIL.Image.Image]` or `np.array`: - When returning a `List[PIL.Image.Image]`, the list consists of a batch of single-channel binary images + `list[PIL.Image.Image]` or `np.array`: + When returning a `list[PIL.Image.Image]`, the list consists of a batch of single-channel binary images with dimensions `(height // self.vae_scale_factor, width // self.vae_scale_factor)`. If it's `np.array`, the shape is `(batch_size, height // self.vae_scale_factor, width // self.vae_scale_factor)`. @@ -1063,21 +1063,21 @@ def generate_mask( @replace_example_docstring(EXAMPLE_INVERT_DOC_STRING) def invert( self, - prompt: Optional[Union[str, List[str]]] = None, - image: Union[torch.Tensor, PIL.Image.Image] = None, + prompt: str | list[str] | None = None, + image: torch.Tensor | PIL.Image.Image = None, num_inference_steps: int = 50, inpaint_strength: float = 0.8, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt: str | list[str] | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, decode_latents: bool = False, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: Optional[int] = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, + callback_steps: int | None = 1, + cross_attention_kwargs: dict[str, Any] | None = None, lambda_auto_corr: float = 20.0, lambda_kl: float = 20.0, num_reg_steps: int = 0, @@ -1087,7 +1087,7 @@ def invert( Generate inverted latents given a prompt and image. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. image (`PIL.Image.Image`): `Image` or tensor representing an image batch to produce the inverted latents guided by `prompt`. @@ -1102,7 +1102,7 @@ def invert( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). generator (`torch.Generator`, *optional*): @@ -1301,31 +1301,31 @@ def invert( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Optional[Union[str, List[str]]] = None, - mask_image: Union[torch.Tensor, PIL.Image.Image] = None, - image_latents: Union[torch.Tensor, PIL.Image.Image] = None, - inpaint_strength: Optional[float] = 0.8, + prompt: str | list[str] | None = None, + mask_image: torch.Tensor | PIL.Image.Image = None, + image_latents: torch.Tensor | PIL.Image.Image = None, + inpaint_strength: float | None = 0.8, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, clip_skip: int = None, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. mask_image (`PIL.Image.Image`): `Image` or tensor representing an image batch to mask the generated image. White pixels in the mask are @@ -1345,7 +1345,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): diff --git a/src/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen.py b/src/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen.py index 78b026684cfa..fa5bc9376e53 100644 --- a/src/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen.py +++ b/src/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen.py @@ -14,7 +14,7 @@ import inspect import warnings -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import PIL.Image import torch @@ -191,9 +191,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -224,16 +224,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -241,7 +241,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -340,7 +340,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -540,34 +540,34 @@ def target_size_center_crop(self, im, new_hw): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, gligen_scheduled_sampling_beta: float = 0.3, - gligen_phrases: List[str] = None, - gligen_boxes: List[List[float]] = None, - gligen_inpaint_image: Optional[PIL.Image.Image] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + gligen_phrases: list[str] = None, + gligen_boxes: list[list[float]] = None, + gligen_inpaint_image: PIL.Image.Image | None = None, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, + cross_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -579,13 +579,13 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - gligen_phrases (`List[str]`): + gligen_phrases (`list[str]`): The phrases to guide what to include in each of the regions defined by the corresponding `gligen_boxes`. There should only be one phrase per bounding box. - gligen_boxes (`List[List[float]]`): + gligen_boxes (`list[list[float]]`): The bounding boxes that identify rectangular regions of the image that are going to be filled with the content described by the corresponding `gligen_phrases`. Each rectangular box is defined as a - `List[float]` of 4 elements `[xmin, ymin, xmax, ymax]` where each value is between [0,1]. + `list[float]` of 4 elements `[xmin, ymin, xmax, ymax]` where each value is between [0,1]. gligen_inpaint_image (`PIL.Image.Image`, *optional*): The input image, if provided, is inpainted with objects described by the `gligen_boxes` and `gligen_phrases`. Otherwise, it is treated as a generation task on a blank input image. @@ -593,7 +593,7 @@ def __call__( Scheduled Sampling factor from [GLIGEN: Open-Set Grounded Text-to-Image Generation](https://huggingface.co/papers/2301.07093). Scheduled Sampling factor is only varied for scheduled sampling during inference for improved quality and controllability. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -601,7 +601,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py b/src/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py index 05cbad139d92..62e8a9fa95ae 100644 --- a/src/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py +++ b/src/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py @@ -14,7 +14,7 @@ import inspect import warnings -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import PIL.Image import torch @@ -255,16 +255,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -272,7 +272,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -371,7 +371,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -715,30 +715,30 @@ def get_cross_attention_kwargs_without_grounded(self, hidden_size, repeat_batch, @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, gligen_scheduled_sampling_beta: float = 0.3, - gligen_phrases: List[str] = None, - gligen_images: List[PIL.Image.Image] = None, - input_phrases_mask: Union[int, List[int]] = None, - input_images_mask: Union[int, List[int]] = None, - gligen_boxes: List[List[float]] = None, - gligen_inpaint_image: Optional[PIL.Image.Image] = None, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + gligen_phrases: list[str] = None, + gligen_images: list[PIL.Image.Image] = None, + input_phrases_mask: int | list[int] = None, + input_images_mask: int | list[int] = None, + gligen_boxes: list[list[float]] = None, + gligen_inpaint_image: PIL.Image.Image | None = None, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, gligen_normalize_constant: float = 28.7, clip_skip: int = None, ): @@ -746,7 +746,7 @@ def __call__( The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -758,20 +758,20 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - gligen_phrases (`List[str]`): + gligen_phrases (`list[str]`): The phrases to guide what to include in each of the regions defined by the corresponding `gligen_boxes`. There should only be one phrase per bounding box. - gligen_images (`List[PIL.Image.Image]`): + gligen_images (`list[PIL.Image.Image]`): The images to guide what to include in each of the regions defined by the corresponding `gligen_boxes`. There should only be one image per bounding box - input_phrases_mask (`int` or `List[int]`): + input_phrases_mask (`int` or `list[int]`): pre phrases mask input defined by the correspongding `input_phrases_mask` - input_images_mask (`int` or `List[int]`): + input_images_mask (`int` or `list[int]`): pre images mask input defined by the correspongding `input_images_mask` - gligen_boxes (`List[List[float]]`): + gligen_boxes (`list[list[float]]`): The bounding boxes that identify rectangular regions of the image that are going to be filled with the content described by the corresponding `gligen_phrases`. Each rectangular box is defined as a - `List[float]` of 4 elements `[xmin, ymin, xmax, ymax]` where each value is between [0,1]. + `list[float]` of 4 elements `[xmin, ymin, xmax, ymax]` where each value is between [0,1]. gligen_inpaint_image (`PIL.Image.Image`, *optional*): The input image, if provided, is inpainted with objects described by the `gligen_boxes` and `gligen_phrases`. Otherwise, it is treated as a generation task on a blank input image. @@ -779,7 +779,7 @@ def __call__( Scheduled Sampling factor from [GLIGEN: Open-Set Grounded Text-to-Image Generation](https://huggingface.co/papers/2301.07093). Scheduled Sampling factor is only varied for scheduled sampling during inference for improved quality and controllability. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -787,7 +787,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/stable_diffusion_k_diffusion/__init__.py b/src/diffusers/pipelines/stable_diffusion_k_diffusion/__init__.py deleted file mode 100644 index 7eb5bf8c229b..000000000000 --- a/src/diffusers/pipelines/stable_diffusion_k_diffusion/__init__.py +++ /dev/null @@ -1,62 +0,0 @@ -from typing import TYPE_CHECKING - -from ...utils import ( - DIFFUSERS_SLOW_IMPORT, - OptionalDependencyNotAvailable, - _LazyModule, - get_objects_from_module, - is_k_diffusion_available, - is_k_diffusion_version, - is_torch_available, - is_transformers_available, -) - - -_dummy_objects = {} -_import_structure = {} - - -try: - if not ( - is_transformers_available() - and is_torch_available() - and is_k_diffusion_available() - and is_k_diffusion_version(">=", "0.0.12") - ): - raise OptionalDependencyNotAvailable() -except OptionalDependencyNotAvailable: - from ...utils import dummy_torch_and_transformers_and_k_diffusion_objects # noqa F403 - - _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_and_k_diffusion_objects)) -else: - _import_structure["pipeline_stable_diffusion_k_diffusion"] = ["StableDiffusionKDiffusionPipeline"] - _import_structure["pipeline_stable_diffusion_xl_k_diffusion"] = ["StableDiffusionXLKDiffusionPipeline"] - -if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - try: - if not ( - is_transformers_available() - and is_torch_available() - and is_k_diffusion_available() - and is_k_diffusion_version(">=", "0.0.12") - ): - raise OptionalDependencyNotAvailable() - - except OptionalDependencyNotAvailable: - from ...utils.dummy_torch_and_transformers_and_k_diffusion_objects import * - else: - from .pipeline_stable_diffusion_k_diffusion import StableDiffusionKDiffusionPipeline - from .pipeline_stable_diffusion_xl_k_diffusion import StableDiffusionXLKDiffusionPipeline - -else: - import sys - - sys.modules[__name__] = _LazyModule( - __name__, - globals()["__file__"], - _import_structure, - module_spec=__spec__, - ) - - for name, value in _dummy_objects.items(): - setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py b/src/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py deleted file mode 100755 index feebd6adf8f8..000000000000 --- a/src/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py +++ /dev/null @@ -1,689 +0,0 @@ -# Copyright 2025 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import importlib -import inspect -from typing import Callable, List, Optional, Union - -import torch -from k_diffusion.external import CompVisDenoiser, CompVisVDenoiser -from k_diffusion.sampling import BrownianTreeNoiseSampler, get_sigmas_karras -from transformers import ( - CLIPImageProcessor, - CLIPTextModel, - CLIPTokenizer, - CLIPTokenizerFast, -) - -from ...image_processor import VaeImageProcessor -from ...loaders import ( - StableDiffusionLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ...models import AutoencoderKL, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers, LMSDiscreteScheduler -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - logging, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion import StableDiffusionPipelineOutput, StableDiffusionSafetyChecker - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class ModelWrapper: - def __init__(self, model, alphas_cumprod): - self.model = model - self.alphas_cumprod = alphas_cumprod - - def apply_model(self, *args, **kwargs): - if len(args) == 3: - encoder_hidden_states = args[-1] - args = args[:2] - if kwargs.get("cond", None) is not None: - encoder_hidden_states = kwargs.pop("cond") - return self.model(*args, encoder_hidden_states=encoder_hidden_states, **kwargs).sample - - -class StableDiffusionKDiffusionPipeline( - DeprecatedPipelineMixin, - DiffusionPipeline, - StableDiffusionMixin, - TextualInversionLoaderMixin, - StableDiffusionLoraLoaderMixin, -): - r""" - Pipeline for text-to-image generation using Stable Diffusion. - - This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the - library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.) - - The pipeline also inherits the following loading methods: - - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings - - [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for loading LoRA weights - - [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for saving LoRA weights - - > [!WARNING] > This is an experimental pipeline and is likely to change in the future. - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`CLIPTextModel`]): - Frozen text-encoder. Stable Diffusion uses the text portion of - [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically - the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant. - tokenizer (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). - unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of - [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. - safety_checker ([`StableDiffusionSafetyChecker`]): - Classification module that estimates whether generated images could be considered offensive or harmful. - Please, refer to the [model card](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) for - details. - feature_extractor ([`CLIPImageProcessor`]): - Model that extracts features from generated images to be used as inputs for the `safety_checker`. - """ - - _last_supported_version = "0.33.1" - - model_cpu_offload_seq = "text_encoder->unet->vae" - _optional_components = ["safety_checker", "feature_extractor"] - _exclude_from_cpu_offload = ["safety_checker"] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - tokenizer: Union[CLIPTokenizer, CLIPTokenizerFast], - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - safety_checker: StableDiffusionSafetyChecker, - feature_extractor: CLIPImageProcessor, - requires_safety_checker: bool = True, - ): - super().__init__() - - logger.info( - f"{self.__class__} is an experimental pipeline and is likely to change in the future. We recommend to use" - " this pipeline for fast experimentation / iteration if needed, but advice to rely on existing pipelines" - " as defined in https://huggingface.co/docs/diffusers/api/schedulers#implemented-schedulers for" - " production settings." - ) - - # get correct sigmas from LMS - scheduler = LMSDiscreteScheduler.from_config(scheduler.config) - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - ) - self.register_to_config(requires_safety_checker=requires_safety_checker) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) - - model = ModelWrapper(unet, scheduler.alphas_cumprod) - if scheduler.config.prediction_type == "v_prediction": - self.k_diffusion_model = CompVisVDenoiser(model) - else: - self.k_diffusion_model = CompVisDenoiser(model) - - def set_scheduler(self, scheduler_type: str): - library = importlib.import_module("k_diffusion") - sampling = getattr(library, "sampling") - try: - self.sampler = getattr(sampling, scheduler_type) - except Exception: - valid_samplers = [] - for s in dir(sampling): - if "sample_" in s: - valid_samplers.append(s) - - raise ValueError(f"Invalid scheduler type {scheduler_type}. Please choose one of {valid_samplers}.") - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt - def _encode_prompt( - self, - prompt, - device, - num_images_per_prompt, - do_classifier_free_guidance, - negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - **kwargs, - ): - deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." - deprecate("_encode_prompt()", "1.0.0", deprecation_message, standard_warn=False) - - prompt_embeds_tuple = self.encode_prompt( - prompt=prompt, - device=device, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=do_classifier_free_guidance, - negative_prompt=negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - lora_scale=lora_scale, - **kwargs, - ) - - # concatenate for backwards comp - prompt_embeds = torch.cat([prompt_embeds_tuple[1], prompt_embeds_tuple[0]]) - - return prompt_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_prompt - def encode_prompt( - self, - prompt, - device, - num_images_per_prompt, - do_classifier_free_guidance, - negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` or `List[str]`, *optional*): - prompt to be encoded - device: (`torch.device`): - torch device - num_images_per_prompt (`int`): - number of images that should be generated per prompt - do_classifier_free_guidance (`bool`): - whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): - The prompt or prompts not to guide the image generation. If not defined, one has to pass - `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is - less than `1`). - prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not - provided, text embeddings will be generated from `prompt` input argument. - negative_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input - argument. - lora_scale (`float`, *optional*): - A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded. - clip_skip (`int`, *optional*): - Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that - the output of the pre-final layer will be used for computing the prompt embeddings. - """ - # set lora scale so that monkey patched LoRA - # function of text encoder can correctly access it - if lora_scale is not None and isinstance(self, StableDiffusionLoraLoaderMixin): - self._lora_scale = lora_scale - - # dynamically adjust the LoRA scale - if not USE_PEFT_BACKEND: - adjust_lora_scale_text_encoder(self.text_encoder, lora_scale) - else: - scale_lora_layers(self.text_encoder, lora_scale) - - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - if prompt_embeds is None: - # textual inversion: process multi-vector tokens if necessary - if isinstance(self, TextualInversionLoaderMixin): - prompt = self.maybe_convert_prompt(prompt, self.tokenizer) - - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=self.tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids - untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids - - if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal( - text_input_ids, untruncated_ids - ): - removed_text = self.tokenizer.batch_decode( - untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1] - ) - logger.warning( - "The following part of your input was truncated because CLIP can only handle sequences up to" - f" {self.tokenizer.model_max_length} tokens: {removed_text}" - ) - - if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask: - attention_mask = text_inputs.attention_mask.to(device) - else: - attention_mask = None - - if clip_skip is None: - prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask) - prompt_embeds = prompt_embeds[0] - else: - prompt_embeds = self.text_encoder( - text_input_ids.to(device), attention_mask=attention_mask, output_hidden_states=True - ) - # Access the `hidden_states` first, that contains a tuple of - # all the hidden states from the encoder layers. Then index into - # the tuple to access the hidden states from the desired layer. - prompt_embeds = prompt_embeds[-1][-(clip_skip + 1)] - # We also need to apply the final LayerNorm here to not mess with the - # representations. The `last_hidden_states` that we typically use for - # obtaining the final prompt representations passes through the LayerNorm - # layer. - prompt_embeds = self.text_encoder.text_model.final_layer_norm(prompt_embeds) - - if self.text_encoder is not None: - prompt_embeds_dtype = self.text_encoder.dtype - elif self.unet is not None: - prompt_embeds_dtype = self.unet.dtype - else: - prompt_embeds_dtype = prompt_embeds.dtype - - prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device) - - bs_embed, seq_len, _ = prompt_embeds.shape - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) - prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1) - - # get unconditional embeddings for classifier free guidance - if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] - if negative_prompt is None: - uncond_tokens = [""] * batch_size - elif prompt is not None and type(prompt) is not type(negative_prompt): - raise TypeError( - f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" - f" {type(prompt)}." - ) - elif isinstance(negative_prompt, str): - uncond_tokens = [negative_prompt] - elif batch_size != len(negative_prompt): - raise ValueError( - f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" - f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" - " the batch size of `prompt`." - ) - else: - uncond_tokens = negative_prompt - - # textual inversion: process multi-vector tokens if necessary - if isinstance(self, TextualInversionLoaderMixin): - uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer) - - max_length = prompt_embeds.shape[1] - uncond_input = self.tokenizer( - uncond_tokens, - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="pt", - ) - - if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask: - attention_mask = uncond_input.attention_mask.to(device) - else: - attention_mask = None - - negative_prompt_embeds = self.text_encoder( - uncond_input.input_ids.to(device), - attention_mask=attention_mask, - ) - negative_prompt_embeds = negative_prompt_embeds[0] - - if do_classifier_free_guidance: - # duplicate unconditional embeddings for each generation per prompt, using mps friendly method - seq_len = negative_prompt_embeds.shape[1] - - negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device) - - negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - - if self.text_encoder is not None: - if isinstance(self, StableDiffusionLoraLoaderMixin) and USE_PEFT_BACKEND: - # Retrieve the original scale by scaling back the LoRA layers - unscale_lora_layers(self.text_encoder, lora_scale) - - return prompt_embeds, negative_prompt_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker - def run_safety_checker(self, image, device, dtype): - if self.safety_checker is None: - has_nsfw_concept = None - else: - if torch.is_tensor(image): - feature_extractor_input = self.image_processor.postprocess(image, output_type="pil") - else: - feature_extractor_input = self.image_processor.numpy_to_pil(image) - safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device) - image, has_nsfw_concept = self.safety_checker( - images=image, clip_input=safety_checker_input.pixel_values.to(dtype) - ) - return image, has_nsfw_concept - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents - def decode_latents(self, latents): - deprecation_message = "The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead" - deprecate("decode_latents", "1.0.0", deprecation_message, standard_warn=False) - - latents = 1 / self.vae.config.scaling_factor * latents - image = self.vae.decode(latents, return_dict=False)[0] - image = (image / 2 + 0.5).clamp(0, 1) - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - return image - - def check_inputs( - self, - prompt, - height, - width, - callback_steps, - negative_prompt=None, - prompt_embeds=None, - negative_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - ): - if height % 8 != 0 or width % 8 != 0: - raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") - - if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0): - raise ValueError( - f"`callback_steps` has to be a positive integer but is {callback_steps} of type" - f" {type(callback_steps)}." - ) - if callback_on_step_end_tensor_inputs is not None and not all( - k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs - ): - raise ValueError( - f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" - ) - - if prompt is not None and prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" - " only forward one of the two." - ) - elif prompt is None and prompt_embeds is None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." - ) - elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - if negative_prompt is not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - if prompt_embeds is not None and negative_prompt_embeds is not None: - if prompt_embeds.shape != negative_prompt_embeds.shape: - raise ValueError( - "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" - f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" - f" {negative_prompt_embeds.shape}." - ) - - def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None): - shape = ( - batch_size, - num_channels_latents, - int(height) // self.vae_scale_factor, - int(width) // self.vae_scale_factor, - ) - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - if latents.shape != shape: - raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}") - latents = latents.to(device) - - # scale the initial noise by the standard deviation required by the scheduler - return latents - - @torch.no_grad() - def __call__( - self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: int = 1, - use_karras_sigmas: Optional[bool] = False, - noise_sampler_seed: Optional[int] = None, - clip_skip: int = None, - ): - r""" - Function invoked when calling the pipeline for generation. - - Args: - prompt (`str` or `List[str]`, *optional*): - The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. - instead. - height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): - The width in pixels of the generated image. - num_inference_steps (`int`, *optional*, defaults to 50): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. - guidance_scale (`float`, *optional*, defaults to 7.5): - Guidance scale as defined in [Classifier-Free Diffusion - Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. - of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting - `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to - the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): - The prompt or prompts not to guide the image generation. If not defined, one has to pass - `negative_prompt_embeds`. instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` - is less than `1`). - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - eta (`float`, *optional*, defaults to 0.0): - Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only - applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator`, *optional*): - One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) - to make generation deterministic. - latents (`torch.Tensor`, *optional*): - Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image - generation. Can be used to tweak the same generation with different prompts. If not provided, a latents - tensor will be generated by sampling using the supplied random `generator`. - prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not - provided, text embeddings will be generated from `prompt` input argument. - negative_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input - argument. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between - [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a - plain tuple. - callback (`Callable`, *optional*): - A function that will be called every `callback_steps` steps during inference. The function will be - called with the following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`. - callback_steps (`int`, *optional*, defaults to 1): - The frequency at which the `callback` function will be called. If not specified, the callback will be - called at every step. - use_karras_sigmas (`bool`, *optional*, defaults to `False`): - Use karras sigmas. For example, specifying `sample_dpmpp_2m` to `set_scheduler` will be equivalent to - `DPM++2M` in stable-diffusion-webui. On top of that, setting this option to True will make it `DPM++2M - Karras`. - noise_sampler_seed (`int`, *optional*, defaults to `None`): - The random seed to use for the noise sampler. If `None`, a random seed will be generated. - clip_skip (`int`, *optional*): - Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that - the output of the pre-final layer will be used for computing the prompt embeddings. - Returns: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple. - When returning a tuple, the first element is a list with the generated images, and the second element is a - list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work" - (nsfw) content, according to the `safety_checker`. - """ - # 0. Default height and width to unet - height = height or self.unet.config.sample_size * self.vae_scale_factor - width = width or self.unet.config.sample_size * self.vae_scale_factor - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds - ) - - # 2. Define call parameters - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - device = self._execution_device - # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) - # of the Imagen paper: https://huggingface.co/papers/2205.11487 . `guidance_scale = 1` - # corresponds to doing no classifier free guidance. - do_classifier_free_guidance = True - if guidance_scale <= 1.0: - raise ValueError("has to use guidance_scale") - - # 3. Encode input prompt - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt, - device, - num_images_per_prompt, - do_classifier_free_guidance, - negative_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - clip_skip=clip_skip, - ) - # For classifier free guidance, we need to do two forward passes. - # Here we concatenate the unconditional and text embeddings into a single batch - # to avoid doing two forward passes - if do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=prompt_embeds.device) - - # 5. Prepare sigmas - if use_karras_sigmas: - sigma_min: float = self.k_diffusion_model.sigmas[0].item() - sigma_max: float = self.k_diffusion_model.sigmas[-1].item() - sigmas = get_sigmas_karras(n=num_inference_steps, sigma_min=sigma_min, sigma_max=sigma_max) - else: - sigmas = self.scheduler.sigmas - sigmas = sigmas.to(device) - sigmas = sigmas.to(prompt_embeds.dtype) - - # 6. Prepare latent variables - num_channels_latents = self.unet.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - latents = latents * sigmas[0] - self.k_diffusion_model.sigmas = self.k_diffusion_model.sigmas.to(latents.device) - self.k_diffusion_model.log_sigmas = self.k_diffusion_model.log_sigmas.to(latents.device) - - # 7. Define model function - def model_fn(x, t): - latent_model_input = torch.cat([x] * 2) - t = torch.cat([t] * 2) - - noise_pred = self.k_diffusion_model(latent_model_input, t, cond=prompt_embeds) - - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - return noise_pred - - # 8. Run k-diffusion solver - sampler_kwargs = {} - - if "noise_sampler" in inspect.signature(self.sampler).parameters: - min_sigma, max_sigma = sigmas[sigmas > 0].min(), sigmas.max() - noise_sampler = BrownianTreeNoiseSampler(latents, min_sigma, max_sigma, noise_sampler_seed) - sampler_kwargs["noise_sampler"] = noise_sampler - - if "generator" in inspect.signature(self.sampler).parameters: - sampler_kwargs["generator"] = generator - - latents = self.sampler(model_fn, latents, sigmas, **sampler_kwargs) - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype) - else: - image = latents - has_nsfw_concept = None - - if has_nsfw_concept is None: - do_denormalize = [True] * image.shape[0] - else: - do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] - - image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (image, has_nsfw_concept) - - return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) diff --git a/src/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_xl_k_diffusion.py b/src/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_xl_k_diffusion.py deleted file mode 100644 index f9a8abfcc568..000000000000 --- a/src/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_xl_k_diffusion.py +++ /dev/null @@ -1,889 +0,0 @@ -# Copyright 2025 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import importlib -import inspect -from typing import List, Optional, Tuple, Union - -import torch -from k_diffusion.external import CompVisDenoiser, CompVisVDenoiser -from k_diffusion.sampling import BrownianTreeNoiseSampler, get_sigmas_karras -from transformers import ( - CLIPTextModel, - CLIPTextModelWithProjection, - CLIPTokenizer, -) - -from ...image_processor import VaeImageProcessor -from ...loaders import ( - FromSingleFileMixin, - IPAdapterMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, -) -from ...models import AutoencoderKL, UNet2DConditionModel -from ...models.lora import adjust_lora_scale_text_encoder -from ...schedulers import KarrasDiffusionSchedulers, LMSDiscreteScheduler -from ...utils import ( - USE_PEFT_BACKEND, - deprecate, - logging, - replace_example_docstring, - scale_lora_layers, - unscale_lora_layers, -) -from ...utils.torch_utils import randn_tensor -from ..pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin -from ..stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import StableDiffusionXLKDiffusionPipeline - - >>> pipe = StableDiffusionXLKDiffusionPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - >>> pipe.set_scheduler("sample_dpmpp_2m_sde") - - >>> prompt = "a photo of an astronaut riding a horse on mars" - >>> image = pipe(prompt).images[0] - ``` -""" - - -# Copied from diffusers.pipelines.stable_diffusion_k_diffusion.pipeline_stable_diffusion_k_diffusion.ModelWrapper -class ModelWrapper: - def __init__(self, model, alphas_cumprod): - self.model = model - self.alphas_cumprod = alphas_cumprod - - def apply_model(self, *args, **kwargs): - if len(args) == 3: - encoder_hidden_states = args[-1] - args = args[:2] - if kwargs.get("cond", None) is not None: - encoder_hidden_states = kwargs.pop("cond") - return self.model(*args, encoder_hidden_states=encoder_hidden_states, **kwargs).sample - - -class StableDiffusionXLKDiffusionPipeline( - DeprecatedPipelineMixin, - DiffusionPipeline, - StableDiffusionMixin, - FromSingleFileMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, - IPAdapterMixin, -): - _last_supported_version = "0.33.1" - - r""" - Pipeline for text-to-image generation using Stable Diffusion XL and k-diffusion. - - This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the - library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.) - - The pipeline also inherits the following loading methods: - - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings - - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - - [`~loaders.StableDiffusionXLLoraLoaderMixin.load_lora_weights`] for loading LoRA weights - - [`~loaders.StableDiffusionXLLoraLoaderMixin.save_lora_weights`] for saving LoRA weights - - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`CLIPTextModel`]): - Frozen text-encoder. Stable Diffusion XL uses the text portion of - [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically - the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant. - text_encoder_2 ([` CLIPTextModelWithProjection`]): - Second frozen text-encoder. Stable Diffusion XL uses the text and pool portion of - [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection), - specifically the - [laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k) - variant. - tokenizer (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). - tokenizer_2 (`CLIPTokenizer`): - Second Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). - unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of - [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. - force_zeros_for_empty_prompt (`bool`, *optional*, defaults to `"True"`): - Whether the negative prompt embeddings shall be forced to always be set to 0. Also see the config of - `stabilityai/stable-diffusion-xl-base-1-0`. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->unet->vae" - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "feature_extractor", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - force_zeros_for_empty_prompt: bool = True, - ): - super().__init__() - - # get correct sigmas from LMS - scheduler = LMSDiscreteScheduler.from_config(scheduler.config) - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - unet=unet, - scheduler=scheduler, - ) - self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) - - self.default_sample_size = ( - self.unet.config.sample_size - if hasattr(self, "unet") and self.unet is not None and hasattr(self.unet.config, "sample_size") - else 128 - ) - - model = ModelWrapper(unet, scheduler.alphas_cumprod) - if scheduler.config.prediction_type == "v_prediction": - self.k_diffusion_model = CompVisVDenoiser(model) - else: - self.k_diffusion_model = CompVisDenoiser(model) - - # Copied from diffusers.pipelines.stable_diffusion_k_diffusion.pipeline_stable_diffusion_k_diffusion.StableDiffusionKDiffusionPipeline.set_scheduler - def set_scheduler(self, scheduler_type: str): - library = importlib.import_module("k_diffusion") - sampling = getattr(library, "sampling") - try: - self.sampler = getattr(sampling, scheduler_type) - except Exception: - valid_samplers = [] - for s in dir(sampling): - if "sample_" in s: - valid_samplers.append(s) - - raise ValueError(f"Invalid scheduler type {scheduler_type}. Please choose one of {valid_samplers}.") - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.encode_prompt - def encode_prompt( - self, - prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, - num_images_per_prompt: int = 1, - do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` or `List[str]`, *optional*): - prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): - The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is - used in both text-encoders - device: (`torch.device`): - torch device - num_images_per_prompt (`int`): - number of images that should be generated per prompt - do_classifier_free_guidance (`bool`): - whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): - The prompt or prompts not to guide the image generation. If not defined, one has to pass - `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is - less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): - The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and - `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders - prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not - provided, text embeddings will be generated from `prompt` input argument. - negative_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input - argument. - pooled_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. - If not provided, pooled text embeddings will be generated from `prompt` input argument. - negative_pooled_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` - input argument. - lora_scale (`float`, *optional*): - A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded. - clip_skip (`int`, *optional*): - Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that - the output of the pre-final layer will be used for computing the prompt embeddings. - """ - device = device or self._execution_device - - # set lora scale so that monkey patched LoRA - # function of text encoder can correctly access it - if lora_scale is not None and isinstance(self, StableDiffusionXLLoraLoaderMixin): - self._lora_scale = lora_scale - - # dynamically adjust the LoRA scale - if self.text_encoder is not None: - if not USE_PEFT_BACKEND: - adjust_lora_scale_text_encoder(self.text_encoder, lora_scale) - else: - scale_lora_layers(self.text_encoder, lora_scale) - - if self.text_encoder_2 is not None: - if not USE_PEFT_BACKEND: - adjust_lora_scale_text_encoder(self.text_encoder_2, lora_scale) - else: - scale_lora_layers(self.text_encoder_2, lora_scale) - - prompt = [prompt] if isinstance(prompt, str) else prompt - - if prompt is not None: - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - # Define tokenizers and text encoders - tokenizers = [self.tokenizer, self.tokenizer_2] if self.tokenizer is not None else [self.tokenizer_2] - text_encoders = ( - [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [self.text_encoder_2] - ) - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - # textual inversion: process multi-vector tokens if necessary - prompt_embeds_list = [] - prompts = [prompt, prompt_2] - for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders): - if isinstance(self, TextualInversionLoaderMixin): - prompt = self.maybe_convert_prompt(prompt, tokenizer) - - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - - text_input_ids = text_inputs.input_ids - untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids - - if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal( - text_input_ids, untruncated_ids - ): - removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because CLIP can only handle sequences up to" - f" {tokenizer.model_max_length} tokens: {removed_text}" - ) - - prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True) - - # We are only ALWAYS interested in the pooled output of the final text encoder - if pooled_prompt_embeds is None and prompt_embeds[0].ndim == 2: - pooled_prompt_embeds = prompt_embeds[0] - - if clip_skip is None: - prompt_embeds = prompt_embeds.hidden_states[-2] - else: - # "2" because SDXL always indexes from the penultimate layer. - prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)] - - prompt_embeds_list.append(prompt_embeds) - - prompt_embeds = torch.concat(prompt_embeds_list, dim=-1) - - # get unconditional embeddings for classifier free guidance - zero_out_negative_prompt = negative_prompt is None and self.config.force_zeros_for_empty_prompt - if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt: - negative_prompt_embeds = torch.zeros_like(prompt_embeds) - negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds) - elif do_classifier_free_guidance and negative_prompt_embeds is None: - negative_prompt = negative_prompt or "" - negative_prompt_2 = negative_prompt_2 or negative_prompt - - # normalize str to list - negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt - negative_prompt_2 = ( - batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 - ) - - uncond_tokens: List[str] - if prompt is not None and type(prompt) is not type(negative_prompt): - raise TypeError( - f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" - f" {type(prompt)}." - ) - elif batch_size != len(negative_prompt): - raise ValueError( - f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" - f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" - " the batch size of `prompt`." - ) - else: - uncond_tokens = [negative_prompt, negative_prompt_2] - - negative_prompt_embeds_list = [] - for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders): - if isinstance(self, TextualInversionLoaderMixin): - negative_prompt = self.maybe_convert_prompt(negative_prompt, tokenizer) - - max_length = prompt_embeds.shape[1] - uncond_input = tokenizer( - negative_prompt, - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="pt", - ) - - negative_prompt_embeds = text_encoder( - uncond_input.input_ids.to(device), - output_hidden_states=True, - ) - - # We are only ALWAYS interested in the pooled output of the final text encoder - if negative_pooled_prompt_embeds is None and negative_prompt_embeds[0].ndim == 2: - negative_pooled_prompt_embeds = negative_prompt_embeds[0] - negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2] - - negative_prompt_embeds_list.append(negative_prompt_embeds) - - negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1) - - if self.text_encoder_2 is not None: - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device) - else: - prompt_embeds = prompt_embeds.to(dtype=self.unet.dtype, device=device) - - bs_embed, seq_len, _ = prompt_embeds.shape - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) - prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1) - - if do_classifier_free_guidance: - # duplicate unconditional embeddings for each generation per prompt, using mps friendly method - seq_len = negative_prompt_embeds.shape[1] - - if self.text_encoder_2 is not None: - negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device) - else: - negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.unet.dtype, device=device) - - negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view( - bs_embed * num_images_per_prompt, -1 - ) - if do_classifier_free_guidance: - negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt).view( - bs_embed * num_images_per_prompt, -1 - ) - - if self.text_encoder is not None: - if isinstance(self, StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND: - # Retrieve the original scale by scaling back the LoRA layers - unscale_lora_layers(self.text_encoder, lora_scale) - - if self.text_encoder_2 is not None: - if isinstance(self, StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND: - # Retrieve the original scale by scaling back the LoRA layers - unscale_lora_layers(self.text_encoder_2, lora_scale) - - return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds - - def check_inputs( - self, - prompt, - prompt_2, - height, - width, - negative_prompt=None, - negative_prompt_2=None, - prompt_embeds=None, - negative_prompt_embeds=None, - pooled_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - ): - if height % 8 != 0 or width % 8 != 0: - raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") - - if prompt is not None and prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" - " only forward one of the two." - ) - elif prompt_2 is not None and prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to" - " only forward one of the two." - ) - elif prompt is None and prompt_embeds is None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." - ) - elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)): - raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}") - - if negative_prompt is not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - elif negative_prompt_2 is not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - if prompt_embeds is not None and negative_prompt_embeds is not None: - if prompt_embeds.shape != negative_prompt_embeds.shape: - raise ValueError( - "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" - f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" - f" {negative_prompt_embeds.shape}." - ) - - if prompt_embeds is not None and pooled_prompt_embeds is None: - raise ValueError( - "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`." - ) - - if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None: - raise ValueError( - "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`." - ) - - def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None): - shape = ( - batch_size, - num_channels_latents, - int(height) // self.vae_scale_factor, - int(width) // self.vae_scale_factor, - ) - if isinstance(generator, list) and len(generator) != batch_size: - raise ValueError( - f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" - f" size of {batch_size}. Make sure the batch size matches the length of the generators." - ) - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - latents = latents.to(device) - - return latents - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline._get_add_time_ids - def _get_add_time_ids( - self, original_size, crops_coords_top_left, target_size, dtype, text_encoder_projection_dim=None - ): - add_time_ids = list(original_size + crops_coords_top_left + target_size) - - passed_add_embed_dim = ( - self.unet.config.addition_time_embed_dim * len(add_time_ids) + text_encoder_projection_dim - ) - expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features - - if expected_add_embed_dim != passed_add_embed_dim: - raise ValueError( - f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`." - ) - - add_time_ids = torch.tensor([add_time_ids], dtype=dtype) - return add_time_ids - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.upcast_vae - def upcast_vae(self): - deprecate( - "upcast_vae", - "1.0.0", - "`upcast_vae` is deprecated. Please use `pipe.vae.to(torch.float32)`. For more details, please refer to: https://github.com/huggingface/diffusers/pull/12619#issue-3606633695.", - ) - self.vae.to(dtype=torch.float32) - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def clip_skip(self): - return self._clip_skip - - # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) - # of the Imagen paper: https://huggingface.co/papers/2205.11487 . `guidance_scale = 1` - # corresponds to doing no classifier free guidance. - @property - def do_classifier_free_guidance(self): - return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", - return_dict: bool = True, - original_size: Optional[Tuple[int, int]] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Optional[Tuple[int, int]] = None, - negative_original_size: Optional[Tuple[int, int]] = None, - negative_crops_coords_top_left: Tuple[int, int] = (0, 0), - negative_target_size: Optional[Tuple[int, int]] = None, - use_karras_sigmas: Optional[bool] = False, - noise_sampler_seed: Optional[int] = None, - clip_skip: Optional[int] = None, - ): - r""" - Function invoked when calling the pipeline for generation. - - Args: - prompt (`str` or `List[str]`, *optional*): - The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. - instead. - prompt_2 (`str` or `List[str]`, *optional*): - The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is - used in both text-encoders - height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): - The height in pixels of the generated image. This is set to 1024 by default for the best results. - Anything below 512 pixels won't work well for - [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) - and checkpoints that are not specifically fine-tuned on low resolutions. - width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): - The width in pixels of the generated image. This is set to 1024 by default for the best results. - Anything below 512 pixels won't work well for - [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) - and checkpoints that are not specifically fine-tuned on low resolutions. - num_inference_steps (`int`, *optional*, defaults to 50): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. - guidance_scale (`float`, *optional*, defaults to 5.0): - Guidance scale as defined in [Classifier-Free Diffusion - Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. - of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting - `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to - the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): - The prompt or prompts not to guide the image generation. If not defined, one has to pass - `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is - less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): - The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and - `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): - One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) - to make generation deterministic. - latents (`torch.Tensor`, *optional*): - Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image - generation. Can be used to tweak the same generation with different prompts. If not provided, a latents - tensor will be generated by sampling using the supplied random `generator`. - prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not - provided, text embeddings will be generated from `prompt` input argument. - negative_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input - argument. - pooled_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. - If not provided, pooled text embeddings will be generated from `prompt` input argument. - negative_pooled_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` - input argument. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between - [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead - of a plain tuple. - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): - If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. - `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as - explained in section 2.2 of - [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): - `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position - `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting - `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of - [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): - For most cases, `target_size` should be set to the desired height and width of the generated image. If - not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in - section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): - To negatively condition the generation process based on a specific image resolution. Part of SDXL's - micro-conditioning as explained in section 2.2 of - [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more - information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): - To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's - micro-conditioning as explained in section 2.2 of - [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more - information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): - To negatively condition the generation process based on a target image resolution. It should be as same - as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of - [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more - information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - - Examples: - - Returns: - [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, the first element is a list with the generated images. - """ - - # 0. Default height and width to unet - height = height or self.default_sample_size * self.vae_scale_factor - width = width or self.default_sample_size * self.vae_scale_factor - - original_size = original_size or (height, width) - target_size = target_size or (height, width) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - prompt_2, - height, - width, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) - - if guidance_scale <= 1.0: - raise ValueError("has to use guidance_scale") - - self._guidance_scale = guidance_scale - self._clip_skip = clip_skip - - # 2. Define call parameters - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - device = self._execution_device - - # 3. Encode input prompt - lora_scale = None - - ( - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - device=device, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - lora_scale=lora_scale, - clip_skip=self.clip_skip, - ) - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=prompt_embeds.device) - - # 5. Prepare sigmas - if use_karras_sigmas: - sigma_min: float = self.k_diffusion_model.sigmas[0].item() - sigma_max: float = self.k_diffusion_model.sigmas[-1].item() - sigmas = get_sigmas_karras(n=num_inference_steps, sigma_min=sigma_min, sigma_max=sigma_max) - else: - sigmas = self.scheduler.sigmas - sigmas = sigmas.to(dtype=prompt_embeds.dtype, device=device) - - # 6. Prepare latent variables - num_channels_latents = self.unet.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - latents = latents * sigmas[0] - - self.k_diffusion_model.sigmas = self.k_diffusion_model.sigmas.to(latents.device) - self.k_diffusion_model.log_sigmas = self.k_diffusion_model.log_sigmas.to(latents.device) - - # 7. Prepare added time ids & embeddings - add_text_embeds = pooled_prompt_embeds - if self.text_encoder_2 is None: - text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1]) - else: - text_encoder_projection_dim = self.text_encoder_2.config.projection_dim - - add_time_ids = self._get_add_time_ids( - original_size, - crops_coords_top_left, - target_size, - dtype=prompt_embeds.dtype, - text_encoder_projection_dim=text_encoder_projection_dim, - ) - if negative_original_size is not None and negative_target_size is not None: - negative_add_time_ids = self._get_add_time_ids( - negative_original_size, - negative_crops_coords_top_left, - negative_target_size, - dtype=prompt_embeds.dtype, - text_encoder_projection_dim=text_encoder_projection_dim, - ) - else: - negative_add_time_ids = add_time_ids - - if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) - add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0) - add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0) - - prompt_embeds = prompt_embeds.to(device) - add_text_embeds = add_text_embeds.to(device) - add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1) - - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - - # 8. Optionally get Guidance Scale Embedding - timestep_cond = None - if self.unet.config.time_cond_proj_dim is not None: - guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt) - timestep_cond = self.get_guidance_scale_embedding( - guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim - ).to(device=device, dtype=latents.dtype) - - # 9. Define model function - def model_fn(x, t): - latent_model_input = torch.cat([x] * 2) - t = torch.cat([t] * 2) - - noise_pred = self.k_diffusion_model( - latent_model_input, - t, - cond=prompt_embeds, - timestep_cond=timestep_cond, - added_cond_kwargs=added_cond_kwargs, - ) - - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - return noise_pred - - # 10. Run k-diffusion solver - sampler_kwargs = {} - - if "noise_sampler" in inspect.signature(self.sampler).parameters: - min_sigma, max_sigma = sigmas[sigmas > 0].min(), sigmas.max() - noise_sampler = BrownianTreeNoiseSampler(latents, min_sigma, max_sigma, noise_sampler_seed) - sampler_kwargs["noise_sampler"] = noise_sampler - - if "generator" in inspect.signature(self.sampler).parameters: - sampler_kwargs["generator"] = generator - - latents = self.sampler(model_fn, latents, sigmas, **sampler_kwargs) - - if not output_type == "latent": - # make sure the VAE is in float32 mode, as it overflows in float16 - needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast - - if needs_upcasting: - self.upcast_vae() - latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype) - - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - - # cast back to fp16 if needed - if needs_upcasting: - self.vae.to(dtype=torch.float16) - else: - image = latents - - if not output_type == "latent": - image = self.image_processor.postprocess(image, output_type=output_type) - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (image,) - - return StableDiffusionXLPipelineOutput(images=image) diff --git a/src/diffusers/pipelines/stable_diffusion_ldm3d/pipeline_stable_diffusion_ldm3d.py b/src/diffusers/pipelines/stable_diffusion_ldm3d/pipeline_stable_diffusion_ldm3d.py index c32121c88c9b..6de144aa7e8b 100644 --- a/src/diffusers/pipelines/stable_diffusion_ldm3d/pipeline_stable_diffusion_ldm3d.py +++ b/src/diffusers/pipelines/stable_diffusion_ldm3d/pipeline_stable_diffusion_ldm3d.py @@ -14,7 +14,7 @@ import inspect from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -98,10 +98,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -116,15 +116,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -161,20 +161,20 @@ class LDM3DPipelineOutput(BaseOutput): Output class for Stable Diffusion pipelines. Args: - rgb (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, + rgb (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, num_channels)`. - depth (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, + depth (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, num_channels)`. - nsfw_content_detected (`List[bool]`) - List indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or + nsfw_content_detected (`list[bool]`) + list indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or `None` if safety checking could not be performed. """ - rgb: Union[List[PIL.Image.Image], np.ndarray] - depth: Union[List[PIL.Image.Image], np.ndarray] - nsfw_content_detected: Optional[List[bool]] + rgb: list[PIL.Image.Image] | np.ndarray + depth: list[PIL.Image.Image] | np.ndarray + nsfw_content_detected: list[bool] | None class StableDiffusionLDM3DPipeline( @@ -235,7 +235,7 @@ def __init__( scheduler: KarrasDiffusionSchedulers, safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, - image_encoder: Optional[CLIPVisionModelWithProjection], + image_encoder: CLIPVisionModelWithProjection | None, requires_safety_checker: bool = True, ): super().__init__() @@ -278,9 +278,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -311,16 +311,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -328,7 +328,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -427,7 +427,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -748,36 +748,36 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 49, - timesteps: List[int] = None, - sigmas: List[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, guidance_rescale: float = 0.0, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -786,18 +786,18 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. guidance_scale (`float`, *optional*, defaults to 5.0): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -805,7 +805,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -820,7 +820,7 @@ def __call__( not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -841,7 +841,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/stable_diffusion_panorama/pipeline_stable_diffusion_panorama.py b/src/diffusers/pipelines/stable_diffusion_panorama/pipeline_stable_diffusion_panorama.py index 295095947a12..259fbd933430 100644 --- a/src/diffusers/pipelines/stable_diffusion_panorama/pipeline_stable_diffusion_panorama.py +++ b/src/diffusers/pipelines/stable_diffusion_panorama/pipeline_stable_diffusion_panorama.py @@ -13,7 +13,7 @@ import copy import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection @@ -98,10 +98,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -116,15 +116,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -211,7 +211,7 @@ def __init__( scheduler: DDIMScheduler, safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, - image_encoder: Optional[CLIPVisionModelWithProjection] = None, + image_encoder: CLIPVisionModelWithProjection | None = None, requires_safety_checker: bool = True, ): super().__init__() @@ -254,9 +254,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -287,16 +287,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -304,7 +304,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -403,7 +403,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -735,7 +735,7 @@ def get_views( window_size: int = 64, stride: int = 8, circular_padding: bool = False, - ) -> List[Tuple[int, int, int, int]]: + ) -> list[tuple[int, int, int, int]]: """ Generates a list of views based on the given parameters. Here, we define the mappings F_i (see Eq. 7 in the MultiDiffusion paper https://huggingface.co/papers/2302.08113). If panorama's height/width < window_size, @@ -749,7 +749,7 @@ def get_views( circular_padding (bool, optional): Whether to apply circular padding. Defaults to False. Returns: - List[Tuple[int, int, int, int]]: A list of tuples representing the views. Each tuple contains four integers + list[tuple[int, int, int, int]]: A list of tuples representing the views. Each tuple contains four integers representing the start and end coordinates of the window in the panorama. """ @@ -802,37 +802,37 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = 512, - width: Optional[int] = 2048, + prompt: str | list[str] = None, + height: int | None = 512, + width: int | None = 2048, num_inference_steps: int = 50, - timesteps: List[int] = None, + timesteps: list[int] = None, guidance_scale: float = 7.5, view_batch_size: int = 1, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, guidance_rescale: float = 0.0, circular_padding: bool = False, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs: Any, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to 512): The height in pixels of the generated image. @@ -842,7 +842,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): The timesteps at which to generate the images. If not specified, then the default timestep spacing strategy of the scheduler is used. guidance_scale (`float`, *optional*, defaults to 7.5): @@ -851,7 +851,7 @@ def __call__( view_batch_size (`int`, *optional*, defaults to 1): The batch size to denoise split views. For some GPUs with high performance, higher view batch size can speedup the generation and increase the VRAM usage. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -859,7 +859,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -874,7 +874,7 @@ def __call__( not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -902,7 +902,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List[str]`, *optional*): + callback_on_step_end_tensor_inputs (`list[str]`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1000,7 +1000,13 @@ def __call__( prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, timesteps + ) # 5. Prepare latent variables num_channels_latents = self.unet.config.in_channels diff --git a/src/diffusers/pipelines/stable_diffusion_safe/__init__.py b/src/diffusers/pipelines/stable_diffusion_safe/__init__.py index b432b9418c46..b35015a9f729 100644 --- a/src/diffusers/pipelines/stable_diffusion_safe/__init__.py +++ b/src/diffusers/pipelines/stable_diffusion_safe/__init__.py @@ -1,6 +1,6 @@ from dataclasses import dataclass from enum import Enum -from typing import TYPE_CHECKING, List, Optional, Union +from typing import TYPE_CHECKING import numpy as np import PIL diff --git a/src/diffusers/pipelines/stable_diffusion_safe/pipeline_output.py b/src/diffusers/pipelines/stable_diffusion_safe/pipeline_output.py index 69a064d6638d..6b784bb0e102 100644 --- a/src/diffusers/pipelines/stable_diffusion_safe/pipeline_output.py +++ b/src/diffusers/pipelines/stable_diffusion_safe/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Optional, Union import numpy as np import PIL.Image @@ -15,20 +14,20 @@ class StableDiffusionSafePipelineOutput(BaseOutput): Output class for Safe Stable Diffusion pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. - nsfw_content_detected (`List[bool]`) - List of flags denoting whether the corresponding generated image likely represents "not-safe-for-work" + nsfw_content_detected (`list[bool]`) + list of flags denoting whether the corresponding generated image likely represents "not-safe-for-work" (nsfw) content, or `None` if safety checking could not be performed. - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images that were flagged by the safety checker any may contain "not-safe-for-work" + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images that were flagged by the safety checker any may contain "not-safe-for-work" (nsfw) content, or `None` if no safety check was performed or no images were flagged. applied_safety_concept (`str`) The safety concept that was applied for safety guidance, or `None` if safety guidance was disabled """ - images: Union[List[PIL.Image.Image], np.ndarray] - nsfw_content_detected: Optional[List[bool]] - unsafe_images: Optional[Union[List[PIL.Image.Image], np.ndarray]] - applied_safety_concept: Optional[str] + images: list[PIL.Image.Image] | np.ndarray + nsfw_content_detected: list[bool] | None + unsafe_images: list[PIL.Image.Image] | np.ndarray | None + applied_safety_concept: str | None diff --git a/src/diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py b/src/diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py index d334107b0703..26bb5128ba9b 100644 --- a/src/diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py +++ b/src/diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py @@ -1,6 +1,6 @@ import inspect import warnings -from typing import Callable, List, Optional, Union +from typing import Callable import numpy as np import torch @@ -73,11 +73,11 @@ def __init__( scheduler: KarrasDiffusionSchedulers, safety_checker: SafeStableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, - image_encoder: Optional[CLIPVisionModelWithProjection] = None, + image_encoder: CLIPVisionModelWithProjection | None = None, requires_safety_checker: bool = True, ): super().__init__() - safety_concept: Optional[str] = ( + safety_concept: str | None = ( "an image showing hate, harassment, violence, suffering, humiliation, harm, suicide, sexual, nudity," " bodily fluids, blood, obscene gestures, illegal activity, drug use, theft, vandalism, weapons, child" " abuse, brutality, cruelty" @@ -199,7 +199,7 @@ def _encode_prompt( Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): prompt to be encoded device: (`torch.device`): torch device @@ -207,7 +207,7 @@ def _encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`): + negative_prompt (`str` or `list[str]`): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). """ @@ -248,7 +248,7 @@ def _encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -374,7 +374,6 @@ def prepare_extra_step_kwargs(self, generator, eta): extra_step_kwargs["generator"] = generator return extra_step_kwargs - # Copied from diffusers.pipelines.stable_diffusion_k_diffusion.pipeline_stable_diffusion_k_diffusion.StableDiffusionKDiffusionPipeline.check_inputs def check_inputs( self, prompt, @@ -520,32 +519,32 @@ def encode_image(self, image, device, num_images_per_prompt, output_hidden_state @torch.no_grad() def __call__( self, - prompt: Union[str, List[str]], - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str], + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - sld_guidance_scale: Optional[float] = 1000, - sld_warmup_steps: Optional[int] = 10, - sld_threshold: Optional[float] = 0.01, - sld_momentum_scale: Optional[float] = 0.3, - sld_mom_beta: Optional[float] = 0.4, + sld_guidance_scale: float | None = 1000, + sld_warmup_steps: int | None = 10, + sld_threshold: float | None = 0.01, + sld_momentum_scale: float | None = 0.3, + sld_mom_beta: float | None = 0.4, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -557,7 +556,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -565,7 +564,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/stable_diffusion_sag/pipeline_stable_diffusion_sag.py b/src/diffusers/pipelines/stable_diffusion_sag/pipeline_stable_diffusion_sag.py index 48add535a81d..3cf604911f0b 100644 --- a/src/diffusers/pipelines/stable_diffusion_sag/pipeline_stable_diffusion_sag.py +++ b/src/diffusers/pipelines/stable_diffusion_sag/pipeline_stable_diffusion_sag.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch import torch.nn.functional as F @@ -155,7 +155,7 @@ def __init__( scheduler: KarrasDiffusionSchedulers, safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, - image_encoder: Optional[CLIPVisionModelWithProjection] = None, + image_encoder: CLIPVisionModelWithProjection | None = None, requires_safety_checker: bool = True, ): super().__init__() @@ -182,9 +182,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -215,16 +215,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -232,7 +232,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -331,7 +331,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -494,7 +494,6 @@ def prepare_extra_step_kwargs(self, generator, eta): extra_step_kwargs["generator"] = generator return extra_step_kwargs - # Copied from diffusers.pipelines.stable_diffusion_k_diffusion.pipeline_stable_diffusion_k_diffusion.StableDiffusionKDiffusionPipeline.check_inputs def check_inputs( self, prompt, @@ -574,33 +573,33 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, sag_scale: float = 0.75, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: Optional[int] = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, + callback_steps: int | None = 1, + cross_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. @@ -614,7 +613,7 @@ def __call__( `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. sag_scale (`float`, *optional*, defaults to 0.75): Chosen between [0, 1.0] for better quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -622,7 +621,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -637,7 +636,7 @@ def __call__( not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. If not provided, embeddings are computed from the `ip_adapter_image` input argument. output_type (`str`, *optional*, defaults to `"pil"`): diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py index 3227fd9a08a4..f80abc7073ed 100644 --- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py +++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py @@ -12,8 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + from functools import partial -from typing import Dict, List, Optional, Union import jax import jax.numpy as jnp @@ -48,9 +49,7 @@ def __init__( tokenizer: CLIPTokenizer, tokenizer_2: CLIPTokenizer, unet: FlaxUNet2DConditionModel, - scheduler: Union[ - FlaxDDIMScheduler, FlaxPNDMScheduler, FlaxLMSDiscreteScheduler, FlaxDPMSolverMultistepScheduler - ], + scheduler: FlaxDDIMScheduler | FlaxPNDMScheduler | FlaxLMSDiscreteScheduler | FlaxDPMSolverMultistepScheduler, dtype: jnp.dtype = jnp.float32, ): super().__init__() @@ -67,7 +66,7 @@ def __init__( ) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - def prepare_inputs(self, prompt: Union[str, List[str]]): + def prepare_inputs(self, prompt: str | list[str]): if not isinstance(prompt, (str, list)): raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") @@ -88,12 +87,12 @@ def prepare_inputs(self, prompt: Union[str, List[str]]): def __call__( self, prompt_ids: jax.Array, - params: Union[Dict, FrozenDict], + params: dict | FrozenDict, prng_seed: jax.Array, num_inference_steps: int = 50, - guidance_scale: Union[float, jax.Array] = 7.5, - height: Optional[int] = None, - width: Optional[int] = None, + guidance_scale: float | jax.Array = 7.5, + height: int | None = None, + width: int | None = None, latents: jnp.array = None, neg_prompt_ids: jnp.array = None, return_dict: bool = True, @@ -169,14 +168,14 @@ def _get_add_time_ids(self, original_size, crops_coords_top_left, target_size, b def _generate( self, prompt_ids: jnp.array, - params: Union[Dict, FrozenDict], + params: dict | FrozenDict, prng_seed: jax.Array, num_inference_steps: int, height: int, width: int, guidance_scale: float, - latents: Optional[jnp.array] = None, - neg_prompt_ids: Optional[jnp.array] = None, + latents: jnp.array | None = None, + neg_prompt_ids: jnp.array | None = None, return_latents=False, ): if height % 8 != 0 or width % 8 != 0: diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_output.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_output.py index 0783f44486ee..6617e7cd058c 100644 --- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_output.py +++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -13,12 +12,12 @@ class StableDiffusionXLPipelineOutput(BaseOutput): Output class for Stable Diffusion pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image] | np.ndarray if is_flax_available(): diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py index e969d2a21a99..2f6b105702e8 100644 --- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py +++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import torch from transformers import ( @@ -110,10 +110,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -128,15 +128,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -248,7 +248,7 @@ def __init__( image_encoder: CLIPVisionModelWithProjection = None, feature_extractor: CLIPImageProcessor = None, force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, + add_watermarker: bool | None = None, ): super().__init__() @@ -283,26 +283,26 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, + prompt_2: str | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders device: (`torch.device`): @@ -311,11 +311,11 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -433,7 +433,7 @@ def encode_prompt( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) - uncond_tokens: List[str] + uncond_tokens: list[str] if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -822,52 +822,50 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_end: Optional[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, + denoising_end: float | None = None, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, guidance_rescale: float = 0.0, - original_size: Optional[Tuple[int, int]] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Optional[Tuple[int, int]] = None, - negative_original_size: Optional[Tuple[int, int]] = None, - negative_crops_coords_top_left: Tuple[int, int] = (0, 0), - negative_target_size: Optional[Tuple[int, int]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + original_size: tuple[int, int] | None = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] | None = None, + negative_original_size: tuple[int, int] | None = None, + negative_crops_coords_top_left: tuple[int, int] = (0, 0), + negative_target_size: tuple[int, int] | None = None, + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): @@ -883,11 +881,11 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -904,11 +902,11 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -916,7 +914,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -938,7 +936,7 @@ def __call__( weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -959,31 +957,31 @@ def __call__( [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when using zero terminal SNR. - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a specific image resolution. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a target image resolution. It should be as same as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more @@ -993,7 +991,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1094,8 +1092,12 @@ def __call__( ) # 4. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) # 5. Prepare latent variables diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py index 8d1da8dc102c..19ccfab3de0a 100644 --- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py +++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import PIL.Image import torch @@ -113,7 +113,7 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -128,10 +128,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -146,15 +146,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -270,7 +270,7 @@ def __init__( feature_extractor: CLIPImageProcessor = None, requires_aesthetics_score: bool = False, force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, + add_watermarker: bool | None = None, ): super().__init__() @@ -301,26 +301,26 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, + prompt_2: str | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders device: (`torch.device`): @@ -329,11 +329,11 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -451,7 +451,7 @@ def encode_prompt( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) - uncond_tokens: List[str] + uncond_tokens: list[str] if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -974,58 +974,56 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, image: PipelineImageInput = None, strength: float = 0.3, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_start: Optional[float] = None, - denoising_end: Optional[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, + denoising_start: float | None = None, + denoising_end: float | None = None, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, guidance_rescale: float = 0.0, - original_size: Tuple[int, int] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Tuple[int, int] = None, - negative_original_size: Optional[Tuple[int, int]] = None, - negative_crops_coords_top_left: Tuple[int, int] = (0, 0), - negative_target_size: Optional[Tuple[int, int]] = None, + original_size: tuple[int, int] = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] = None, + negative_original_size: tuple[int, int] | None = None, + negative_crops_coords_top_left: tuple[int, int] = (0, 0), + negative_target_size: tuple[int, int] | None = None, aesthetic_score: float = 6.0, negative_aesthetic_score: float = 2.5, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders - image (`torch.Tensor` or `PIL.Image.Image` or `np.ndarray` or `List[torch.Tensor]` or `List[PIL.Image.Image]` or `List[np.ndarray]`): + image (`torch.Tensor` or `PIL.Image.Image` or `np.ndarray` or `list[torch.Tensor]` or `list[PIL.Image.Image]` or `list[np.ndarray]`): The image(s) to modify with the pipeline. strength (`float`, *optional*, defaults to 0.3): Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1. `image` @@ -1037,11 +1035,11 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -1066,11 +1064,11 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -1078,7 +1076,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -1100,7 +1098,7 @@ def __call__( weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -1121,31 +1119,31 @@ def __call__( [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when using zero terminal SNR. - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a specific image resolution. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a target image resolution. It should be as same as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more @@ -1166,7 +1164,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1264,8 +1262,12 @@ def __call__( def denoising_value_valid(dnv): return isinstance(dnv, float) and 0 < dnv < 1 + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) timesteps, num_inference_steps = self.get_timesteps( num_inference_steps, diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py index 54a1e311804c..8de7d4f0bb7d 100644 --- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py +++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -140,7 +140,7 @@ def mask_pil_to_torch(mask, height, width): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -155,10 +155,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -173,15 +173,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -300,7 +300,7 @@ def __init__( feature_extractor: CLIPImageProcessor = None, requires_aesthetics_score: bool = False, force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, + add_watermarker: bool | None = None, ): super().__init__() @@ -405,26 +405,26 @@ def prepare_ip_adapter_image_embeds( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, + prompt_2: str | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders device: (`torch.device`): @@ -433,11 +433,11 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -555,7 +555,7 @@ def encode_prompt( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) - uncond_tokens: List[str] + uncond_tokens: list[str] if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -1079,60 +1079,58 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, image: PipelineImageInput = None, mask_image: PipelineImageInput = None, masked_image_latents: torch.Tensor = None, - height: Optional[int] = None, - width: Optional[int] = None, - padding_mask_crop: Optional[int] = None, + height: int | None = None, + width: int | None = None, + padding_mask_crop: int | None = None, strength: float = 0.9999, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_start: Optional[float] = None, - denoising_end: Optional[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, + denoising_start: float | None = None, + denoising_end: float | None = None, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, guidance_rescale: float = 0.0, - original_size: Tuple[int, int] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Tuple[int, int] = None, - negative_original_size: Optional[Tuple[int, int]] = None, - negative_crops_coords_top_left: Tuple[int, int] = (0, 0), - negative_target_size: Optional[Tuple[int, int]] = None, + original_size: tuple[int, int] = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] = None, + negative_original_size: tuple[int, int] | None = None, + negative_crops_coords_top_left: tuple[int, int] = (0, 0), + negative_target_size: tuple[int, int] | None = None, aesthetic_score: float = 6.0, negative_aesthetic_score: float = 2.5, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + clip_skip: int | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders image (`PIL.Image.Image`): @@ -1171,11 +1169,11 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -1200,11 +1198,11 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -1222,7 +1220,7 @@ def __call__( weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -1249,31 +1247,31 @@ def __call__( A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a specific image resolution. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a target image resolution. It should be as same as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more @@ -1294,7 +1292,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -1399,8 +1397,12 @@ def __call__( def denoising_value_valid(dnv): return isinstance(dnv, float) and 0 < dnv < 1 + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) timesteps, num_inference_steps = self.get_timesteps( num_inference_steps, diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py index 5e13362eb3d1..b79119a94a0c 100644 --- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py +++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import PIL.Image import torch @@ -84,7 +84,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -179,8 +179,8 @@ def __init__( unet: UNet2DConditionModel, scheduler: KarrasDiffusionSchedulers, force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, - is_cosxl_edit: Optional[bool] = False, + add_watermarker: bool | None = None, + is_cosxl_edit: bool | None = False, ): super().__init__() @@ -213,25 +213,25 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, + prompt_2: str | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders device: (`torch.device`): @@ -240,11 +240,11 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -350,7 +350,7 @@ def encode_prompt( negative_prompt = negative_prompt or "" negative_prompt_2 = negative_prompt_2 or negative_prompt - uncond_tokens: List[str] + uncond_tokens: list[str] if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -596,46 +596,46 @@ def upcast_vae(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 100, - denoising_end: Optional[float] = None, + denoising_end: float | None = None, guidance_scale: float = 5.0, image_guidance_scale: float = 1.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, guidance_rescale: float = 0.0, - original_size: Tuple[int, int] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Tuple[int, int] = None, + original_size: tuple[int, int] = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] = None, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders - image (`torch.Tensor` or `PIL.Image.Image` or `np.ndarray` or `List[torch.Tensor]` or `List[PIL.Image.Image]` or `List[np.ndarray]`): + image (`torch.Tensor` or `PIL.Image.Image` or `np.ndarray` or `list[torch.Tensor]` or `list[PIL.Image.Image]` or `list[np.ndarray]`): The image(s) to modify with the pipeline. height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): The height in pixels of the generated image. @@ -662,11 +662,11 @@ def __call__( scale is enabled by setting `image_guidance_scale > 1`. Higher image guidance scale encourages to generate images that are closely linked to the source image `image`, usually at the expense of lower image quality. This pipeline requires a value of at least `1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders. num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -674,7 +674,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -717,17 +717,17 @@ def __call__( [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when using zero terminal SNR. - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). diff --git a/src/diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py b/src/diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py index 6d9053faaec8..05877f69d403 100644 --- a/src/diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py +++ b/src/diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py @@ -14,7 +14,7 @@ import inspect from dataclasses import dataclass -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import numpy as np import PIL.Image @@ -73,10 +73,10 @@ def _append_dims(x, target_dims): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -91,15 +91,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -136,12 +136,12 @@ class StableVideoDiffusionPipelineOutput(BaseOutput): Output class for Stable Video Diffusion pipeline. Args: - frames (`[List[List[PIL.Image.Image]]`, `np.ndarray`, `torch.Tensor`]): - List of denoised PIL images of length `batch_size` or numpy array or torch tensor of shape `(batch_size, + frames (`[list[list[PIL.Image.Image]]`, `np.ndarray`, `torch.Tensor`]): + list of denoised PIL images of length `batch_size` or numpy array or torch tensor of shape `(batch_size, num_frames, height, width, num_channels)`. """ - frames: Union[List[List[PIL.Image.Image]], np.ndarray, torch.Tensor] + frames: list[list[PIL.Image.Image]] | np.ndarray | torch.Tensor class StableVideoDiffusionPipeline(DiffusionPipeline): @@ -191,7 +191,7 @@ def __init__( def _encode_image( self, image: PipelineImageInput, - device: Union[str, torch.device], + device: str | torch.device, num_videos_per_prompt: int, do_classifier_free_guidance: bool, ) -> torch.Tensor: @@ -239,7 +239,7 @@ def _encode_image( def _encode_vae_image( self, image: torch.Tensor, - device: Union[str, torch.device], + device: str | torch.device, num_videos_per_prompt: int, do_classifier_free_guidance: bool, ): @@ -323,7 +323,7 @@ def check_inputs(self, image, height, width): and not isinstance(image, list) ): raise ValueError( - "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is" + "`image` has to be of type `torch.Tensor` or `PIL.Image.Image` or `list[PIL.Image.Image]` but is" f" {type(image)}" ) @@ -338,9 +338,9 @@ def prepare_latents( height: int, width: int, dtype: torch.dtype, - device: Union[str, torch.device], + device: str | torch.device, generator: torch.Generator, - latents: Optional[torch.Tensor] = None, + latents: torch.Tensor | None = None, ): shape = ( batch_size, @@ -385,31 +385,31 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - image: Union[PIL.Image.Image, List[PIL.Image.Image], torch.Tensor], + image: PIL.Image.Image | list[PIL.Image.Image] | torch.Tensor, height: int = 576, width: int = 1024, - num_frames: Optional[int] = None, + num_frames: int | None = None, num_inference_steps: int = 25, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, min_guidance_scale: float = 1.0, max_guidance_scale: float = 3.0, fps: int = 7, motion_bucket_id: int = 127, noise_aug_strength: float = 0.02, - decode_chunk_size: Optional[int] = None, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + decode_chunk_size: int | None = None, + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], return_dict: bool = True, ): r""" The call function to the pipeline for generation. Args: - image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.Tensor`): + image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.Tensor`): Image(s) to guide image generation. If you provide a tensor, the expected value range is between `[0, 1]`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): @@ -422,7 +422,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 25): The number of denoising steps. More denoising steps usually lead to a higher quality video at the expense of slower inference. This parameter is modulated by `strength`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -445,7 +445,7 @@ def __call__( For lower memory usage, reduce `decode_chunk_size`. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of videos to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -460,7 +460,7 @@ def __call__( `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -473,7 +473,7 @@ def __call__( Returns: [`~pipelines.stable_diffusion.StableVideoDiffusionPipelineOutput`] or `tuple`: If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableVideoDiffusionPipelineOutput`] is - returned, otherwise a `tuple` of (`List[List[PIL.Image.Image]]` or `np.ndarray` or `torch.Tensor`) is + returned, otherwise a `tuple` of (`list[list[PIL.Image.Image]]` or `np.ndarray` or `torch.Tensor`) is returned. """ # 0. Default height and width to unet @@ -544,7 +544,13 @@ def __call__( added_time_ids = added_time_ids.to(device) # 6. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, None, sigmas) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, None, sigmas + ) # 7. Prepare latent variables num_channels_latents = self.unet.config.in_channels @@ -691,7 +697,7 @@ def _filter2d(input, kernel): height, width = tmp_kernel.shape[-2:] - padding_shape: List[int] = _compute_padding([height, width]) + padding_shape: list[int] = _compute_padding([height, width]) input = torch.nn.functional.pad(input, padding_shape, mode="reflect") # kernel and input tensor reshape to align element-wise or batch-wise params diff --git a/src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py b/src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py index 1ce6987114a7..ffb877cfd0f6 100644 --- a/src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py +++ b/src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py @@ -14,7 +14,7 @@ import inspect from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -54,16 +54,16 @@ class StableDiffusionAdapterPipelineOutput(BaseOutput): """ Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. - nsfw_content_detected (`List[bool]`) - List of flags denoting whether the corresponding generated image likely represents "not-safe-for-work" + nsfw_content_detected (`list[bool]`) + list of flags denoting whether the corresponding generated image likely represents "not-safe-for-work" (nsfw) content, or `None` if safety checking could not be performed. """ - images: Union[List[PIL.Image.Image], np.ndarray] - nsfw_content_detected: Optional[List[bool]] + images: list[PIL.Image.Image] | np.ndarray + nsfw_content_detected: list[bool] | None logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -131,10 +131,10 @@ def _preprocess_adapter_image(image, height, width): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -149,15 +149,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -197,11 +197,11 @@ class StableDiffusionAdapterPipeline(DiffusionPipeline, StableDiffusionMixin, Fr library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.) Args: - adapter ([`T2IAdapter`] or [`MultiAdapter`] or `List[T2IAdapter]`): + adapter ([`T2IAdapter`] or [`MultiAdapter`] or `list[T2IAdapter]`): Provides additional conditioning to the unet during the denoising process. If you set multiple Adapter as a list, the outputs from each Adapter are added together to create one combined additional conditioning. - adapter_weights (`List[float]`, *optional*, defaults to None): - List of floats representing the weight which will be multiply to each adapter's output before adding them + adapter_weights (`list[float]`, *optional*, defaults to None): + list of floats representing the weight which will be multiply to each adapter's output before adding them together. vae ([`AutoencoderKL`]): Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. @@ -233,7 +233,7 @@ def __init__( text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, unet: UNet2DConditionModel, - adapter: Union[T2IAdapter, MultiAdapter, List[T2IAdapter]], + adapter: T2IAdapter | MultiAdapter | list[T2IAdapter], scheduler: KarrasDiffusionSchedulers, safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, @@ -282,9 +282,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -315,16 +315,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -332,7 +332,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -431,7 +431,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -690,37 +690,37 @@ def do_classifier_free_guidance(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - image: Union[torch.Tensor, PIL.Image.Image, List[PIL.Image.Image]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + image: torch.Tensor | PIL.Image.Image | list[PIL.Image.Image] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - adapter_conditioning_scale: Union[float, List[float]] = 1.0, - clip_skip: Optional[int] = None, + cross_attention_kwargs: dict[str, Any] | None = None, + adapter_conditioning_scale: float | list[float] = 1.0, + clip_skip: int | None = None, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - image (`torch.Tensor`, `PIL.Image.Image`, `List[torch.Tensor]` or `List[PIL.Image.Image]` or `List[List[PIL.Image.Image]]`): + image (`torch.Tensor`, `PIL.Image.Image`, `list[torch.Tensor]` or `list[PIL.Image.Image]` or `list[list[PIL.Image.Image]]`): The Adapter input condition. Adapter uses this input condition to generate guidance to Unet. If the type is specified as `torch.Tensor`, it is passed to Adapter as is. PIL.Image.Image` can also be accepted as an image. The control image is automatically resized to fit the output image. @@ -731,11 +731,11 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -745,7 +745,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -754,7 +754,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -784,7 +784,7 @@ def __call__( A kwargs dictionary that if specified is passed along to the `AttnProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - adapter_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + adapter_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the adapter are multiplied by `adapter_conditioning_scale` before they are added to the residual in the original unet. If multiple adapters are specified in init, you can set the corresponding scale as a list. @@ -848,8 +848,12 @@ def __call__( prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) # 4. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) # 5. Prepare latent variables diff --git a/src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py b/src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py index 0ea3ba5046cf..a6dd07847de2 100644 --- a/src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py +++ b/src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL.Image @@ -153,10 +153,10 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -171,15 +171,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -233,11 +233,11 @@ class StableDiffusionXLAdapterPipeline( - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters Args: - adapter ([`T2IAdapter`] or [`MultiAdapter`] or `List[T2IAdapter]`): + adapter ([`T2IAdapter`] or [`MultiAdapter`] or `list[T2IAdapter]`): Provides additional conditioning to the unet during the denoising process. If you set multiple Adapter as a list, the outputs from each Adapter are added together to create one combined additional conditioning. - adapter_weights (`List[float]`, *optional*, defaults to None): - List of floats representing the weight which will be multiply to each adapter's output before adding them + adapter_weights (`list[float]`, *optional*, defaults to None): + list of floats representing the weight which will be multiply to each adapter's output before adding them together. vae ([`AutoencoderKL`]): Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. @@ -278,7 +278,7 @@ def __init__( tokenizer: CLIPTokenizer, tokenizer_2: CLIPTokenizer, unet: UNet2DConditionModel, - adapter: Union[T2IAdapter, MultiAdapter, List[T2IAdapter]], + adapter: T2IAdapter | MultiAdapter | list[T2IAdapter], scheduler: KarrasDiffusionSchedulers, force_zeros_for_empty_prompt: bool = True, feature_extractor: CLIPImageProcessor = None, @@ -311,26 +311,26 @@ def __init__( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, + prompt_2: str | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders device: (`torch.device`): @@ -339,11 +339,11 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -461,7 +461,7 @@ def encode_prompt( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) - uncond_tokens: List[str] + uncond_tokens: list[str] if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -857,55 +857,55 @@ def do_classifier_free_guidance(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] = None, + prompt_2: str | list[str] | None = None, image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - timesteps: List[int] = None, - sigmas: List[float] = None, - denoising_end: Optional[float] = None, + timesteps: list[int] = None, + sigmas: list[float] = None, + denoising_end: float | None = None, guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + ip_adapter_image: PipelineImageInput | None = None, + ip_adapter_image_embeds: list[torch.Tensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, guidance_rescale: float = 0.0, - original_size: Optional[Tuple[int, int]] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Optional[Tuple[int, int]] = None, - negative_original_size: Optional[Tuple[int, int]] = None, - negative_crops_coords_top_left: Tuple[int, int] = (0, 0), - negative_target_size: Optional[Tuple[int, int]] = None, - adapter_conditioning_scale: Union[float, List[float]] = 1.0, + original_size: tuple[int, int] | None = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] | None = None, + negative_original_size: tuple[int, int] | None = None, + negative_crops_coords_top_left: tuple[int, int] = (0, 0), + negative_target_size: tuple[int, int] | None = None, + adapter_conditioning_scale: float | list[float] = 1.0, adapter_conditioning_factor: float = 1.0, - clip_skip: Optional[int] = None, + clip_skip: int | None = None, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders - image (`torch.Tensor`, `PIL.Image.Image`, `List[torch.Tensor]` or `List[PIL.Image.Image]` or `List[List[PIL.Image.Image]]`): + image (`torch.Tensor`, `PIL.Image.Image`, `list[torch.Tensor]` or `list[PIL.Image.Image]` or `list[list[PIL.Image.Image]]`): The Adapter input condition. Adapter uses this input condition to generate guidance to Unet. If the type is specified as `torch.Tensor`, it is passed to Adapter as is. PIL.Image.Image` can also be accepted as an image. The control image is automatically resized to fit the output image. @@ -920,11 +920,11 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -941,11 +941,11 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -953,7 +953,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -975,7 +975,7 @@ def __call__( weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` input argument. ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*): + ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*): Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not @@ -1002,37 +1002,37 @@ def __call__( [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when using zero terminal SNR. - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a specific image resolution. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + negative_crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + negative_target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): To negatively condition the generation process based on a target image resolution. It should be as same as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - adapter_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + adapter_conditioning_scale (`float` or `list[float]`, *optional*, defaults to 1.0): The outputs of the adapter are multiplied by `adapter_conditioning_scale` before they are added to the residual in the original unet. If multiple adapters are specified in init, you can set the corresponding scale as a list. @@ -1130,8 +1130,12 @@ def __call__( ) # 4. Prepare timesteps + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps, sigmas + self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas ) # 5. Prepare latent variables diff --git a/src/diffusers/pipelines/text_to_video_synthesis/pipeline_output.py b/src/diffusers/pipelines/text_to_video_synthesis/pipeline_output.py index 040bf0efba84..c94c5d2d144a 100644 --- a/src/diffusers/pipelines/text_to_video_synthesis/pipeline_output.py +++ b/src/diffusers/pipelines/text_to_video_synthesis/pipeline_output.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL @@ -16,11 +15,11 @@ class TextToVideoSDPipelineOutput(BaseOutput): Output class for text-to-video pipelines. Args: - frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): - List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): + list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape `(batch_size, num_frames, channels, height, width)` """ - frames: Union[torch.Tensor, np.ndarray, List[List[PIL.Image.Image]]] + frames: torch.Tensor | np.ndarray | list[list[PIL.Image.Image]] diff --git a/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py b/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py index 3ce7b4d1990f..0ca64d33acda 100644 --- a/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py +++ b/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch from transformers import CLIPTextModel, CLIPTokenizer @@ -130,9 +130,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -163,16 +163,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -180,7 +180,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -279,7 +279,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -368,7 +368,6 @@ def prepare_extra_step_kwargs(self, generator, eta): extra_step_kwargs["generator"] = generator return extra_step_kwargs - # Copied from diffusers.pipelines.stable_diffusion_k_diffusion.pipeline_stable_diffusion_k_diffusion.StableDiffusionKDiffusionPipeline.check_inputs def check_inputs( self, prompt, @@ -450,30 +449,30 @@ def prepare_latents( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_frames: int = 16, num_inference_steps: int = 50, guidance_scale: float = 9.0, - negative_prompt: Optional[Union[str, List[str]]] = None, + negative_prompt: str | list[str] | None = None, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "np", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "np", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, + cross_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated video. @@ -488,7 +487,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -496,7 +495,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py b/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py index 9d0b7e3dbc32..6908f51eb21b 100644 --- a/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py +++ b/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -90,7 +90,7 @@ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -165,9 +165,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -198,16 +198,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -215,7 +215,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -314,7 +314,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -515,31 +515,31 @@ def prepare_latents(self, video, timestep, batch_size, dtype, device, generator= @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - video: Union[List[np.ndarray], torch.Tensor] = None, + prompt: str | list[str] = None, + video: list[np.ndarray] | torch.Tensor = None, strength: float = 0.6, num_inference_steps: int = 50, guidance_scale: float = 15.0, - negative_prompt: Optional[Union[str, List[str]]] = None, + negative_prompt: str | list[str] | None = None, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "np", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "np", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, + cross_attention_kwargs: dict[str, Any] | None = None, + clip_skip: int | None = None, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. - video (`List[np.ndarray]` or `torch.Tensor`): + video (`list[np.ndarray]` or `torch.Tensor`): `video` frames or tensor representing a video batch to be used as the starting point for the process. Can also accept video latents as `image`, if passing latents directly, it will not be encoded again. strength (`float`, *optional*, defaults to 0.8): @@ -554,13 +554,13 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in video generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py b/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py index 96316f8e91e5..66defb2f3745 100644 --- a/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py +++ b/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py @@ -1,7 +1,7 @@ import copy import inspect from dataclasses import dataclass -from typing import Callable, List, Optional, Union +from typing import Callable import numpy as np import PIL.Image @@ -199,16 +199,16 @@ class TextToVideoPipelineOutput(BaseOutput): Output class for zero-shot text-to-video pipeline. Args: - images (`[List[PIL.Image.Image]`, `np.ndarray`]): - List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, + images (`[list[PIL.Image.Image]`, `np.ndarray`]): + list of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, num_channels)`. - nsfw_content_detected (`[List[bool]]`): - List indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or + nsfw_content_detected (`[list[bool]]`): + list indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or `None` if safety checking could not be performed. """ - images: Union[List[PIL.Image.Image], np.ndarray] - nsfw_content_detected: Optional[List[bool]] + images: list[PIL.Image.Image] | np.ndarray + nsfw_content_detected: list[bool] | None def coords_grid(batch, ht, wd, device): @@ -374,7 +374,7 @@ def forward_loop(self, x_t0, t0, t1, generator): Timestep at t0. t1: Timestamp at t1. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. @@ -466,7 +466,6 @@ def backward_loop( return latents.clone().detach() - # Copied from diffusers.pipelines.stable_diffusion_k_diffusion.pipeline_stable_diffusion_k_diffusion.StableDiffusionKDiffusionPipeline.check_inputs def check_inputs( self, prompt, @@ -545,32 +544,32 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype @torch.no_grad() def __call__( self, - prompt: Union[str, List[str]], - video_length: Optional[int] = 8, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str], + video_length: int | None = 8, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_videos_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_videos_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, motion_field_strength_x: float = 12, motion_field_strength_y: float = 12, - output_type: Optional[str] = "tensor", + output_type: str | None = "tensor", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, - callback_steps: Optional[int] = 1, + callback: Callable[[int, int, torch.Tensor], None] | None = None, + callback_steps: int | None = 1, t0: int = 44, t1: int = 47, - frame_ids: Optional[List[int]] = None, + frame_ids: list[int] | None = None, ): """ The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. video_length (`int`, *optional*, defaults to 8): The number of generated video frames. @@ -584,7 +583,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 7.5): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in video generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_videos_per_prompt (`int`, *optional*, defaults to 1): @@ -592,7 +591,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -623,7 +622,7 @@ def __call__( t1 (`int`, *optional*, defaults to 47): Timestep t0. Should be in the range [t0 + 1, num_inference_steps - 1]. See the [paper](https://huggingface.co/papers/2303.13439), Sect. 3.3.1. - frame_ids (`List[int]`, *optional*): + frame_ids (`list[int]`, *optional*): Indexes of the frames that are being generated. This is used when generating longer videos chunk-by-chunk. @@ -821,16 +820,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -838,7 +837,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -937,7 +936,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): diff --git a/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py b/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py index c8dce75e2671..a3286cd940fd 100644 --- a/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py +++ b/src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py @@ -1,7 +1,7 @@ import copy import inspect from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import numpy as np import PIL @@ -218,12 +218,12 @@ class TextToVideoSDXLPipelineOutput(BaseOutput): Output class for zero-shot text-to-video pipeline. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image] | np.ndarray # Copied from diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.coords_grid @@ -402,7 +402,7 @@ def __init__( image_encoder: CLIPVisionModelWithProjection = None, feature_extractor: CLIPImageProcessor = None, force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, + add_watermarker: bool | None = None, ): super().__init__() self.register_modules( @@ -585,26 +585,26 @@ def check_inputs( def encode_prompt( self, prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, + prompt_2: str | None = None, + device: torch.device | None = None, num_images_per_prompt: int = 1, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + negative_prompt: str | None = None, + negative_prompt_2: str | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders device: (`torch.device`): @@ -613,11 +613,11 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders prompt_embeds (`torch.Tensor`, *optional*): @@ -735,7 +735,7 @@ def encode_prompt( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) - uncond_tokens: List[str] + uncond_tokens: list[str] if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" @@ -832,7 +832,7 @@ def forward_loop(self, x_t0, t0, t1, generator): Timestep at t0. t1: Timestamp at t1. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. @@ -937,36 +937,36 @@ def backward_loop( @torch.no_grad() def __call__( self, - prompt: Union[str, List[str]], - prompt_2: Optional[Union[str, List[str]]] = None, - video_length: Optional[int] = 8, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str], + prompt_2: str | list[str] | None = None, + video_length: int | None = 8, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - denoising_end: Optional[float] = None, + denoising_end: float | None = None, guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_videos_per_prompt: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + negative_prompt_2: str | list[str] | None = None, + num_videos_per_prompt: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - frame_ids: Optional[List[int]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - latents: Optional[torch.Tensor] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + frame_ids: list[int] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + latents: torch.Tensor | None = None, motion_field_strength_x: float = 12, motion_field_strength_y: float = 12, - output_type: Optional[str] = "tensor", + output_type: str | None = "tensor", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: dict[str, Any] | None = None, guidance_rescale: float = 0.0, - original_size: Optional[Tuple[int, int]] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Optional[Tuple[int, int]] = None, + original_size: tuple[int, int] | None = None, + crops_coords_top_left: tuple[int, int] = (0, 0), + target_size: tuple[int, int] | None = None, t0: int = 44, t1: int = 47, ): @@ -974,10 +974,10 @@ def __call__( Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - prompt_2 (`str` or `List[str]`, *optional*): + prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders video_length (`int`, *optional*, defaults to 8): @@ -1002,11 +1002,11 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): + negative_prompt_2 (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders num_videos_per_prompt (`int`, *optional*, defaults to 1): @@ -1014,10 +1014,10 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. - frame_ids (`List[int]`, *optional*): + frame_ids (`list[int]`, *optional*): Indexes of the frames that are being generated. This is used when generating longer videos chunk-by-chunk. prompt_embeds (`torch.Tensor`, *optional*): @@ -1066,17 +1066,17 @@ def __call__( [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when using zero terminal SNR. - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + original_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. `original_size` defaults to `(width, height)` if not specified. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + crops_coords_top_left (`tuple[int]`, *optional*, defaults to (0, 0)): `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + target_size (`tuple[int]`, *optional*, defaults to (1024, 1024)): For most cases, `target_size` should be set to the desired height and width of the generated image. If not specified it will default to `(width, height)`. Part of SDXL's micro-conditioning as explained in section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). diff --git a/src/diffusers/pipelines/transformers_loading_utils.py b/src/diffusers/pipelines/transformers_loading_utils.py index b52d154d6ba2..02f8ee612601 100644 --- a/src/diffusers/pipelines/transformers_loading_utils.py +++ b/src/diffusers/pipelines/transformers_loading_utils.py @@ -15,7 +15,7 @@ import contextlib import os import tempfile -from typing import TYPE_CHECKING, Dict +from typing import TYPE_CHECKING from huggingface_hub import DDUFEntry from tqdm import tqdm @@ -34,7 +34,7 @@ def _load_tokenizer_from_dduf( - cls: "PreTrainedTokenizer", name: str, dduf_entries: Dict[str, DDUFEntry], **kwargs + cls: "PreTrainedTokenizer", name: str, dduf_entries: dict[str, DDUFEntry], **kwargs ) -> "PreTrainedTokenizer": """ Load a tokenizer from a DDUF archive. @@ -57,7 +57,7 @@ def _load_tokenizer_from_dduf( def _load_transformers_model_from_dduf( - cls: "PreTrainedModel", name: str, dduf_entries: Dict[str, DDUFEntry], **kwargs + cls: "PreTrainedModel", name: str, dduf_entries: dict[str, DDUFEntry], **kwargs ) -> "PreTrainedModel": """ Load a transformers model from a DDUF archive. @@ -112,6 +112,8 @@ def _load_transformers_model_from_dduf( tensors = safetensors.torch.load(mmap) # Update the state dictionary with tensors state_dict.update(tensors) + # `from_pretrained` sets the model to eval mode by default, which is the + # correct behavior for inference. Do not call `model.train()` here. return cls.from_pretrained( pretrained_model_name_or_path=None, config=config, diff --git a/src/diffusers/pipelines/unclip/pipeline_unclip.py b/src/diffusers/pipelines/unclip/pipeline_unclip.py index bbb9b0eb3ab2..430f1a1e5265 100644 --- a/src/diffusers/pipelines/unclip/pipeline_unclip.py +++ b/src/diffusers/pipelines/unclip/pipeline_unclip.py @@ -13,7 +13,6 @@ # limitations under the License. import inspect -from typing import List, Optional, Tuple, Union import torch from torch.nn import functional as F @@ -131,8 +130,8 @@ def _encode_prompt( device, num_images_per_prompt, do_classifier_free_guidance, - text_model_output: Optional[Union[CLIPTextModelOutput, Tuple]] = None, - text_attention_mask: Optional[torch.Tensor] = None, + text_model_output: CLIPTextModelOutput | tuple | None = None, + text_attention_mask: torch.Tensor | None = None, ): if text_model_output is None: batch_size = len(prompt) if isinstance(prompt, list) else 1 @@ -219,27 +218,27 @@ def _encode_prompt( @torch.no_grad() def __call__( self, - prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] | None = None, num_images_per_prompt: int = 1, prior_num_inference_steps: int = 25, decoder_num_inference_steps: int = 25, super_res_num_inference_steps: int = 7, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - prior_latents: Optional[torch.Tensor] = None, - decoder_latents: Optional[torch.Tensor] = None, - super_res_latents: Optional[torch.Tensor] = None, - text_model_output: Optional[Union[CLIPTextModelOutput, Tuple]] = None, - text_attention_mask: Optional[torch.Tensor] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + prior_latents: torch.Tensor | None = None, + decoder_latents: torch.Tensor | None = None, + super_res_latents: torch.Tensor | None = None, + text_model_output: CLIPTextModelOutput | tuple | None = None, + text_attention_mask: torch.Tensor | None = None, prior_guidance_scale: float = 4.0, decoder_guidance_scale: float = 8.0, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, ): """ The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide image generation. This can only be left undefined if `text_model_output` and `text_attention_mask` is passed. num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -253,7 +252,7 @@ def __call__( super_res_num_inference_steps (`int`, *optional*, defaults to 7): The number of denoising steps for super resolution. More denoising steps usually lead to a higher quality image at the expense of slower inference. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. prior_latents (`torch.Tensor` of shape (batch size, embeddings dimension), *optional*): diff --git a/src/diffusers/pipelines/unclip/pipeline_unclip_image_variation.py b/src/diffusers/pipelines/unclip/pipeline_unclip_image_variation.py index 31710a000e0a..d0d8bdc44787 100644 --- a/src/diffusers/pipelines/unclip/pipeline_unclip_image_variation.py +++ b/src/diffusers/pipelines/unclip/pipeline_unclip_image_variation.py @@ -13,7 +13,6 @@ # limitations under the License. import inspect -from typing import List, Optional, Union import PIL.Image import torch @@ -190,7 +189,7 @@ def _encode_prompt(self, prompt, device, num_images_per_prompt, do_classifier_fr return prompt_embeds, text_encoder_hidden_states, text_mask - def _encode_image(self, image, device, num_images_per_prompt, image_embeddings: Optional[torch.Tensor] = None): + def _encode_image(self, image, device, num_images_per_prompt, image_embeddings: torch.Tensor | None = None): dtype = next(self.image_encoder.parameters()).dtype if image_embeddings is None: @@ -207,23 +206,23 @@ def _encode_image(self, image, device, num_images_per_prompt, image_embeddings: @torch.no_grad() def __call__( self, - image: Optional[Union[PIL.Image.Image, List[PIL.Image.Image], torch.Tensor]] = None, + image: PIL.Image.Image | list[PIL.Image.Image] | torch.Tensor | None = None, num_images_per_prompt: int = 1, decoder_num_inference_steps: int = 25, super_res_num_inference_steps: int = 7, - generator: Optional[torch.Generator] = None, - decoder_latents: Optional[torch.Tensor] = None, - super_res_latents: Optional[torch.Tensor] = None, - image_embeddings: Optional[torch.Tensor] = None, + generator: torch.Generator | None = None, + decoder_latents: torch.Tensor | None = None, + super_res_latents: torch.Tensor | None = None, + image_embeddings: torch.Tensor | None = None, decoder_guidance_scale: float = 8.0, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, ): """ The call function to the pipeline for generation. Args: - image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.Tensor`): + image (`PIL.Image.Image` or `list[PIL.Image.Image]` or `torch.Tensor`): `Image` or tensor representing an image batch to be used as the starting point. If you provide a tensor, it needs to be compatible with the [`CLIPImageProcessor`] [configuration](https://huggingface.co/fusing/karlo-image-variations-diffusers/blob/main/feature_extractor/preprocessor_config.json). diff --git a/src/diffusers/pipelines/unidiffuser/modeling_text_decoder.py b/src/diffusers/pipelines/unidiffuser/modeling_text_decoder.py index 0ddcbf735770..c68b5d9ab5a8 100644 --- a/src/diffusers/pipelines/unidiffuser/modeling_text_decoder.py +++ b/src/diffusers/pipelines/unidiffuser/modeling_text_decoder.py @@ -1,5 +1,3 @@ -from typing import Optional - import numpy as np import torch from torch import nn @@ -68,13 +66,13 @@ def __init__( self, prefix_length: int, prefix_inner_dim: int, - prefix_hidden_dim: Optional[int] = None, + prefix_hidden_dim: int | None = None, vocab_size: int = 50257, # Start of GPT2 config args n_positions: int = 1024, n_embd: int = 768, n_layer: int = 12, n_head: int = 12, - n_inner: Optional[int] = None, + n_inner: int | None = None, activation_function: str = "gelu_new", resid_pdrop: float = 0.1, embd_pdrop: float = 0.1, @@ -132,8 +130,8 @@ def forward( self, input_ids: torch.Tensor, prefix_embeds: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, + attention_mask: torch.Tensor | None = None, + labels: torch.Tensor | None = None, ): """ Args: @@ -180,7 +178,7 @@ def generate_captions(self, features, eos_token_id, device): Device to perform text generation on. Returns: - `List[str]`: A list of strings generated from the decoder model. + `list[str]`: A list of strings generated from the decoder model. """ features = torch.split(features, 1, dim=0) @@ -207,7 +205,7 @@ def generate_beam( beam_size: int = 5, entry_length: int = 67, temperature: float = 1.0, - eos_token_id: Optional[int] = None, + eos_token_id: int | None = None, ): """ Generates text using the given tokenizer and text prompt or token embedding via beam search. This @@ -233,7 +231,7 @@ def generate_beam( The temperature to use when performing the softmax over logits from the decoding model. Returns: - `Tuple(torch.Tensor, torch.Tensor)`: A tuple of tensors where the first element is a tensor of generated + `tuple(torch.Tensor, torch.Tensor)`: A tuple of tensors where the first element is a tensor of generated token sequences sorted by score in descending order, and the second element is the sequence lengths corresponding to those sequences. """ diff --git a/src/diffusers/pipelines/unidiffuser/modeling_uvit.py b/src/diffusers/pipelines/unidiffuser/modeling_uvit.py index 2a04ec2e4030..125188196c1e 100644 --- a/src/diffusers/pipelines/unidiffuser/modeling_uvit.py +++ b/src/diffusers/pipelines/unidiffuser/modeling_uvit.py @@ -1,5 +1,4 @@ import math -from typing import Optional, Union import torch from torch import nn @@ -178,9 +177,9 @@ def __init__( num_attention_heads: int, attention_head_dim: int, dropout=0.0, - cross_attention_dim: Optional[int] = None, + cross_attention_dim: int | None = None, activation_fn: str = "geglu", - num_embeds_ada_norm: Optional[int] = None, + num_embeds_ada_norm: int | None = None, attention_bias: bool = False, only_cross_attention: bool = False, double_self_attention: bool = False, @@ -373,9 +372,9 @@ def __init__( num_attention_heads: int, attention_head_dim: int, dropout=0.0, - cross_attention_dim: Optional[int] = None, + cross_attention_dim: int | None = None, activation_fn: str = "geglu", - num_embeds_ada_norm: Optional[int] = None, + num_embeds_ada_norm: int | None = None, attention_bias: bool = False, only_cross_attention: bool = False, double_self_attention: bool = False, @@ -591,18 +590,18 @@ def __init__( self, num_attention_heads: int = 16, attention_head_dim: int = 88, - in_channels: Optional[int] = None, - out_channels: Optional[int] = None, + in_channels: int | None = None, + out_channels: int | None = None, num_layers: int = 1, dropout: float = 0.0, norm_num_groups: int = 32, - cross_attention_dim: Optional[int] = None, + cross_attention_dim: int | None = None, attention_bias: bool = False, - sample_size: Optional[int] = None, - num_vector_embeds: Optional[int] = None, - patch_size: Optional[int] = 2, + sample_size: int | None = None, + num_vector_embeds: int | None = None, + patch_size: int | None = 2, activation_fn: str = "geglu", - num_embeds_ada_norm: Optional[int] = None, + num_embeds_ada_norm: int | None = None, use_linear_projection: bool = False, only_cross_attention: bool = False, upcast_attention: bool = False, @@ -902,18 +901,18 @@ def __init__( num_text_tokens: int = 77, num_attention_heads: int = 16, attention_head_dim: int = 88, - in_channels: Optional[int] = None, - out_channels: Optional[int] = None, + in_channels: int | None = None, + out_channels: int | None = None, num_layers: int = 1, dropout: float = 0.0, norm_num_groups: int = 32, - cross_attention_dim: Optional[int] = None, + cross_attention_dim: int | None = None, attention_bias: bool = False, - sample_size: Optional[int] = None, - num_vector_embeds: Optional[int] = None, - patch_size: Optional[int] = None, + sample_size: int | None = None, + num_vector_embeds: int | None = None, + patch_size: int | None = None, activation_fn: str = "geglu", - num_embeds_ada_norm: Optional[int] = None, + num_embeds_ada_norm: int | None = None, use_linear_projection: bool = False, only_cross_attention: bool = False, upcast_attention: bool = False, @@ -1040,9 +1039,9 @@ def forward( latent_image_embeds: torch.Tensor, image_embeds: torch.Tensor, prompt_embeds: torch.Tensor, - timestep_img: Union[torch.Tensor, float, int], - timestep_text: Union[torch.Tensor, float, int], - data_type: Optional[Union[torch.Tensor, float, int]] = 1, + timestep_img: torch.Tensor | float | int, + timestep_text: torch.Tensor | float | int, + data_type: torch.Tensor | float | int | None = 1, encoder_hidden_states=None, cross_attention_kwargs=None, ): diff --git a/src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py b/src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py index f9298d5b86f8..81d2ce95dc53 100644 --- a/src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py +++ b/src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py @@ -1,6 +1,6 @@ import inspect from dataclasses import dataclass -from typing import Callable, List, Optional, Union +from typing import Callable import numpy as np import PIL.Image @@ -50,16 +50,16 @@ class ImageTextPipelineOutput(BaseOutput): Output class for joint image-text pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) - List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, + images (`list[PIL.Image.Image]` or `np.ndarray`) + list of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, num_channels)`. - text (`List[str]` or `List[List[str]]`) - List of generated text strings of length `batch_size` or a list of list of strings whose outer list has + text (`list[str]` or `list[list[str]]`) + list of generated text strings of length `batch_size` or a list of list of strings whose outer list has length `batch_size`. """ - images: Optional[Union[List[PIL.Image.Image], np.ndarray]] - text: Optional[Union[List[str], List[List[str]]]] + images: list[PIL.Image.Image] | np.ndarray | None + text: list[str] | list[list[str]] | None class UniDiffuserPipeline(DeprecatedPipelineMixin, DiffusionPipeline): @@ -392,9 +392,9 @@ def _encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, **kwargs, ): deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." @@ -425,16 +425,16 @@ def encode_prompt( num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + lora_scale: float | None = None, + clip_skip: int | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device @@ -442,7 +442,7 @@ def encode_prompt( number of images that should be generated per prompt do_classifier_free_guidance (`bool`): whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -541,7 +541,7 @@ def encode_prompt( # get unconditional embeddings for classifier free guidance if do_classifier_free_guidance and negative_prompt_embeds is None: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif prompt is not None and type(prompt) is not type(negative_prompt): @@ -1119,34 +1119,34 @@ def check_inputs( @torch.no_grad() def __call__( self, - prompt: Optional[Union[str, List[str]]] = None, - image: Optional[Union[torch.Tensor, PIL.Image.Image]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - data_type: Optional[int] = 1, + prompt: str | list[str] | None = None, + image: torch.Tensor | PIL.Image.Image | None = None, + height: int | None = None, + width: int | None = None, + data_type: int | None = 1, num_inference_steps: int = 50, guidance_scale: float = 8.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - num_prompts_per_image: Optional[int] = 1, + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + num_prompts_per_image: int | None = 1, eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_latents: Optional[torch.Tensor] = None, - vae_latents: Optional[torch.Tensor] = None, - clip_latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_latents: torch.Tensor | None = None, + vae_latents: torch.Tensor | None = None, + clip_latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback: Callable[[int, int, torch.Tensor], None] | None = None, callback_steps: int = 1, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. Required for text-conditioned image generation (`text2img`) mode. image (`torch.Tensor` or `PIL.Image.Image`, *optional*): @@ -1166,7 +1166,7 @@ def __call__( guidance_scale (`float`, *optional*, defaults to 8.0): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). Used in text-conditioned image generation (`text2img`) mode. @@ -1181,7 +1181,7 @@ def __call__( eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_combined.py b/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_combined.py index 91a54e1ae82f..f640fddc2bc5 100644 --- a/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_combined.py +++ b/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_combined.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch from PIL import Image @@ -250,24 +250,24 @@ def check_inputs( @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - task_prompt: Union[str, List[str]] = None, - content_prompt: Union[str, List[str]] = None, - image: Optional[torch.FloatTensor] = None, - upsampling_height: Optional[int] = None, - upsampling_width: Optional[int] = None, + task_prompt: str | list[str] = None, + content_prompt: str | list[str] = None, + image: torch.FloatTensor | None = None, + upsampling_height: int | None = None, + upsampling_width: int | None = None, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 30.0, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, upsampling_strength: float = 1.0, ): @@ -275,11 +275,11 @@ def __call__( Function invoked when calling the VisualCloze pipeline for generation. Args: - task_prompt (`str` or `List[str]`, *optional*): + task_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to define the task intention. - content_prompt (`str` or `List[str]`, *optional*): + content_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to define the content or caption of the target image to be generated. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a @@ -297,7 +297,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -309,7 +309,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -336,7 +336,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py b/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py index e12995106bcf..dd5d0603d6d0 100644 --- a/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py +++ b/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import numpy as np import torch @@ -192,11 +192,11 @@ def __init__( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_images_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -242,9 +242,9 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds def _get_clip_prompt_embeds( self, - prompt: Union[str, List[str]], + prompt: str | list[str], num_images_per_prompt: int = 1, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): device = device or self._execution_device @@ -287,25 +287,25 @@ def _get_clip_prompt_embeds( # Modified from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt def encode_prompt( self, - layout_prompt: Union[str, List[str]], - task_prompt: Union[str, List[str]], - content_prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + layout_prompt: str | list[str], + task_prompt: str | list[str], + content_prompt: str | list[str], + device: torch.device | None = None, num_images_per_prompt: int = 1, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 512, - lora_scale: Optional[float] = None, + lora_scale: float | None = None, ): r""" Args: - layout_prompt (`str` or `List[str]`, *optional*): + layout_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to define the number of in-context examples and the number of images involved in the task. - task_prompt (`str` or `List[str]`, *optional*): + task_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to define the task intention. - content_prompt (`str` or `List[str]`, *optional*): + content_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to define the content or caption of the target image to be generated. device: (`torch.device`): torch device @@ -709,33 +709,33 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - task_prompt: Union[str, List[str]] = None, - content_prompt: Union[str, List[str]] = None, - image: Optional[torch.FloatTensor] = None, + task_prompt: str | list[str] = None, + content_prompt: str | list[str] = None, + image: torch.FloatTensor | None = None, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 30.0, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: torch.FloatTensor | None = None, + pooled_prompt_embeds: torch.FloatTensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the VisualCloze pipeline for generation. Args: - task_prompt (`str` or `List[str]`, *optional*): + task_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to define the task intention. - content_prompt (`str` or `List[str]`, *optional*): + content_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to define the content or caption of the target image to be generated. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a @@ -743,7 +743,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -755,7 +755,7 @@ def __call__( the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): @@ -782,7 +782,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/visualcloze/visualcloze_utils.py b/src/diffusers/pipelines/visualcloze/visualcloze_utils.py index efe5dff47623..1a96ff8d2d13 100644 --- a/src/diffusers/pipelines/visualcloze/visualcloze_utils.py +++ b/src/diffusers/pipelines/visualcloze/visualcloze_utils.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Dict, List, Optional, Tuple, Union import torch from PIL import Image @@ -40,8 +39,8 @@ def __init__(self, *args, resolution: int = 384, **kwargs): self.resolution = resolution def preprocess_image( - self, input_images: List[List[Optional[Image.Image]]], vae_scale_factor: int - ) -> Tuple[List[List[torch.Tensor]], List[List[List[int]]], List[int]]: + self, input_images: list[list[Image.Image | None]], vae_scale_factor: int + ) -> tuple[list[list[torch.Tensor]], list[list[list[int]]], list[int]]: """ Preprocesses input images for the VisualCloze pipeline. @@ -52,7 +51,7 @@ def preprocess_image( 4. Tracking image sizes and positions of target images Args: - input_images (List[List[Optional[Image.Image]]]): + input_images (list[list[Image.Image | None]]): A nested list of PIL Images where: - Outer list represents different samples, including in-context examples and the query - Inner list contains images for the task @@ -61,17 +60,17 @@ def preprocess_image( The scale factor used by the VAE for resizing images Returns: - Tuple containing: - - List[List[torch.Tensor]]: Preprocessed images in tensor format - - List[List[List[int]]]: Dimensions of each processed image [height, width] - - List[int]: Target positions indicating which images are to be generated + tuple containing: + - list[list[torch.Tensor]]: Preprocessed images in tensor format + - list[list[list[int]]]: Dimensions of each processed image [height, width] + - list[int]: Target positions indicating which images are to be generated """ n_samples, n_task_images = len(input_images), len(input_images[0]) divisible = 2 * vae_scale_factor - processed_images: List[List[Image.Image]] = [[] for _ in range(n_samples)] - resize_size: List[Optional[Tuple[int, int]]] = [None for _ in range(n_samples)] - target_position: List[int] = [] + processed_images: list[list[Image.Image]] = [[] for _ in range(n_samples)] + resize_size: list[tuple[int, int] | None] = [None for _ in range(n_samples)] + target_position: list[int] = [] # Process each sample for i in range(n_samples): @@ -125,19 +124,19 @@ def preprocess_image( return processed_images, image_sizes, target_position def preprocess_mask( - self, input_images: List[List[Image.Image]], target_position: List[int] - ) -> List[List[torch.Tensor]]: + self, input_images: list[list[Image.Image]], target_position: list[int] + ) -> list[list[torch.Tensor]]: """ Generate masks for the VisualCloze pipeline. Args: - input_images (List[List[Image.Image]]): + input_images (list[list[Image.Image]]): Processed images from preprocess_image - target_position (List[int]): + target_position (list[int]): Binary list marking the positions of target images (1 for target, 0 for condition) Returns: - List[List[torch.Tensor]]: + list[list[torch.Tensor]]: A nested list of mask tensors (1 for target positions, 0 for condition images) """ mask = [] @@ -155,10 +154,10 @@ def preprocess_mask( def preprocess_image_upsampling( self, - input_images: List[List[Image.Image]], + input_images: list[list[Image.Image]], height: int, width: int, - ) -> Tuple[List[List[Image.Image]], List[List[List[int]]]]: + ) -> tuple[list[list[Image.Image]], list[list[list[int]]]]: """Process images for the upsampling stage in the VisualCloze pipeline. Args: @@ -167,7 +166,7 @@ def preprocess_image_upsampling( width: Target width Returns: - Tuple of processed image and its size + tuple of processed image and its size """ image = self.resize(input_images[0][0], height, width) image = self.pil_to_numpy(image) # to np @@ -178,10 +177,10 @@ def preprocess_image_upsampling( image_sizes = [[[height, width]]] return input_images, image_sizes - def preprocess_mask_upsampling(self, input_images: List[List[Image.Image]]) -> List[List[torch.Tensor]]: + def preprocess_mask_upsampling(self, input_images: list[list[Image.Image]]) -> list[list[torch.Tensor]]: return [[torch.ones((1, 1, input_images[0][0].shape[2], input_images[0][0].shape[3]))]] - def get_layout_prompt(self, size: Tuple[int, int]) -> str: + def get_layout_prompt(self, size: tuple[int, int]) -> str: layout_instruction = ( f"A grid layout with {size[0]} rows and {size[1]} columns, displaying {size[0] * size[1]} images arranged side by side.", ) @@ -189,26 +188,26 @@ def get_layout_prompt(self, size: Tuple[int, int]) -> str: def preprocess( self, - task_prompt: Union[str, List[str]], - content_prompt: Union[str, List[str]], - input_images: Optional[List[List[List[Optional[str]]]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + task_prompt: str | list[str], + content_prompt: str | list[str], + input_images: list[list[list[str | None]]] | None = None, + height: int | None = None, + width: int | None = None, upsampling: bool = False, vae_scale_factor: int = 16, - ) -> Dict: + ) -> dict: """Process visual cloze inputs. Args: task_prompt: Task description(s) content_prompt: Content description(s) - input_images: List of images or None for the target images + input_images: list of images or None for the target images height: Optional target height for upsampling stage width: Optional target width for upsampling stage upsampling: Whether this is in the upsampling processing stage Returns: - Dictionary containing processed images, masks, prompts and metadata + dictionary containing processed images, masks, prompts and metadata """ if isinstance(task_prompt, str): task_prompt = [task_prompt] diff --git a/src/diffusers/pipelines/wan/image_processor.py b/src/diffusers/pipelines/wan/image_processor.py index b1594d08630f..fa18150fcc6e 100644 --- a/src/diffusers/pipelines/wan/image_processor.py +++ b/src/diffusers/pipelines/wan/image_processor.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Tuple, Union import numpy as np import PIL.Image @@ -36,7 +35,7 @@ class WanAnimateImageProcessor(VaeImageProcessor): this factor. vae_latent_channels (`int`, *optional*, defaults to `16`): VAE latent channels. - spatial_patch_size (`Tuple[int, int]`, *optional*, defaults to `(2, 2)`): + spatial_patch_size (`tuple[int, int]`, *optional*, defaults to `(2, 2)`): The spatial patch size used by the diffusion transformer. For Wan models, this is typically (2, 2). resample (`str`, *optional*, defaults to `lanczos`): Resampling filter to use when resizing the image. @@ -48,7 +47,7 @@ class WanAnimateImageProcessor(VaeImageProcessor): Whether to convert the images to RGB format. do_convert_grayscale (`bool`, *optional*, defaults to be `False`): Whether to convert the images to grayscale format. - fill_color (`str` or `float` or `Tuple[float, ...]`, *optional*, defaults to `None`): + fill_color (`str` or `float` or `tuple[float, ...]`, *optional*, defaults to `None`): An optional fill color when `resize_mode` is set to `"fill"`. This will fill the empty space with that color instead of filling with data from the image. Any valid `color` argument to `PIL.Image.new` is valid; if `None`, will default to filling with data from `image`. @@ -60,14 +59,14 @@ def __init__( do_resize: bool = True, vae_scale_factor: int = 8, vae_latent_channels: int = 16, - spatial_patch_size: Tuple[int, int] = (2, 2), + spatial_patch_size: tuple[int, int] = (2, 2), resample: str = "lanczos", reducing_gap: int = None, do_normalize: bool = True, do_binarize: bool = False, do_convert_rgb: bool = False, do_convert_grayscale: bool = False, - fill_color: Optional[Union[str, float, Tuple[float, ...]]] = 0, + fill_color: str | float | tuple[float, ...] | None = 0, ): super().__init__() if do_convert_rgb and do_convert_grayscale: @@ -134,25 +133,25 @@ def _resize_and_fill( def get_default_height_width( self, - image: Union[PIL.Image.Image, np.ndarray, torch.Tensor], - height: Optional[int] = None, - width: Optional[int] = None, - ) -> Tuple[int, int]: + image: PIL.Image.Image | np.ndarray | torch.Tensor, + height: int | None = None, + width: int | None = None, + ) -> tuple[int, int]: r""" Returns the height and width of the image, downscaled to the next integer multiple of `vae_scale_factor`. Args: - image (`Union[PIL.Image.Image, np.ndarray, torch.Tensor]`): + image (`PIL.Image.Image | np.ndarray | torch.Tensor`): The image input, which can be a PIL image, NumPy array, or PyTorch tensor. If it is a NumPy array, it should have shape `[batch, height, width]` or `[batch, height, width, channels]`. If it is a PyTorch tensor, it should have shape `[batch, channels, height, width]`. - height (`Optional[int]`, *optional*, defaults to `None`): + height (`int | None`, *optional*, defaults to `None`): The height of the preprocessed image. If `None`, the height of the `image` input will be used. - width (`Optional[int]`, *optional*, defaults to `None`): + width (`int | None`, *optional*, defaults to `None`): The width of the preprocessed image. If `None`, the width of the `image` input will be used. Returns: - `Tuple[int, int]`: + `tuple[int, int]`: A tuple containing the height and width, both resized to the nearest integer multiple of `vae_scale_factor * spatial_patch_size`. """ diff --git a/src/diffusers/pipelines/wan/pipeline_output.py b/src/diffusers/pipelines/wan/pipeline_output.py index 88907ad0f0a1..7c7f8a4f0ad9 100644 --- a/src/diffusers/pipelines/wan/pipeline_output.py +++ b/src/diffusers/pipelines/wan/pipeline_output.py @@ -11,8 +11,8 @@ class WanPipelineOutput(BaseOutput): Output class for Wan pipelines. Args: - frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]): - List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): + list of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape `(batch_size, num_frames, channels, height, width)`. """ diff --git a/src/diffusers/pipelines/wan/pipeline_wan.py b/src/diffusers/pipelines/wan/pipeline_wan.py index 78fe71ea9138..d4edff01ad66 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan.py +++ b/src/diffusers/pipelines/wan/pipeline_wan.py @@ -13,7 +13,7 @@ # limitations under the License. import html -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import regex as re import torch @@ -76,7 +76,8 @@ def basic_clean(text): - text = ftfy.fix_text(text) + if is_ftfy_available(): + text = ftfy.fix_text(text) text = html.unescape(html.unescape(text)) return text.strip() @@ -133,9 +134,9 @@ def __init__( text_encoder: UMT5EncoderModel, vae: AutoencoderKLWan, scheduler: FlowMatchEulerDiscreteScheduler, - transformer: Optional[WanTransformer3DModel] = None, - transformer_2: Optional[WanTransformer3DModel] = None, - boundary_ratio: Optional[float] = None, + transformer: WanTransformer3DModel | None = None, + transformer_2: WanTransformer3DModel | None = None, + boundary_ratio: float | None = None, expand_timesteps: bool = False, # Wan2.2 ti2v ): super().__init__() @@ -156,11 +157,11 @@ def __init__( def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -197,23 +198,23 @@ def _get_t5_prompt_embeds( def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -328,10 +329,10 @@ def prepare_latents( height: int = 480, width: int = 832, num_frames: int = 81, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: if latents is not None: return latents.to(device=device, dtype=dtype) @@ -381,35 +382,33 @@ def attention_kwargs(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, height: int = 480, width: int = 832, num_frames: int = 81, num_inference_steps: int = 50, guidance_scale: float = 5.0, - guidance_scale_2: Optional[float] = None, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "np", + guidance_scale_2: float | None = None, + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "np", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, pass `prompt_embeds` instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to avoid during image generation. If not defined, pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale` < `1`). height (`int`, defaults to `480`): @@ -433,7 +432,7 @@ def __call__( and the pipeline's `boundary_ratio` are not None. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -456,7 +455,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -495,6 +494,22 @@ def __call__( num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1 num_frames = max(num_frames, 1) + patch_size = ( + self.transformer.config.patch_size + if self.transformer is not None + else self.transformer_2.config.patch_size + ) + h_multiple_of = self.vae_scale_factor_spatial * patch_size[1] + w_multiple_of = self.vae_scale_factor_spatial * patch_size[2] + calc_height = height // h_multiple_of * h_multiple_of + calc_width = width // w_multiple_of * w_multiple_of + if height != calc_height or width != calc_width: + logger.warning( + f"`height` and `width` must be multiples of ({h_multiple_of}, {w_multiple_of}) for proper patchification. " + f"Adjusting ({height}, {width}) -> ({calc_height}, {calc_width})." + ) + height, width = calc_height, calc_width + if self.config.boundary_ratio is not None and guidance_scale_2 is None: guidance_scale_2 = guidance_scale diff --git a/src/diffusers/pipelines/wan/pipeline_wan_animate.py b/src/diffusers/pipelines/wan/pipeline_wan_animate.py index c7c983b2f7d4..5806032c0142 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_animate.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_animate.py @@ -14,7 +14,7 @@ import html from copy import deepcopy -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import PIL import regex as re @@ -135,7 +135,7 @@ def prompt_clean(text): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -234,11 +234,11 @@ def __init__( def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -277,7 +277,7 @@ def _get_t5_prompt_embeds( def encode_image( self, image: PipelineImageInput, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): device = device or self._execution_device image = self.image_processor(images=image, return_tensors="pt").to(device) @@ -287,23 +287,23 @@ def encode_image( # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -463,9 +463,9 @@ def get_i2v_mask( latent_h: int, latent_w: int, mask_len: int = 1, - mask_pixel_values: Optional[torch.Tensor] = None, - dtype: Optional[torch.dtype] = None, - device: Union[str, torch.device] = "cuda", + mask_pixel_values: torch.Tensor | None = None, + dtype: torch.dtype | None = None, + device: str | torch.device = "cuda", ) -> torch.Tensor: # mask_pixel_values shape (if supplied): [B, C = 1, T, latent_h, latent_w] if mask_pixel_values is None: @@ -490,9 +490,9 @@ def prepare_reference_image_latents( image: torch.Tensor, batch_size: int = 1, sample_mode: int = "argmax", - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, ) -> torch.Tensor: # image shape: (B, C, H, W) or (B, C, T, H, W) dtype = dtype or self.vae.dtype @@ -537,9 +537,9 @@ def prepare_reference_image_latents( def prepare_prev_segment_cond_latents( self, - prev_segment_cond_video: Optional[torch.Tensor] = None, - background_video: Optional[torch.Tensor] = None, - mask_video: Optional[torch.Tensor] = None, + prev_segment_cond_video: torch.Tensor | None = None, + background_video: torch.Tensor | None = None, + mask_video: torch.Tensor | None = None, batch_size: int = 1, segment_frame_length: int = 77, start_frame: int = 0, @@ -549,9 +549,9 @@ def prepare_prev_segment_cond_latents( task: str = "animate", interpolation_mode: str = "bicubic", sample_mode: str = "argmax", - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, ) -> torch.Tensor: # prev_segment_cond_video shape: (B, C, T, H, W) in pixel space if supplied # background_video shape: (B, C, T, H, W) (same as prev_segment_cond_video shape) @@ -655,9 +655,9 @@ def prepare_pose_latents( pose_video: torch.Tensor, batch_size: int = 1, sample_mode: int = "argmax", - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, + generator: torch.Generator | list[torch.Generator] | None = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, ) -> torch.Tensor: # pose_video shape: (B, C, T, H, W) pose_video = pose_video.to(device=device, dtype=dtype if dtype is not None else self.vae.dtype) @@ -689,11 +689,11 @@ def prepare_latents( height: int = 720, width: int = 1280, num_frames: int = 77, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 latent_height = height // self.vae_scale_factor_spatial latent_width = width // self.vae_scale_factor_spatial @@ -712,7 +712,7 @@ def prepare_latents( return latents - def pad_video_frames(self, frames: List[Any], num_target_frames: int) -> List[Any]: + def pad_video_frames(self, frames: list[Any], num_target_frames: int) -> list[Any]: """ Pads an array-like video `frames` to `num_target_frames` using a "reflect"-like strategy. The frame dimension is assumed to be the first dimension. In the 1D case, we can visualize this strategy as follows: @@ -762,33 +762,31 @@ def attention_kwargs(self): def __call__( self, image: PipelineImageInput, - pose_video: List[PIL.Image.Image], - face_video: List[PIL.Image.Image], - background_video: Optional[List[PIL.Image.Image]] = None, - mask_video: Optional[List[PIL.Image.Image]] = None, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + pose_video: list[PIL.Image.Image], + face_video: list[PIL.Image.Image], + background_video: list[PIL.Image.Image] | None = None, + mask_video: list[PIL.Image.Image] | None = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, height: int = 720, width: int = 1280, segment_frame_length: int = 77, num_inference_steps: int = 20, mode: str = "animate", prev_segment_conditioning_frames: int = 1, - motion_encode_batch_size: Optional[int] = None, + motion_encode_batch_size: int | None = None, guidance_scale: float = 1.0, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - image_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "np", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + image_embeds: torch.Tensor | None = None, + output_type: str | None = "np", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int, None], PipelineCallback | MultiPipelineCallbacks] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" @@ -798,20 +796,20 @@ def __call__( image (`PipelineImageInput`): The input character image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`. - pose_video (`List[PIL.Image.Image]`): + pose_video (`list[PIL.Image.Image]`): The input pose video to condition the generation on. Must be a list of PIL images. - face_video (`List[PIL.Image.Image]`): + face_video (`list[PIL.Image.Image]`): The input face video to condition the generation on. Must be a list of PIL images. - background_video (`List[PIL.Image.Image]`, *optional*): + background_video (`list[PIL.Image.Image]`, *optional*): When mode is `"replace"`, the input background video to condition the generation on. Must be a list of PIL images. - mask_video (`List[PIL.Image.Image]`, *optional*): + mask_video (`list[PIL.Image.Image]`, *optional*): When mode is `"replace"`, the input mask video to condition the generation on. Must be a list of PIL images. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -844,7 +842,7 @@ def __call__( Animate inference. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): diff --git a/src/diffusers/pipelines/wan/pipeline_wan_i2v.py b/src/diffusers/pipelines/wan/pipeline_wan_i2v.py index b7fd0b05980f..f669e9b1d0ec 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_i2v.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_i2v.py @@ -13,7 +13,7 @@ # limitations under the License. import html -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable import PIL import regex as re @@ -112,7 +112,7 @@ def prompt_clean(text): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -174,7 +174,7 @@ def __init__( image_encoder: CLIPVisionModel = None, transformer: WanTransformer3DModel = None, transformer_2: WanTransformer3DModel = None, - boundary_ratio: Optional[float] = None, + boundary_ratio: float | None = None, expand_timesteps: bool = False, ): super().__init__() @@ -198,11 +198,11 @@ def __init__( def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 512, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -240,7 +240,7 @@ def _get_t5_prompt_embeds( def encode_image( self, image: PipelineImageInput, - device: Optional[torch.device] = None, + device: torch.device | None = None, ): device = device or self._execution_device image = self.image_processor(images=image, return_tensors="pt").to(device) @@ -250,23 +250,23 @@ def encode_image( # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -398,12 +398,12 @@ def prepare_latents( height: int = 480, width: int = 832, num_frames: int = 81, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - last_image: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + last_image: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 latent_height = height // self.vae_scale_factor_spatial latent_width = width // self.vae_scale_factor_spatial @@ -509,28 +509,26 @@ def attention_kwargs(self): def __call__( self, image: PipelineImageInput, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, height: int = 480, width: int = 832, num_frames: int = 81, num_inference_steps: int = 50, guidance_scale: float = 5.0, - guidance_scale_2: Optional[float] = None, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - image_embeds: Optional[torch.Tensor] = None, - last_image: Optional[torch.Tensor] = None, - output_type: Optional[str] = "np", + guidance_scale_2: float | None = None, + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + image_embeds: torch.Tensor | None = None, + last_image: torch.Tensor | None = None, + output_type: str | None = "np", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" @@ -539,10 +537,10 @@ def __call__( Args: image (`PipelineImageInput`): The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -567,7 +565,7 @@ def __call__( and the pipeline's `boundary_ratio` are not None. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -596,7 +594,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -637,6 +635,22 @@ def __call__( num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1 num_frames = max(num_frames, 1) + patch_size = ( + self.transformer.config.patch_size + if self.transformer is not None + else self.transformer_2.config.patch_size + ) + h_multiple_of = self.vae_scale_factor_spatial * patch_size[1] + w_multiple_of = self.vae_scale_factor_spatial * patch_size[2] + calc_height = height // h_multiple_of * h_multiple_of + calc_width = width // w_multiple_of * w_multiple_of + if height != calc_height or width != calc_width: + logger.warning( + f"`height` and `width` must be multiples of ({h_multiple_of}, {w_multiple_of}) for proper patchification. " + f"Adjusting ({height}, {width}) -> ({calc_height}, {calc_width})." + ) + height, width = calc_height, calc_width + if self.config.boundary_ratio is not None and guidance_scale_2 is None: guidance_scale_2 = guidance_scale diff --git a/src/diffusers/pipelines/wan/pipeline_wan_vace.py b/src/diffusers/pipelines/wan/pipeline_wan_vace.py index 351ae2e70563..c016eec1b535 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_vace.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_vace.py @@ -13,7 +13,7 @@ # limitations under the License. import html -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import PIL.Image import regex as re @@ -126,7 +126,7 @@ def prompt_clean(text): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -183,7 +183,7 @@ def __init__( scheduler: FlowMatchEulerDiscreteScheduler, transformer: WanVACETransformer3DModel = None, transformer_2: WanVACETransformer3DModel = None, - boundary_ratio: Optional[float] = None, + boundary_ratio: float | None = None, ): super().__init__() @@ -203,11 +203,11 @@ def __init__( # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -245,23 +245,23 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -413,15 +413,15 @@ def check_inputs( def preprocess_conditions( self, - video: Optional[List[PipelineImageInput]] = None, - mask: Optional[List[PipelineImageInput]] = None, - reference_images: Optional[Union[PIL.Image.Image, List[PIL.Image.Image], List[List[PIL.Image.Image]]]] = None, + video: list[PipelineImageInput] | None = None, + mask: list[PipelineImageInput] | None = None, + reference_images: PIL.Image.Image | list[PIL.Image.Image] | list[list[PIL.Image.Image]] | None = None, batch_size: int = 1, height: int = 480, width: int = 832, num_frames: int = 81, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, ): if video is not None: base = self.vae_scale_factor_spatial * ( @@ -515,9 +515,9 @@ def prepare_video_latents( self, video: torch.Tensor, mask: torch.Tensor, - reference_images: Optional[List[List[torch.Tensor]]] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - device: Optional[torch.device] = None, + reference_images: list[list[torch.Tensor]] | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + device: torch.device | None = None, ) -> torch.Tensor: device = device or self._execution_device @@ -581,8 +581,8 @@ def prepare_video_latents( def prepare_masks( self, mask: torch.Tensor, - reference_images: Optional[List[torch.Tensor]] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + reference_images: list[torch.Tensor] | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, ) -> torch.Tensor: if isinstance(generator, list): # TODO: support this @@ -637,10 +637,10 @@ def prepare_latents( height: int = 480, width: int = 832, num_frames: int = 81, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, ) -> torch.Tensor: if latents is not None: return latents.to(device=device, dtype=dtype) @@ -690,60 +690,58 @@ def attention_kwargs(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, - video: Optional[List[PipelineImageInput]] = None, - mask: Optional[List[PipelineImageInput]] = None, - reference_images: Optional[List[PipelineImageInput]] = None, - conditioning_scale: Union[float, List[float], torch.Tensor] = 1.0, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, + video: list[PipelineImageInput] | None = None, + mask: list[PipelineImageInput] | None = None, + reference_images: list[PipelineImageInput] | None = None, + conditioning_scale: float | list[float] | torch.Tensor = 1.0, height: int = 480, width: int = 832, num_frames: int = 81, num_inference_steps: int = 50, guidance_scale: float = 5.0, - guidance_scale_2: Optional[float] = None, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "np", + guidance_scale_2: float | None = None, + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "np", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds` instead. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). - video (`List[PIL.Image.Image]`, *optional*): + video (`list[PIL.Image.Image]`, *optional*): The input video or videos to be used as a starting point for the generation. The video should be a list of PIL images, a numpy array, or a torch tensor. Currently, the pipeline only supports generating one video at a time. - mask (`List[PIL.Image.Image]`, *optional*): + mask (`list[PIL.Image.Image]`, *optional*): The input mask defines which video regions to condition on and which to generate. Black areas in the mask indicate conditioning regions, while white areas indicate regions for generation. The mask should be a list of PIL images, a numpy array, or a torch tensor. Currently supports generating a single video at a time. - reference_images (`List[PIL.Image.Image]`, *optional*): + reference_images (`list[PIL.Image.Image]`, *optional*): A list of one or more reference images as extra conditioning for the generation. For example, if you are trying to inpaint a video to change the character, you can pass reference images of the new character here. Refer to the Diffusers [examples](https://github.com/huggingface/diffusers/pull/11582) and original [user guide](https://github.com/ali-vilab/VACE/blob/0897c6d055d7d9ea9e191dce763006664d9780f8/UserGuide.md) for a full list of supported tasks and use cases. - conditioning_scale (`float`, `List[float]`, `torch.Tensor`, defaults to `1.0`): + conditioning_scale (`float`, `list[float]`, `torch.Tensor`, defaults to `1.0`): The conditioning scale to be applied when adding the control conditioning latent stream to the denoising latent stream in each control layer of the model. If a float is provided, it will be applied uniformly to all layers. If a list or tensor is provided, it should have the same length as the number @@ -769,7 +767,7 @@ def __call__( and the pipeline's `boundary_ratio` are not None. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -792,7 +790,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/wan/pipeline_wan_video2video.py b/src/diffusers/pipelines/wan/pipeline_wan_video2video.py index a976126da7fe..3d7c5297f4c4 100644 --- a/src/diffusers/pipelines/wan/pipeline_wan_video2video.py +++ b/src/diffusers/pipelines/wan/pipeline_wan_video2video.py @@ -14,7 +14,7 @@ import html import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import regex as re import torch @@ -100,10 +100,10 @@ def prompt_clean(text): # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -118,15 +118,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -159,7 +159,7 @@ def retrieve_timesteps( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -221,11 +221,11 @@ def __init__( # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline._get_t5_prompt_embeds def _get_t5_prompt_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, num_videos_per_prompt: int = 1, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -263,23 +263,23 @@ def _get_t5_prompt_embeds( # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 226, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, ): r""" Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): prompt to be encoded - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). @@ -390,16 +390,16 @@ def check_inputs( def prepare_latents( self, - video: Optional[torch.Tensor] = None, + video: torch.Tensor | None = None, batch_size: int = 1, num_channels_latents: int = 16, height: int = 480, width: int = 832, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - generator: Optional[torch.Generator] = None, - latents: Optional[torch.Tensor] = None, - timestep: Optional[torch.Tensor] = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, + timestep: torch.Tensor | None = None, ): if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( @@ -480,34 +480,32 @@ def attention_kwargs(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - video: List[Image.Image] = None, - prompt: Union[str, List[str]] = None, - negative_prompt: Union[str, List[str]] = None, + video: list[Image.Image] = None, + prompt: str | list[str] = None, + negative_prompt: str | list[str] = None, height: int = 480, width: int = 832, num_inference_steps: int = 50, - timesteps: Optional[List[int]] = None, + timesteps: list[int] | None = None, guidance_scale: float = 5.0, strength: float = 0.8, - num_videos_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - output_type: Optional[str] = "np", + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "np", return_dict: bool = True, - attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[ - Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] - ] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | PipelineCallback | MultiPipelineCallbacks | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" The call function to the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds` instead. height (`int`, defaults to `480`): @@ -529,7 +527,7 @@ def __call__( Higher strength leads to more differences between original image and generated video. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -552,7 +550,7 @@ def __call__( each denoising step during the inference. with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -622,7 +620,13 @@ def __call__( negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype) # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) + if XLA_AVAILABLE: + timestep_device = "cpu" + else: + timestep_device = device + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, timestep_device, timesteps + ) timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, timesteps, strength, device) latent_timestep = timesteps[:1].repeat(batch_size * num_videos_per_prompt) self._num_timesteps = len(timesteps) diff --git a/src/diffusers/pipelines/wuerstchen/modeling_paella_vq_model.py b/src/diffusers/pipelines/wuerstchen/modeling_paella_vq_model.py index 5ab206b15176..932c7ac618f6 100644 --- a/src/diffusers/pipelines/wuerstchen/modeling_paella_vq_model.py +++ b/src/diffusers/pipelines/wuerstchen/modeling_paella_vq_model.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Union import torch import torch.nn as nn @@ -142,7 +141,7 @@ def encode(self, x: torch.Tensor, return_dict: bool = True) -> VQEncoderOutput: @apply_forward_hook def decode( self, h: torch.Tensor, force_not_quantize: bool = True, return_dict: bool = True - ) -> Union[DecoderOutput, torch.Tensor]: + ) -> DecoderOutput | torch.Tensor: if not force_not_quantize: quant, _, _ = self.vquantizer(h) else: @@ -155,7 +154,7 @@ def decode( return DecoderOutput(sample=dec) - def forward(self, sample: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: + def forward(self, sample: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: r""" Args: sample (`torch.Tensor`): Input sample. diff --git a/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py b/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py index bbdb60471fd1..cce05c189201 100644 --- a/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py +++ b/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import numpy as np import torch @@ -156,7 +156,7 @@ def encode_prompt( uncond_text_encoder_hidden_states = None if do_classifier_free_guidance: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -217,33 +217,33 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - image_embeddings: Union[torch.Tensor, List[torch.Tensor]], - prompt: Union[str, List[str]] = None, + image_embeddings: torch.Tensor | list[torch.Tensor], + prompt: str | list[str] = None, num_inference_steps: int = 12, - timesteps: Optional[List[float]] = None, + timesteps: list[float] | None = None, guidance_scale: float = 0.0, - negative_prompt: Optional[Union[str, List[str]]] = None, + negative_prompt: str | list[str] | None = None, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): """ Function invoked when calling the pipeline for generation. Args: - image_embedding (`torch.Tensor` or `List[torch.Tensor]`): + image_embedding (`torch.Tensor` or `list[torch.Tensor]`): Image Embeddings either extracted from an image or generated by a Prior Model. - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. num_inference_steps (`int`, *optional*, defaults to 12): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps` timesteps are used. Must be in descending order. guidance_scale (`float`, *optional*, defaults to 0.0): @@ -252,12 +252,12 @@ def __call__( equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `decoder_guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `decoder_guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -274,7 +274,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py b/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py index c54c1fefe8fe..16300a7c71d2 100644 --- a/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py +++ b/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import torch from transformers import CLIPTextModel, CLIPTokenizer @@ -110,10 +110,10 @@ def __init__( vqgan=vqgan, ) - def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None): + def enable_xformers_memory_efficient_attention(self, attention_op: Callable | None = None): self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op) - def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None): + def enable_model_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None): r""" Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward` @@ -123,7 +123,7 @@ def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[t self.prior_pipe.enable_model_cpu_offload(gpu_id=gpu_id, device=device) self.decoder_pipe.enable_model_cpu_offload(gpu_id=gpu_id, device=device) - def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = None): + def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None): r""" Offloads all models (`unet`, `text_encoder`, `vae`, and `safety checker` state dicts) to CPU using 🤗 Accelerate, significantly reducing memory usage. Models are moved to a `torch.device('meta')` and loaded on a @@ -145,36 +145,36 @@ def set_progress_bar_config(self, **kwargs): @replace_example_docstring(TEXT2IMAGE_EXAMPLE_DOC_STRING) def __call__( self, - prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] | None = None, height: int = 512, width: int = 512, prior_num_inference_steps: int = 60, - prior_timesteps: Optional[List[float]] = None, + prior_timesteps: list[float] | None = None, prior_guidance_scale: float = 4.0, num_inference_steps: int = 12, - decoder_timesteps: Optional[List[float]] = None, + decoder_timesteps: list[float] | None = None, decoder_guidance_scale: float = 0.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt: str | list[str] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, num_images_per_prompt: int = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pil", + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", return_dict: bool = True, - prior_callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - prior_callback_on_step_end_tensor_inputs: List[str] = ["latents"], - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + prior_callback_on_step_end: Callable[[int, int], None] | None = None, + prior_callback_on_step_end_tensor_inputs: list[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation for the prior and decoder. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). prompt_embeds (`torch.Tensor`, *optional*): @@ -196,7 +196,7 @@ def __call__( equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `prior_guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - prior_num_inference_steps (`Union[int, Dict[float, int]]`, *optional*, defaults to 60): + prior_num_inference_steps (`int | dict[float, int]`, *optional*, defaults to 60): The number of prior denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. For more specific timestep spacing, you can pass customized `prior_timesteps` @@ -204,10 +204,10 @@ def __call__( The number of decoder denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. For more specific timestep spacing, you can pass customized `timesteps` - prior_timesteps (`List[float]`, *optional*): + prior_timesteps (`list[float]`, *optional*): Custom timesteps to use for the denoising process for the prior. If not defined, equal spaced `prior_num_inference_steps` timesteps are used. Must be in descending order. - decoder_timesteps (`List[float]`, *optional*): + decoder_timesteps (`list[float]`, *optional*): Custom timesteps to use for the denoising process for the decoder. If not defined, equal spaced `num_inference_steps` timesteps are used. Must be in descending order. decoder_guidance_scale (`float`, *optional*, defaults to 0.0): @@ -216,7 +216,7 @@ def __call__( of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -232,7 +232,7 @@ def __call__( A function that calls at the end of each denoising steps during the inference. The function is called with the following arguments: `prior_callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. - prior_callback_on_step_end_tensor_inputs (`List`, *optional*): + prior_callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `prior_callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. @@ -241,7 +241,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py b/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py index e138b6e805c8..e79fcf8378aa 100644 --- a/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py +++ b/src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py @@ -14,7 +14,7 @@ from dataclasses import dataclass from math import ceil -from typing import Callable, Dict, List, Optional, Union +from typing import Callable import numpy as np import torch @@ -67,7 +67,7 @@ class WuerstchenPriorPipelineOutput(BaseOutput): """ - image_embeddings: Union[torch.Tensor, np.ndarray] + image_embeddings: torch.Tensor | np.ndarray class WuerstchenPriorPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin): @@ -145,8 +145,8 @@ def encode_prompt( do_classifier_free_guidance, prompt=None, negative_prompt=None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, ): if prompt is not None and isinstance(prompt, str): batch_size = 1 @@ -191,7 +191,7 @@ def encode_prompt( prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) if negative_prompt_embeds is None and do_classifier_free_guidance: - uncond_tokens: List[str] + uncond_tokens: list[str] if negative_prompt is None: uncond_tokens = [""] * batch_size elif type(prompt) is not type(negative_prompt): @@ -290,29 +290,29 @@ def num_timesteps(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Optional[Union[str, List[str]]] = None, + prompt: str | list[str] | None = None, height: int = 1024, width: int = 1024, num_inference_steps: int = 60, - timesteps: List[float] = None, + timesteps: list[float] = None, guidance_scale: float = 8.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "pt", + negative_prompt: str | list[str] | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pt", return_dict: bool = True, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], **kwargs, ): """ Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`): + prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. height (`int`, *optional*, defaults to 1024): The height in pixels of the generated image. @@ -321,7 +321,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 60): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps` timesteps are used. Must be in descending order. guidance_scale (`float`, *optional*, defaults to 8.0): @@ -330,7 +330,7 @@ def __call__( equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `decoder_guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `decoder_guidance_scale` is less than `1`). prompt_embeds (`torch.Tensor`, *optional*): @@ -342,7 +342,7 @@ def __call__( argument. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): @@ -359,7 +359,7 @@ def __call__( with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): + callback_on_step_end_tensor_inputs (`list`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. diff --git a/src/diffusers/pipelines/z_image/__init__.py b/src/diffusers/pipelines/z_image/__init__.py index 78bd3bfacbec..14629a6e2160 100644 --- a/src/diffusers/pipelines/z_image/__init__.py +++ b/src/diffusers/pipelines/z_image/__init__.py @@ -26,6 +26,7 @@ _import_structure["pipeline_z_image_controlnet"] = ["ZImageControlNetPipeline"] _import_structure["pipeline_z_image_controlnet_inpaint"] = ["ZImageControlNetInpaintPipeline"] _import_structure["pipeline_z_image_img2img"] = ["ZImageImg2ImgPipeline"] + _import_structure["pipeline_z_image_inpaint"] = ["ZImageInpaintPipeline"] _import_structure["pipeline_z_image_omni"] = ["ZImageOmniPipeline"] @@ -42,6 +43,7 @@ from .pipeline_z_image_controlnet import ZImageControlNetPipeline from .pipeline_z_image_controlnet_inpaint import ZImageControlNetInpaintPipeline from .pipeline_z_image_img2img import ZImageImg2ImgPipeline + from .pipeline_z_image_inpaint import ZImageInpaintPipeline from .pipeline_z_image_omni import ZImageOmniPipeline else: import sys diff --git a/src/diffusers/pipelines/z_image/pipeline_output.py b/src/diffusers/pipelines/z_image/pipeline_output.py index 69a320fc036a..5cbcbee4bfab 100644 --- a/src/diffusers/pipelines/z_image/pipeline_output.py +++ b/src/diffusers/pipelines/z_image/pipeline_output.py @@ -13,7 +13,6 @@ # limitations under the License. from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -27,9 +26,9 @@ class ZImagePipelineOutput(BaseOutput): Output class for Z-Image pipelines. Args: - images (`List[PIL.Image.Image]` or `np.ndarray`) + images (`list[PIL.Image.Image]` or `np.ndarray`) List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline. """ - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image, np.ndarray] diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py index 82bdd7d361b7..959368ec1cd1 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch from transformers import AutoTokenizer, PreTrainedModel @@ -77,10 +77,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -95,15 +95,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -163,12 +163,12 @@ def __init__( def encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + device: torch.device | None = None, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[List[torch.FloatTensor]] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt: str | list[str] | None = None, + prompt_embeds: list[torch.FloatTensor] | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 512, ): prompt = [prompt] if isinstance(prompt, str) else prompt @@ -197,11 +197,11 @@ def encode_prompt( def _encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, - prompt_embeds: Optional[List[torch.FloatTensor]] = None, + prompt: str | list[str], + device: torch.device | None = None, + prompt_embeds: list[torch.FloatTensor] | None = None, max_sequence_length: int = 512, - ) -> List[torch.FloatTensor]: + ) -> list[torch.FloatTensor]: device = device or self._execution_device if prompt_embeds is not None: @@ -276,7 +276,7 @@ def guidance_scale(self): @property def do_classifier_free_guidance(self): - return self._guidance_scale > 1 + return self._guidance_scale > 0 @property def joint_attention_kwargs(self): @@ -294,32 +294,32 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 5.0, cfg_normalization: bool = False, cfg_truncation: float = 1.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[List[torch.FloatTensor]] = None, - negative_prompt_embeds: Optional[List[torch.FloatTensor]] = None, - output_type: Optional[str] = "pil", + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: list[torch.FloatTensor] | None = None, + negative_prompt_embeds: list[torch.FloatTensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. height (`int`, *optional*, defaults to 1024): @@ -329,7 +329,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -343,23 +343,23 @@ def __call__( Whether to apply configuration normalization. cfg_truncation (`float`, *optional*, defaults to 1.0): The truncation value for configuration. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image generation. Can be used to tweak the same generation with different prompts. If not provided, a latents tensor will be generated by sampling using the supplied random `generator`. - prompt_embeds (`List[torch.FloatTensor]`, *optional*): + prompt_embeds (`list[torch.FloatTensor]`, *optional*): Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, text embeddings will be generated from `prompt` input argument. - negative_prompt_embeds (`List[torch.FloatTensor]`, *optional*): + negative_prompt_embeds (`list[torch.FloatTensor]`, *optional*): Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py index 5e26862b018e..1e49737bb5b0 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py @@ -13,13 +13,13 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch from transformers import AutoTokenizer, PreTrainedModel from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin +from ...loaders import FromSingleFileMixin, ZImageLoraLoaderMixin from ...models.autoencoders import AutoencoderKL from ...models.controlnets import ZImageControlNetModel from ...models.transformers import ZImageTransformer2DModel @@ -58,14 +58,13 @@ >>> # torch_dtype=torch.bfloat16, >>> # ) - >>> # 2.0 - `config` is required + >>> # 2.0 >>> # controlnet = ZImageControlNetModel.from_single_file( >>> # hf_hub_download( >>> # "alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union-2.0", >>> # filename="Z-Image-Turbo-Fun-Controlnet-Union-2.0.safetensors", >>> # ), >>> # torch_dtype=torch.bfloat16, - >>> # config="hlky/Z-Image-Turbo-Fun-Controlnet-Union-2.0", >>> # ) >>> pipe = ZImageControlNetPipeline.from_pretrained( @@ -114,7 +113,7 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -129,10 +128,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -147,15 +146,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -186,7 +185,7 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class ZImageControlNetPipeline(DiffusionPipeline, FromSingleFileMixin): +class ZImageControlNetPipeline(DiffusionPipeline, ZImageLoraLoaderMixin, FromSingleFileMixin): model_cpu_offload_seq = "text_encoder->transformer->vae" _optional_components = [] _callback_tensor_inputs = ["latents", "prompt_embeds"] @@ -218,12 +217,12 @@ def __init__( def encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + device: torch.device | None = None, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[List[torch.FloatTensor]] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt: str | list[str] | None = None, + prompt_embeds: list[torch.FloatTensor] | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 512, ): prompt = [prompt] if isinstance(prompt, str) else prompt @@ -252,11 +251,11 @@ def encode_prompt( def _encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, - prompt_embeds: Optional[List[torch.FloatTensor]] = None, + prompt: str | list[str], + device: torch.device | None = None, + prompt_embeds: list[torch.FloatTensor] | None = None, max_sequence_length: int = 512, - ) -> List[torch.FloatTensor]: + ) -> list[torch.FloatTensor]: device = device or self._execution_device if prompt_embeds is not None: @@ -366,7 +365,7 @@ def guidance_scale(self): @property def do_classifier_free_guidance(self): - return self._guidance_scale > 1 + return self._guidance_scale > 0 @property def joint_attention_kwargs(self): @@ -384,34 +383,34 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 5.0, control_image: PipelineImageInput = None, - controlnet_conditioning_scale: Union[float, List[float]] = 0.75, + controlnet_conditioning_scale: float | list[float] = 0.75, cfg_normalization: bool = False, cfg_truncation: float = 1.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[List[torch.FloatTensor]] = None, - negative_prompt_embeds: Optional[List[torch.FloatTensor]] = None, - output_type: Optional[str] = "pil", + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: list[torch.FloatTensor] | None = None, + negative_prompt_embeds: list[torch.FloatTensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. height (`int`, *optional*, defaults to 1024): @@ -421,7 +420,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -435,23 +434,23 @@ def __call__( Whether to apply configuration normalization. cfg_truncation (`float`, *optional*, defaults to 1.0): The truncation value for configuration. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image generation. Can be used to tweak the same generation with different prompts. If not provided, a latents tensor will be generated by sampling using the supplied random `generator`. - prompt_embeds (`List[torch.FloatTensor]`, *optional*): + prompt_embeds (`list[torch.FloatTensor]`, *optional*): Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, text embeddings will be generated from `prompt` input argument. - negative_prompt_embeds (`List[torch.FloatTensor]`, *optional*): + negative_prompt_embeds (`list[torch.FloatTensor]`, *optional*): Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. @@ -636,10 +635,12 @@ def __call__( latent_model_input = latents_typed.repeat(2, 1, 1, 1) prompt_embeds_model_input = prompt_embeds + negative_prompt_embeds timestep_model_input = timestep.repeat(2) + control_image_input = control_image.repeat(2, 1, 1, 1, 1) else: latent_model_input = latents.to(self.transformer.dtype) prompt_embeds_model_input = prompt_embeds timestep_model_input = timestep + control_image_input = control_image latent_model_input = latent_model_input.unsqueeze(2) latent_model_input_list = list(latent_model_input.unbind(dim=0)) @@ -648,7 +649,7 @@ def __call__( latent_model_input_list, timestep_model_input, prompt_embeds_model_input, - control_image, + control_image_input, conditioning_scale=controlnet_conditioning_scale, ) diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py index 73ea7d0fddec..09f9b2395458 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py @@ -13,14 +13,14 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch import torch.nn.functional as F from transformers import AutoTokenizer, PreTrainedModel from ...image_processor import PipelineImageInput, VaeImageProcessor -from ...loaders import FromSingleFileMixin +from ...loaders import FromSingleFileMixin, ZImageLoraLoaderMixin from ...models.autoencoders import AutoencoderKL from ...models.controlnets import ZImageControlNetModel from ...models.transformers import ZImageTransformer2DModel @@ -50,14 +50,13 @@ ... torch_dtype=torch.bfloat16, ... ) - >>> # 2.0 - `config` is required + >>> # 2.0 >>> # controlnet = ZImageControlNetModel.from_single_file( >>> # hf_hub_download( >>> # "alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union-2.0", >>> # filename="Z-Image-Turbo-Fun-Controlnet-Union-2.0.safetensors", >>> # ), >>> # torch_dtype=torch.bfloat16, - >>> # config="hlky/Z-Image-Turbo-Fun-Controlnet-Union-2.0", >>> # ) >>> pipe = ZImageControlNetInpaintPipeline.from_pretrained( @@ -114,7 +113,7 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -129,10 +128,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -147,15 +146,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -186,7 +185,7 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class ZImageControlNetInpaintPipeline(DiffusionPipeline, FromSingleFileMixin): +class ZImageControlNetInpaintPipeline(DiffusionPipeline, ZImageLoraLoaderMixin, FromSingleFileMixin): model_cpu_offload_seq = "text_encoder->transformer->vae" _optional_components = [] _callback_tensor_inputs = ["latents", "prompt_embeds"] @@ -225,12 +224,12 @@ def __init__( def encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + device: torch.device | None = None, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[List[torch.FloatTensor]] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt: str | list[str] | None = None, + prompt_embeds: list[torch.FloatTensor] | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 512, ): prompt = [prompt] if isinstance(prompt, str) else prompt @@ -259,11 +258,11 @@ def encode_prompt( def _encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, - prompt_embeds: Optional[List[torch.FloatTensor]] = None, + prompt: str | list[str], + device: torch.device | None = None, + prompt_embeds: list[torch.FloatTensor] | None = None, max_sequence_length: int = 512, - ) -> List[torch.FloatTensor]: + ) -> list[torch.FloatTensor]: device = device or self._execution_device if prompt_embeds is not None: @@ -373,7 +372,7 @@ def guidance_scale(self): @property def do_classifier_free_guidance(self): - return self._guidance_scale > 1 + return self._guidance_scale > 0 @property def joint_attention_kwargs(self): @@ -391,36 +390,36 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 5.0, image: PipelineImageInput = None, mask_image: PipelineImageInput = None, control_image: PipelineImageInput = None, - controlnet_conditioning_scale: Union[float, List[float]] = 0.75, + controlnet_conditioning_scale: float | list[float] = 0.75, cfg_normalization: bool = False, cfg_truncation: float = 1.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[List[torch.FloatTensor]] = None, - negative_prompt_embeds: Optional[List[torch.FloatTensor]] = None, - output_type: Optional[str] = "pil", + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: list[torch.FloatTensor] | None = None, + negative_prompt_embeds: list[torch.FloatTensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. height (`int`, *optional*, defaults to 1024): @@ -430,7 +429,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -444,23 +443,23 @@ def __call__( Whether to apply configuration normalization. cfg_truncation (`float`, *optional*, defaults to 1.0): The truncation value for configuration. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image generation. Can be used to tweak the same generation with different prompts. If not provided, a latents tensor will be generated by sampling using the supplied random `generator`. - prompt_embeds (`List[torch.FloatTensor]`, *optional*): + prompt_embeds (`list[torch.FloatTensor]`, *optional*): Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, text embeddings will be generated from `prompt` input argument. - negative_prompt_embeds (`List[torch.FloatTensor]`, *optional*): + negative_prompt_embeds (`list[torch.FloatTensor]`, *optional*): Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. @@ -658,10 +657,12 @@ def __call__( latent_model_input = latents_typed.repeat(2, 1, 1, 1) prompt_embeds_model_input = prompt_embeds + negative_prompt_embeds timestep_model_input = timestep.repeat(2) + control_image_input = control_image.repeat(2, 1, 1, 1, 1) else: latent_model_input = latents.to(self.transformer.dtype) prompt_embeds_model_input = prompt_embeds timestep_model_input = timestep + control_image_input = control_image latent_model_input = latent_model_input.unsqueeze(2) latent_model_input_list = list(latent_model_input.unbind(dim=0)) @@ -670,7 +671,7 @@ def __call__( latent_model_input_list, timestep_model_input, prompt_embeds_model_input, - control_image, + control_image_input, conditioning_scale=controlnet_conditioning_scale, ) diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_img2img.py b/src/diffusers/pipelines/z_image/pipeline_z_image_img2img.py index 2b3e80a2082b..ee57f51dd957 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image_img2img.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image_img2img.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import torch from transformers import AutoTokenizer, PreTrainedModel @@ -74,7 +74,7 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) @@ -89,10 +89,10 @@ def retrieve_latents( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -107,15 +107,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -192,12 +192,12 @@ def __init__( # Copied from diffusers.pipelines.z_image.pipeline_z_image.ZImagePipeline.encode_prompt def encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + device: torch.device | None = None, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[List[torch.FloatTensor]] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt: str | list[str] | None = None, + prompt_embeds: list[torch.FloatTensor] | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 512, ): prompt = [prompt] if isinstance(prompt, str) else prompt @@ -227,11 +227,11 @@ def encode_prompt( # Copied from diffusers.pipelines.z_image.pipeline_z_image.ZImagePipeline._encode_prompt def _encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, - prompt_embeds: Optional[List[torch.FloatTensor]] = None, + prompt: str | list[str], + device: torch.device | None = None, + prompt_embeds: list[torch.FloatTensor] | None = None, max_sequence_length: int = 512, - ) -> List[torch.FloatTensor]: + ) -> list[torch.FloatTensor]: device = device or self._execution_device if prompt_embeds is not None: @@ -347,7 +347,7 @@ def guidance_scale(self): @property def do_classifier_free_guidance(self): - return self._guidance_scale > 1 + return self._guidance_scale > 0 @property def joint_attention_kwargs(self): @@ -365,37 +365,37 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - prompt: Union[str, List[str]] = None, + prompt: str | list[str] = None, image: PipelineImageInput = None, strength: float = 0.6, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 5.0, cfg_normalization: bool = False, cfg_truncation: float = 1.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[List[torch.FloatTensor]] = None, - negative_prompt_embeds: Optional[List[torch.FloatTensor]] = None, - output_type: Optional[str] = "pil", + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: list[torch.FloatTensor] | None = None, + negative_prompt_embeds: list[torch.FloatTensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for image-to-image generation. Args: - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]`. If it's a tensor or a list of tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or @@ -413,7 +413,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -427,23 +427,23 @@ def __call__( Whether to apply configuration normalization. cfg_truncation (`float`, *optional*, defaults to 1.0): The truncation value for configuration. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image generation. Can be used to tweak the same generation with different prompts. If not provided, a latents tensor will be generated by sampling using the supplied random `generator`. - prompt_embeds (`List[torch.FloatTensor]`, *optional*): + prompt_embeds (`list[torch.FloatTensor]`, *optional*): Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, text embeddings will be generated from `prompt` input argument. - negative_prompt_embeds (`List[torch.FloatTensor]`, *optional*): + negative_prompt_embeds (`list[torch.FloatTensor]`, *optional*): Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_inpaint.py b/src/diffusers/pipelines/z_image/pipeline_z_image_inpaint.py new file mode 100644 index 000000000000..e740a48e65ec --- /dev/null +++ b/src/diffusers/pipelines/z_image/pipeline_z_image_inpaint.py @@ -0,0 +1,932 @@ +# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +from typing import Any, Callable, Dict, List, Optional, Union + +import torch +from transformers import AutoTokenizer, PreTrainedModel + +from ...image_processor import PipelineImageInput, VaeImageProcessor +from ...loaders import FromSingleFileMixin, ZImageLoraLoaderMixin +from ...models.autoencoders import AutoencoderKL +from ...models.transformers import ZImageTransformer2DModel +from ...pipelines.pipeline_utils import DiffusionPipeline +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import is_torch_xla_available, logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from .pipeline_output import ZImagePipelineOutput + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import ZImageInpaintPipeline + >>> from diffusers.utils import load_image + + >>> pipe = ZImageInpaintPipeline.from_pretrained("Tongyi-MAI/Z-Image-Turbo", torch_dtype=torch.bfloat16) + >>> pipe.to("cuda") + + >>> url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg" + >>> init_image = load_image(url).resize((1024, 1024)) + + >>> # Create a mask (white = inpaint, black = preserve) + >>> import numpy as np + >>> from PIL import Image + + >>> mask = np.zeros((1024, 1024), dtype=np.uint8) + >>> mask[256:768, 256:768] = 255 # Inpaint center region + >>> mask_image = Image.fromarray(mask) + + >>> prompt = "A beautiful lake with mountains in the background" + >>> image = pipe( + ... prompt, + ... image=init_image, + ... mask_image=mask_image, + ... strength=1.0, + ... num_inference_steps=9, + ... guidance_scale=0.0, + ... generator=torch.Generator("cuda").manual_seed(42), + ... ).images[0] + >>> image.save("zimage_inpaint.png") + ``` +""" + + +# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents +def retrieve_latents( + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" +): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`list[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`list[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class ZImageInpaintPipeline(DiffusionPipeline, ZImageLoraLoaderMixin, FromSingleFileMixin): + r""" + The ZImage pipeline for inpainting. + + Args: + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + vae ([`AutoencoderKL`]): + Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. + text_encoder ([`PreTrainedModel`]): + A text encoder model to encode text prompts. + tokenizer ([`AutoTokenizer`]): + A tokenizer to tokenize text prompts. + transformer ([`ZImageTransformer2DModel`]): + A ZImage transformer model to denoise the encoded image latents. + """ + + model_cpu_offload_seq = "text_encoder->transformer->vae" + _optional_components = [] + _callback_tensor_inputs = ["latents", "prompt_embeds", "mask", "masked_image_latents"] + + def __init__( + self, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKL, + text_encoder: PreTrainedModel, + tokenizer: AutoTokenizer, + transformer: ZImageTransformer2DModel, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + scheduler=scheduler, + transformer=transformer, + ) + self.vae_scale_factor = ( + 2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8 + ) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) + self.mask_processor = VaeImageProcessor( + vae_scale_factor=self.vae_scale_factor * 2, + do_normalize=False, + do_binarize=True, + do_convert_grayscale=True, + ) + + # Copied from diffusers.pipelines.z_image.pipeline_z_image.ZImagePipeline.encode_prompt + def encode_prompt( + self, + prompt: str | list[str], + device: torch.device | None = None, + do_classifier_free_guidance: bool = True, + negative_prompt: str | list[str] | None = None, + prompt_embeds: list[torch.FloatTensor] | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, + max_sequence_length: int = 512, + ): + prompt = [prompt] if isinstance(prompt, str) else prompt + prompt_embeds = self._encode_prompt( + prompt=prompt, + device=device, + prompt_embeds=prompt_embeds, + max_sequence_length=max_sequence_length, + ) + + if do_classifier_free_guidance: + if negative_prompt is None: + negative_prompt = ["" for _ in prompt] + else: + negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + assert len(prompt) == len(negative_prompt) + negative_prompt_embeds = self._encode_prompt( + prompt=negative_prompt, + device=device, + prompt_embeds=negative_prompt_embeds, + max_sequence_length=max_sequence_length, + ) + else: + negative_prompt_embeds = [] + return prompt_embeds, negative_prompt_embeds + + # Copied from diffusers.pipelines.z_image.pipeline_z_image.ZImagePipeline._encode_prompt + def _encode_prompt( + self, + prompt: str | list[str], + device: torch.device | None = None, + prompt_embeds: list[torch.FloatTensor] | None = None, + max_sequence_length: int = 512, + ) -> list[torch.FloatTensor]: + device = device or self._execution_device + + if prompt_embeds is not None: + return prompt_embeds + + if isinstance(prompt, str): + prompt = [prompt] + + for i, prompt_item in enumerate(prompt): + messages = [ + {"role": "user", "content": prompt_item}, + ] + prompt_item = self.tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=True, + ) + prompt[i] = prompt_item + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + return_tensors="pt", + ) + + text_input_ids = text_inputs.input_ids.to(device) + prompt_masks = text_inputs.attention_mask.to(device).bool() + + prompt_embeds = self.text_encoder( + input_ids=text_input_ids, + attention_mask=prompt_masks, + output_hidden_states=True, + ).hidden_states[-2] + + embeddings_list = [] + + for i in range(len(prompt_embeds)): + embeddings_list.append(prompt_embeds[i][prompt_masks[i]]) + + return embeddings_list + + # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3_img2img.StableDiffusion3Img2ImgPipeline.get_timesteps + def get_timesteps(self, num_inference_steps, strength, device): + # get the original timestep using init_timestep + init_timestep = min(num_inference_steps * strength, num_inference_steps) + + t_start = int(max(num_inference_steps - init_timestep, 0)) + timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :] + if hasattr(self.scheduler, "set_begin_index"): + self.scheduler.set_begin_index(t_start * self.scheduler.order) + + return timesteps, num_inference_steps - t_start + + def prepare_mask_latents( + self, + mask, + masked_image, + batch_size, + height, + width, + dtype, + device, + generator, + ): + """Prepare mask and masked image latents for inpainting. + + Args: + mask: Binary mask tensor where 1 = inpaint region, 0 = preserve region. + masked_image: Original image with masked regions zeroed out. + batch_size: Number of images to generate. + height: Output image height. + width: Output image width. + dtype: Data type for the tensors. + device: Device to place tensors on. + generator: Random generator for reproducibility. + + Returns: + Tuple of (mask, masked_image_latents) prepared for the denoising loop. + """ + # Calculate latent dimensions + latent_height = 2 * (int(height) // (self.vae_scale_factor * 2)) + latent_width = 2 * (int(width) // (self.vae_scale_factor * 2)) + + # Resize mask to latent dimensions + mask = torch.nn.functional.interpolate(mask, size=(latent_height, latent_width), mode="nearest") + mask = mask.to(device=device, dtype=dtype) + + # Encode masked image to latents + masked_image = masked_image.to(device=device, dtype=dtype) + if isinstance(generator, list): + masked_image_latents = [ + retrieve_latents(self.vae.encode(masked_image[i : i + 1]), generator=generator[i]) + for i in range(masked_image.shape[0]) + ] + masked_image_latents = torch.cat(masked_image_latents, dim=0) + else: + masked_image_latents = retrieve_latents(self.vae.encode(masked_image), generator=generator) + + # Apply VAE scaling + masked_image_latents = (masked_image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor + + # Expand for batch size + if mask.shape[0] < batch_size: + if not batch_size % mask.shape[0] == 0: + raise ValueError( + "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to" + f" a total batch size of {batch_size}, but {mask.shape[0]} masks were passed. Make sure the number" + " of masks that you pass is divisible by the total requested batch size." + ) + mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1) + if masked_image_latents.shape[0] < batch_size: + if not batch_size % masked_image_latents.shape[0] == 0: + raise ValueError( + "The passed images and the required batch size don't match. Images are supposed to be duplicated" + f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed." + " Make sure the number of images that you pass is divisible by the total requested batch size." + ) + masked_image_latents = masked_image_latents.repeat(batch_size // masked_image_latents.shape[0], 1, 1, 1) + + return mask, masked_image_latents + + def prepare_latents( + self, + image, + timestep, + batch_size, + num_channels_latents, + height, + width, + dtype, + device, + generator, + latents=None, + ): + """Prepare latents for inpainting, returning noise and image_latents for blending. + + Returns: + Tuple of (latents, noise, image_latents) where: + - latents: Noised image latents for denoising + - noise: The noise tensor used for blending + - image_latents: Clean image latents for blending + """ + height = 2 * (int(height) // (self.vae_scale_factor * 2)) + width = 2 * (int(width) // (self.vae_scale_factor * 2)) + + shape = (batch_size, num_channels_latents, height, width) + + if latents is not None: + # Generate noise for blending even if latents are provided + noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + # Encode image for blending + image = image.to(device=device, dtype=dtype) + if isinstance(generator, list): + image_latents = [ + retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) + for i in range(image.shape[0]) + ] + image_latents = torch.cat(image_latents, dim=0) + else: + image_latents = retrieve_latents(self.vae.encode(image), generator=generator) + image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor + if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0: + image_latents = torch.cat([image_latents] * (batch_size // image_latents.shape[0]), dim=0) + return latents.to(device=device, dtype=dtype), noise, image_latents + + # Encode the input image + image = image.to(device=device, dtype=dtype) + if image.shape[1] != num_channels_latents: + if isinstance(generator, list): + image_latents = [ + retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) + for i in range(image.shape[0]) + ] + image_latents = torch.cat(image_latents, dim=0) + else: + image_latents = retrieve_latents(self.vae.encode(image), generator=generator) + + # Apply scaling (inverse of decoding: decode does latents/scaling_factor + shift_factor) + image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor + else: + image_latents = image + + # Handle batch size expansion + if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0: + additional_image_per_prompt = batch_size // image_latents.shape[0] + image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0) + elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0: + raise ValueError( + f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts." + ) + + # Generate noise for both initial noising and later blending + noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + + # Add noise using flow matching scale_noise + latents = self.scheduler.scale_noise(image_latents, timestep, noise) + + return latents, noise, image_latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 0 + + @property + def joint_attention_kwargs(self): + return self._joint_attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def interrupt(self): + return self._interrupt + + def check_inputs( + self, + prompt, + image, + mask_image, + strength, + height, + width, + output_type, + negative_prompt=None, + prompt_embeds=None, + negative_prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + ): + if strength < 0 or strength > 1: + raise ValueError(f"The value of strength should be in [0.0, 1.0] but is {strength}") + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if image is None: + raise ValueError("`image` input cannot be undefined for inpainting.") + + if mask_image is None: + raise ValueError("`mask_image` input cannot be undefined for inpainting.") + + if output_type not in ["latent", "pil", "np", "pt"]: + raise ValueError(f"`output_type` must be one of 'latent', 'pil', 'np', or 'pt', but got {output_type}") + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + image: PipelineImageInput = None, + mask_image: PipelineImageInput = None, + masked_image_latents: Optional[torch.FloatTensor] = None, + strength: float = 1.0, + height: int | None = None, + width: int | None = None, + num_inference_steps: int = 50, + sigmas: list[float] | None = None, + guidance_scale: float = 5.0, + cfg_normalization: bool = False, + cfg_truncation: float = 1.0, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[List[torch.FloatTensor]] = None, + negative_prompt_embeds: Optional[List[torch.FloatTensor]] = None, + output_type: str = "pil", + return_dict: bool = True, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + max_sequence_length: int = 512, + ): + r""" + Function invoked when calling the pipeline for inpainting. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both + numpy array and pytorch tensor, the expected value range is between `[0, 1]`. If it's a tensor or a + list of tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or + a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)`. + mask_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + `Image`, numpy array or tensor representing a mask image for inpainting. White pixels (value 1) in the + mask will be inpainted, black pixels (value 0) will be preserved from the original image. + masked_image_latents (`torch.FloatTensor`, *optional*): + Pre-encoded masked image latents. If provided, the masked image encoding step will be skipped. + strength (`float`, *optional*, defaults to 1.0): + Indicates extent to transform the reference `image`. Must be between 0 and 1. `image` is used as a + starting point and more noise is added the higher the `strength`. The number of denoising steps depends + on the amount of noise initially added. When `strength` is 1, added noise is maximum and the denoising + process runs for the full number of iterations specified in `num_inference_steps`. A value of 1 + essentially ignores `image` in the masked region. + height (`int`, *optional*, defaults to 1024): + The height in pixels of the generated image. If not provided, uses the input image height. + width (`int`, *optional*, defaults to 1024): + The width in pixels of the generated image. If not provided, uses the input image width. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + guidance_scale (`float`, *optional*, defaults to 5.0): + Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). + `guidance_scale` is defined as `w` of equation 2. of [Imagen + Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > + 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, + usually at the expense of lower image quality. + cfg_normalization (`bool`, *optional*, defaults to False): + Whether to apply configuration normalization. + cfg_truncation (`float`, *optional*, defaults to 1.0): + The truncation value for configuration. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + prompt_embeds (`List[torch.FloatTensor]`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`List[torch.FloatTensor]`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.stable_diffusion.ZImagePipelineOutput`] instead of a plain + tuple. + joint_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int`, *optional*, defaults to 512): + Maximum sequence length to use with the `prompt`. + + Examples: + + Returns: + [`~pipelines.z_image.ZImagePipelineOutput`] or `tuple`: [`~pipelines.z_image.ZImagePipelineOutput`] if + `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the + generated images. + """ + # 1. Check inputs + self.check_inputs( + prompt=prompt, + image=image, + mask_image=mask_image, + strength=strength, + height=height, + width=width, + output_type=output_type, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + ) + + # 2. Preprocess image and mask + init_image = self.image_processor.preprocess(image) + init_image = init_image.to(dtype=torch.float32) + + # Get dimensions from the preprocessed image if not specified + if height is None: + height = init_image.shape[-2] + if width is None: + width = init_image.shape[-1] + + vae_scale = self.vae_scale_factor * 2 + if height % vae_scale != 0: + raise ValueError( + f"Height must be divisible by {vae_scale} (got {height}). " + f"Please adjust the height to a multiple of {vae_scale}." + ) + if width % vae_scale != 0: + raise ValueError( + f"Width must be divisible by {vae_scale} (got {width}). " + f"Please adjust the width to a multiple of {vae_scale}." + ) + + # Preprocess mask + mask = self.mask_processor.preprocess(mask_image, height=height, width=width) + + device = self._execution_device + + self._guidance_scale = guidance_scale + self._joint_attention_kwargs = joint_attention_kwargs + self._interrupt = False + self._cfg_normalization = cfg_normalization + self._cfg_truncation = cfg_truncation + + # 3. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = len(prompt_embeds) + + # If prompt_embeds is provided and prompt is None, skip encoding + if prompt_embeds is not None and prompt is None: + if self.do_classifier_free_guidance and negative_prompt_embeds is None: + raise ValueError( + "When `prompt_embeds` is provided without `prompt`, " + "`negative_prompt_embeds` must also be provided for classifier-free guidance." + ) + else: + ( + prompt_embeds, + negative_prompt_embeds, + ) = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + device=device, + max_sequence_length=max_sequence_length, + ) + + # 4. Prepare latent variables + num_channels_latents = self.transformer.in_channels + + # Repeat prompt_embeds for num_images_per_prompt + if num_images_per_prompt > 1: + prompt_embeds = [pe for pe in prompt_embeds for _ in range(num_images_per_prompt)] + if self.do_classifier_free_guidance and negative_prompt_embeds: + negative_prompt_embeds = [npe for npe in negative_prompt_embeds for _ in range(num_images_per_prompt)] + + actual_batch_size = batch_size * num_images_per_prompt + + # Calculate latent dimensions for image_seq_len + latent_height = 2 * (int(height) // (self.vae_scale_factor * 2)) + latent_width = 2 * (int(width) // (self.vae_scale_factor * 2)) + image_seq_len = (latent_height // 2) * (latent_width // 2) + + # 5. Prepare timesteps + mu = calculate_shift( + image_seq_len, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.5), + self.scheduler.config.get("max_shift", 1.15), + ) + self.scheduler.sigma_min = 0.0 + scheduler_kwargs = {"mu": mu} + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + sigmas=sigmas, + **scheduler_kwargs, + ) + + # 6. Adjust timesteps based on strength + timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) + if num_inference_steps < 1: + raise ValueError( + f"After adjusting the num_inference_steps by strength parameter: {strength}, the number of pipeline " + f"steps is {num_inference_steps} which is < 1 and not appropriate for this pipeline." + ) + latent_timestep = timesteps[:1].repeat(actual_batch_size) + + # 7. Prepare latents from image (returns noise and image_latents for blending) + latents, noise, image_latents = self.prepare_latents( + init_image, + latent_timestep, + actual_batch_size, + num_channels_latents, + height, + width, + prompt_embeds[0].dtype, + device, + generator, + latents, + ) + + # 8. Prepare mask and masked image latents + # Create masked image: preserve only unmasked regions (mask=0) + if masked_image_latents is None: + masked_image = init_image * (mask < 0.5) + else: + masked_image = None # Will use provided masked_image_latents + + mask, masked_image_latents = self.prepare_mask_latents( + mask, + masked_image if masked_image is not None else init_image, + actual_batch_size, + height, + width, + prompt_embeds[0].dtype, + device, + generator, + ) + + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + # 9. Denoising loop + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latents.shape[0]) + timestep = (1000 - timestep) / 1000 + # Normalized time for time-aware config (0 at start, 1 at end) + t_norm = timestep[0].item() + + # Handle cfg truncation + current_guidance_scale = self.guidance_scale + if ( + self.do_classifier_free_guidance + and self._cfg_truncation is not None + and float(self._cfg_truncation) <= 1 + ): + if t_norm > self._cfg_truncation: + current_guidance_scale = 0.0 + + # Run CFG only if configured AND scale is non-zero + apply_cfg = self.do_classifier_free_guidance and current_guidance_scale > 0 + + if apply_cfg: + latents_typed = latents.to(self.transformer.dtype) + latent_model_input = latents_typed.repeat(2, 1, 1, 1) + prompt_embeds_model_input = prompt_embeds + negative_prompt_embeds + timestep_model_input = timestep.repeat(2) + else: + latent_model_input = latents.to(self.transformer.dtype) + prompt_embeds_model_input = prompt_embeds + timestep_model_input = timestep + + latent_model_input = latent_model_input.unsqueeze(2) + latent_model_input_list = list(latent_model_input.unbind(dim=0)) + + model_out_list = self.transformer( + latent_model_input_list, + timestep_model_input, + prompt_embeds_model_input, + )[0] + + if apply_cfg: + # Perform CFG + pos_out = model_out_list[:actual_batch_size] + neg_out = model_out_list[actual_batch_size:] + + noise_pred = [] + for j in range(actual_batch_size): + pos = pos_out[j].float() + neg = neg_out[j].float() + + pred = pos + current_guidance_scale * (pos - neg) + + # Renormalization + if self._cfg_normalization and float(self._cfg_normalization) > 0.0: + ori_pos_norm = torch.linalg.vector_norm(pos) + new_pos_norm = torch.linalg.vector_norm(pred) + max_new_norm = ori_pos_norm * float(self._cfg_normalization) + if new_pos_norm > max_new_norm: + pred = pred * (max_new_norm / new_pos_norm) + + noise_pred.append(pred) + + noise_pred = torch.stack(noise_pred, dim=0) + else: + noise_pred = torch.stack([t.float() for t in model_out_list], dim=0) + + noise_pred = noise_pred.squeeze(2) + noise_pred = -noise_pred + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred.to(torch.float32), t, latents, return_dict=False)[0] + assert latents.dtype == torch.float32 + + # Inpainting blend: combine denoised latents with original image latents + init_latents_proper = image_latents + + # Re-scale original latents to current noise level for proper blending + if i < len(timesteps) - 1: + noise_timestep = timesteps[i + 1] + init_latents_proper = self.scheduler.scale_noise( + init_latents_proper, torch.tensor([noise_timestep]), noise + ) + + # Blend: mask=1 for inpaint region (use denoised), mask=0 for preserve region (use original) + latents = (1 - mask) * init_latents_proper + mask * latents + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + mask = callback_outputs.pop("mask", mask) + masked_image_latents = callback_outputs.pop("masked_image_latents", masked_image_latents) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + if output_type == "latent": + image = latents + + else: + latents = latents.to(self.vae.dtype) + latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor + + image = self.vae.decode(latents, return_dict=False)[0] + image = self.image_processor.postprocess(image, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return ZImagePipelineOutput(images=image) diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_omni.py b/src/diffusers/pipelines/z_image/pipeline_z_image_omni.py index 26848bea0a9e..6d04202162f9 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image_omni.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image_omni.py @@ -13,7 +13,7 @@ # limitations under the License. import inspect -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable import PIL import torch @@ -78,10 +78,10 @@ def calculate_shift( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, **kwargs, ): r""" @@ -96,15 +96,15 @@ def retrieve_timesteps( must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: @@ -169,12 +169,12 @@ def __init__( def encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, + prompt: str | list[str], + device: torch.device | None = None, do_classifier_free_guidance: bool = True, - negative_prompt: Optional[Union[str, List[str]]] = None, - prompt_embeds: Optional[List[torch.FloatTensor]] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt: str | list[str] | None = None, + prompt_embeds: list[torch.FloatTensor] | None = None, + negative_prompt_embeds: torch.FloatTensor | None = None, max_sequence_length: int = 512, num_condition_images: int = 0, ): @@ -206,12 +206,12 @@ def encode_prompt( def _encode_prompt( self, - prompt: Union[str, List[str]], - device: Optional[torch.device] = None, - prompt_embeds: Optional[List[torch.FloatTensor]] = None, + prompt: str | list[str], + device: torch.device | None = None, + prompt_embeds: list[torch.FloatTensor] | None = None, max_sequence_length: int = 512, num_condition_images: int = 0, - ) -> List[torch.FloatTensor]: + ) -> list[torch.FloatTensor]: device = device or self._execution_device if prompt_embeds is not None: @@ -292,7 +292,7 @@ def prepare_latents( def prepare_image_latents( self, - images: List[torch.Tensor], + images: list[torch.Tensor], batch_size, device, dtype, @@ -313,7 +313,7 @@ def prepare_image_latents( def prepare_siglip_embeds( self, - images: List[torch.Tensor], + images: list[torch.Tensor], batch_size, device, dtype, @@ -339,7 +339,7 @@ def guidance_scale(self): @property def do_classifier_free_guidance(self): - return self._guidance_scale > 1 + return self._guidance_scale > 0 @property def joint_attention_kwargs(self): @@ -357,39 +357,39 @@ def interrupt(self): @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - image: Optional[Union[List[PIL.Image.Image], PIL.Image.Image]] = None, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, + image: list[PIL.Image.Image, PIL.Image.Image] | None = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, num_inference_steps: int = 50, - sigmas: Optional[List[float]] = None, + sigmas: list[float] | None = None, guidance_scale: float = 5.0, cfg_normalization: bool = False, cfg_truncation: float = 1.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[List[torch.FloatTensor]] = None, - negative_prompt_embeds: Optional[List[torch.FloatTensor]] = None, - output_type: Optional[str] = "pil", + negative_prompt: str | list[str] | None = None, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.FloatTensor | None = None, + prompt_embeds: list[torch.FloatTensor] | None = None, + negative_prompt_embeds: list[torch.FloatTensor] | None = None, + output_type: str | None = "pil", return_dict: bool = True, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], + joint_attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ): r""" Function invoked when calling the pipeline for generation. Args: - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`): + image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `list[torch.Tensor]`, `list[PIL.Image.Image]`, or `list[np.ndarray]`): `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image latents as `image`, but if passing latents directly it is not encoded again. - prompt (`str` or `List[str]`, *optional*): + prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. height (`int`, *optional*, defaults to 1024): @@ -399,7 +399,7 @@ def __call__( num_inference_steps (`int`, *optional*, defaults to 50): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. @@ -413,23 +413,23 @@ def __call__( Whether to apply configuration normalization. cfg_truncation (`float`, *optional*, defaults to 1.0): The truncation value for configuration. - negative_prompt (`str` or `List[str]`, *optional*): + negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image generation. Can be used to tweak the same generation with different prompts. If not provided, a latents tensor will be generated by sampling using the supplied random `generator`. - prompt_embeds (`List[torch.FloatTensor]`, *optional*): + prompt_embeds (`list[torch.FloatTensor]`, *optional*): Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, text embeddings will be generated from `prompt` input argument. - negative_prompt_embeds (`List[torch.FloatTensor]`, *optional*): + negative_prompt_embeds (`list[torch.FloatTensor]`, *optional*): Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. diff --git a/src/diffusers/quantizers/auto.py b/src/diffusers/quantizers/auto.py index 070bcd0b2151..6cd24c459c9d 100644 --- a/src/diffusers/quantizers/auto.py +++ b/src/diffusers/quantizers/auto.py @@ -17,7 +17,6 @@ """ import warnings -from typing import Dict, Optional, Union from .bitsandbytes import BnB4BitDiffusersQuantizer, BnB8BitDiffusersQuantizer from .gguf import GGUFQuantizer @@ -61,7 +60,7 @@ class DiffusersAutoQuantizer: """ @classmethod - def from_dict(cls, quantization_config_dict: Dict): + def from_dict(cls, quantization_config_dict: dict): quant_method = quantization_config_dict.get("quant_method", None) # We need a special care for bnb models to make sure everything is BC .. if quantization_config_dict.get("load_in_8bit", False) or quantization_config_dict.get("load_in_4bit", False): @@ -82,7 +81,7 @@ def from_dict(cls, quantization_config_dict: Dict): return target_cls.from_dict(quantization_config_dict) @classmethod - def from_config(cls, quantization_config: Union[QuantizationConfigMixin, Dict], **kwargs): + def from_config(cls, quantization_config: QuantizationConfigMixin | dict, **kwargs): # Convert it to a QuantizationConfig if the q_config is a dict if isinstance(quantization_config, dict): quantization_config = cls.from_dict(quantization_config) @@ -123,8 +122,8 @@ def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): @classmethod def merge_quantization_configs( cls, - quantization_config: Union[dict, QuantizationConfigMixin], - quantization_config_from_args: Optional[QuantizationConfigMixin], + quantization_config: dict | QuantizationConfigMixin, + quantization_config_from_args: QuantizationConfigMixin | None, ): """ handles situations where both quantization_config from args and quantization_config from model config are diff --git a/src/diffusers/quantizers/base.py b/src/diffusers/quantizers/base.py index 24fc724b4c88..b0988284b648 100644 --- a/src/diffusers/quantizers/base.py +++ b/src/diffusers/quantizers/base.py @@ -18,7 +18,7 @@ """ from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any from ..utils import is_torch_available from .quantization_config import QuantizationConfigMixin @@ -40,9 +40,9 @@ class DiffusersQuantizer(ABC): Attributes quantization_config (`diffusers.quantizers.quantization_config.QuantizationConfigMixin`): The quantization config that defines the quantization parameters of your model that you want to quantize. - modules_to_not_convert (`List[str]`, *optional*): + modules_to_not_convert (`list[str]`, *optional*): The list of module names to not convert when quantizing the model. - required_packages (`List[str]`, *optional*): + required_packages (`list[str]`, *optional*): The list of required pip packages to install prior to using the quantizer requires_calibration (`bool`): Whether the quantization method requires to calibrate the model before using it. @@ -76,14 +76,14 @@ def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": """ return torch_dtype - def update_device_map(self, device_map: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]: + def update_device_map(self, device_map: dict[str, Any] | None) -> dict[str, Any] | None: """ Override this method if you want to pass a override the existing device map with a new one. E.g. for bitsandbytes, since `accelerate` is a hard requirement, if no device_map is passed, the device_map is set to `"auto"`` Args: - device_map (`Union[dict, str]`, *optional*): + device_map (`dict | str`, *optional*): The device_map that is passed through the `from_pretrained` method. """ return device_map @@ -100,17 +100,17 @@ def adjust_target_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": """ return torch_dtype - def update_missing_keys(self, model, missing_keys: List[str], prefix: str) -> List[str]: + def update_missing_keys(self, model, missing_keys: list[str], prefix: str) -> list[str]: """ Override this method if you want to adjust the `missing_keys`. Args: - missing_keys (`List[str]`, *optional*): + missing_keys (`list[str]`, *optional*): The list of missing keys in the checkpoint compared to the state dict of the model """ return missing_keys - def get_special_dtypes_update(self, model, torch_dtype: "torch.dtype") -> Dict[str, "torch.dtype"]: + def get_special_dtypes_update(self, model, torch_dtype: "torch.dtype") -> dict[str, "torch.dtype"]: """ returns dtypes for modules that are not quantized - used for the computation of the device_map in case one passes a str as a device_map. The method will use the `modules_to_not_convert` that is modified in @@ -130,7 +130,7 @@ def get_special_dtypes_update(self, model, torch_dtype: "torch.dtype") -> Dict[s if any(m in name for m in self.modules_to_not_convert) } - def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]: + def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]: """adjust max_memory argument for infer_auto_device_map() if extra memory is needed for quantization""" return max_memory @@ -139,7 +139,7 @@ def check_if_quantized_param( model: "ModelMixin", param_value: "torch.Tensor", param_name: str, - state_dict: Dict[str, Any], + state_dict: dict[str, Any], **kwargs, ) -> bool: """ diff --git a/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py b/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py index 0dfdff019b79..16ff0d83b8c4 100644 --- a/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py +++ b/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py @@ -16,7 +16,7 @@ https://github.com/huggingface/transformers/blob/c409cd81777fb27aadc043ed3d8339dbc020fb3b/src/transformers/quantizers/quantizer_bnb_4bit.py """ -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any from ...utils import get_module_from_name from ..base import DiffusersQuantizer @@ -111,7 +111,7 @@ def check_if_quantized_param( model: "ModelMixin", param_value: "torch.Tensor", param_name: str, - state_dict: Dict[str, Any], + state_dict: dict[str, Any], **kwargs, ) -> bool: import bitsandbytes as bnb @@ -133,8 +133,8 @@ def create_quantized_param( param_value: "torch.Tensor", param_name: str, target_device: "torch.device", - state_dict: Dict[str, Any], - unexpected_keys: Optional[List[str]] = None, + state_dict: dict[str, Any], + unexpected_keys: list[str] | None = None, **kwargs, ): import bitsandbytes as bnb @@ -218,7 +218,7 @@ def check_quantized_param_shape(self, param_name, current_param, loaded_param): else: return True - def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]: + def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]: # need more space for buffers that are created during quantization max_memory = {key: val * 0.90 for key, val in max_memory.items()} return max_memory @@ -255,7 +255,7 @@ def _process_model_before_weight_loading( self, model: "ModelMixin", device_map, - keep_in_fp32_modules: List[str] = [], + keep_in_fp32_modules: list[str] = [], **kwargs, ): from .utils import replace_with_bnb_linear @@ -387,7 +387,7 @@ def validate_environment(self, *args, **kwargs): ) # Copied from diffusers.quantizers.bitsandbytes.bnb_quantizer.BnB4BitDiffusersQuantizer.adjust_max_memory - def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]: + def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]: # need more space for buffers that are created during quantization max_memory = {key: val * 0.90 for key, val in max_memory.items()} return max_memory @@ -432,7 +432,7 @@ def check_if_quantized_param( model: "ModelMixin", param_value: "torch.Tensor", param_name: str, - state_dict: Dict[str, Any], + state_dict: dict[str, Any], **kwargs, ): import bitsandbytes as bnb @@ -455,8 +455,8 @@ def create_quantized_param( param_value: "torch.Tensor", param_name: str, target_device: "torch.device", - state_dict: Dict[str, Any], - unexpected_keys: Optional[List[str]] = None, + state_dict: dict[str, Any], + unexpected_keys: list[str] | None = None, **kwargs, ): import bitsandbytes as bnb @@ -513,7 +513,7 @@ def _process_model_before_weight_loading( self, model: "ModelMixin", device_map, - keep_in_fp32_modules: List[str] = [], + keep_in_fp32_modules: list[str] = [], **kwargs, ): from .utils import replace_with_bnb_linear diff --git a/src/diffusers/quantizers/bitsandbytes/utils.py b/src/diffusers/quantizers/bitsandbytes/utils.py index 429aabb8fae6..aea4e7dda57f 100644 --- a/src/diffusers/quantizers/bitsandbytes/utils.py +++ b/src/diffusers/quantizers/bitsandbytes/utils.py @@ -18,7 +18,6 @@ import inspect from inspect import signature -from typing import Union from ...utils import is_accelerate_available, is_bitsandbytes_available, is_torch_available, logging from ..quantization_config import QuantizationMethod @@ -128,10 +127,10 @@ def replace_with_bnb_linear(model, modules_to_not_convert=None, current_key_name Parameters: model (`torch.nn.Module`): Input model or `torch.nn.Module` as the function is run recursively. - modules_to_not_convert (`List[`str`]`, *optional*, defaults to `[]`): + modules_to_not_convert (`list[`str`]`, *optional*, defaults to `[]`): Names of the modules to not convert in `Linear8bitLt`. In practice we keep the `modules_to_not_convert` in full precision for numerical stability reasons. - current_key_name (`List[`str`]`, *optional*): + current_key_name (`list[`str`]`, *optional*): An array to track the current key of the recursion. This is used to check whether the current key (part of it) is not in the list of modules to not convert (for instances modules that are offloaded to `cpu` or `disk`). @@ -305,7 +304,7 @@ def dequantize_and_replace( return model -def _check_bnb_status(module) -> Union[bool, bool]: +def _check_bnb_status(module) -> bool | bool: is_loaded_in_4bit_bnb = ( hasattr(module, "is_loaded_in_4bit") and module.is_loaded_in_4bit diff --git a/src/diffusers/quantizers/gguf/gguf_quantizer.py b/src/diffusers/quantizers/gguf/gguf_quantizer.py index aa5ebf5711a3..15f39dd9605e 100644 --- a/src/diffusers/quantizers/gguf/gguf_quantizer.py +++ b/src/diffusers/quantizers/gguf/gguf_quantizer.py @@ -1,4 +1,6 @@ -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from __future__ import annotations + +from typing import TYPE_CHECKING, Any from ..base import DiffusersQuantizer @@ -41,7 +43,7 @@ def __init__(self, quantization_config, **kwargs): self.compute_dtype = quantization_config.compute_dtype self.pre_quantized = quantization_config.pre_quantized - self.modules_to_not_convert = quantization_config.modules_to_not_convert + self.modules_to_not_convert = quantization_config.modules_to_not_convert or [] if not isinstance(self.modules_to_not_convert, list): self.modules_to_not_convert = [self.modules_to_not_convert] @@ -57,7 +59,7 @@ def validate_environment(self, *args, **kwargs): ) # Copied from diffusers.quantizers.bitsandbytes.bnb_quantizer.BnB4BitDiffusersQuantizer.adjust_max_memory - def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]: + def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]: # need more space for buffers that are created during quantization max_memory = {key: val * 0.90 for key, val in max_memory.items()} return max_memory @@ -90,9 +92,9 @@ def check_quantized_param_shape(self, param_name, current_param, loaded_param): def check_if_quantized_param( self, model: "ModelMixin", - param_value: Union["GGUFParameter", "torch.Tensor"], + param_value: "GGUFParameter" | "torch.Tensor", param_name: str, - state_dict: Dict[str, Any], + state_dict: dict[str, Any], **kwargs, ) -> bool: if isinstance(param_value, GGUFParameter): @@ -103,11 +105,11 @@ def check_if_quantized_param( def create_quantized_param( self, model: "ModelMixin", - param_value: Union["GGUFParameter", "torch.Tensor"], + param_value: "GGUFParameter" | "torch.Tensor", param_name: str, target_device: "torch.device", - state_dict: Optional[Dict[str, Any]] = None, - unexpected_keys: Optional[List[str]] = None, + state_dict: dict[str, Any] | None = None, + unexpected_keys: list[str] | None = None, **kwargs, ): module, tensor_name = get_module_from_name(model, param_name) @@ -123,7 +125,7 @@ def _process_model_before_weight_loading( self, model: "ModelMixin", device_map, - keep_in_fp32_modules: List[str] = [], + keep_in_fp32_modules: list[str] = [], **kwargs, ): state_dict = kwargs.get("state_dict", None) diff --git a/src/diffusers/quantizers/gguf/utils.py b/src/diffusers/quantizers/gguf/utils.py index 2fba9986e825..e0ad0e1cce42 100644 --- a/src/diffusers/quantizers/gguf/utils.py +++ b/src/diffusers/quantizers/gguf/utils.py @@ -79,7 +79,8 @@ def _fused_mul_mat_gguf(x: torch.Tensor, qweight: torch.Tensor, qweight_type: int) -> torch.Tensor: # there is no need to call any kernel for fp16/bf16 if qweight_type in UNQUANTIZED_TYPES: - return x @ qweight.T + weight = dequantize_gguf_tensor(qweight) + return x @ weight.T # TODO(Isotr0py): GGUF's MMQ and MMVQ implementation are designed for # contiguous batching and inefficient with diffusers' batching, @@ -515,6 +516,9 @@ def dequantize_gguf_tensor(tensor): block_size, type_size = GGML_QUANT_SIZES[quant_type] + # Conver to plain tensor to avoid unnecessary __torch_function__ overhead. + tensor = tensor.as_tensor() + tensor = tensor.view(torch.uint8) shape = _quant_shape_from_byte_shape(tensor.shape, type_size, block_size) @@ -524,7 +528,7 @@ def dequantize_gguf_tensor(tensor): dequant = dequant_fn(blocks, block_size, type_size) dequant = dequant.reshape(shape) - return dequant.as_tensor() + return dequant class GGUFParameter(torch.nn.Parameter): diff --git a/src/diffusers/quantizers/modelopt/modelopt_quantizer.py b/src/diffusers/quantizers/modelopt/modelopt_quantizer.py index 7312036f52d0..5c9e198c82b2 100644 --- a/src/diffusers/quantizers/modelopt/modelopt_quantizer.py +++ b/src/diffusers/quantizers/modelopt/modelopt_quantizer.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Dict, List, Union +from typing import TYPE_CHECKING, Any from ...utils import ( get_module_from_name, @@ -61,7 +61,7 @@ def check_if_quantized_param( model: "ModelMixin", param_value: "torch.Tensor", param_name: str, - state_dict: Dict[str, Any], + state_dict: dict[str, Any], **kwargs, ): # ModelOpt imports diffusers internally. This is here to prevent circular imports @@ -101,7 +101,7 @@ def create_quantized_param( mtq.compress(module) module.weight.requires_grad = False - def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]: + def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]: max_memory = {key: val * 0.90 for key, val in max_memory.items()} return max_memory @@ -116,7 +116,7 @@ def update_torch_dtype(self, torch_dtype: "torch.dtype" = None) -> "torch.dtype" torch_dtype = torch.float32 return torch_dtype - def get_conv_param_names(self, model: "ModelMixin") -> List[str]: + def get_conv_param_names(self, model: "ModelMixin") -> list[str]: """ Get parameter names for all convolutional layers in a HuggingFace ModelMixin. Includes Conv1d/2d/3d and ConvTranspose1d/2d/3d. @@ -142,7 +142,7 @@ def _process_model_before_weight_loading( self, model: "ModelMixin", device_map, - keep_in_fp32_modules: List[str] = [], + keep_in_fp32_modules: list[str] = [], **kwargs, ): # ModelOpt imports diffusers internally. This is here to prevent circular imports diff --git a/src/diffusers/quantizers/pipe_quant_config.py b/src/diffusers/quantizers/pipe_quant_config.py index f75a337341a9..90be7d796438 100644 --- a/src/diffusers/quantizers/pipe_quant_config.py +++ b/src/diffusers/quantizers/pipe_quant_config.py @@ -12,8 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import inspect -from typing import Dict, List, Optional, Union from ..utils import is_transformers_available, logging from .quantization_config import QuantizationConfigMixin as DiffQuantConfigMixin @@ -47,9 +48,9 @@ class PipelineQuantizationConfig: def __init__( self, quant_backend: str = None, - quant_kwargs: Dict[str, Union[str, float, int, dict]] = None, - components_to_quantize: Optional[Union[List[str], str]] = None, - quant_mapping: Dict[str, Union[DiffQuantConfigMixin, "TransformersQuantConfigMixin"]] = None, + quant_kwargs: dict[str, str | float | int | dict] = None, + components_to_quantize: list[str] | str | None = None, + quant_mapping: dict[str, DiffQuantConfigMixin | "TransformersQuantConfigMixin"] = None, ): self.quant_backend = quant_backend # Initialize kwargs to be {} to set to the defaults. diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index 5dd8f56717df..138ec7b7e989 100644 --- a/src/diffusers/quantizers/quantization_config.py +++ b/src/diffusers/quantizers/quantization_config.py @@ -20,6 +20,8 @@ https://github.com/huggingface/transformers/blob/52cb4034ada381fe1ffe8d428a1076e5411a8026/src/transformers/utils/quantization_config.py """ +from __future__ import annotations + import copy import dataclasses import importlib.metadata @@ -30,7 +32,7 @@ from dataclasses import dataclass, is_dataclass from enum import Enum from functools import partial -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable from packaging import version @@ -76,12 +78,12 @@ def from_dict(cls, config_dict, return_unused_kwargs=False, **kwargs): Instantiates a [`QuantizationConfigMixin`] from a Python dictionary of parameters. Args: - config_dict (`Dict[str, Any]`): + config_dict (`dict[str, Any]`): Dictionary that will be used to instantiate the configuration object. return_unused_kwargs (`bool`, *optional*, defaults to `False`): Whether or not to return a list of unused keyword arguments. Used for `from_pretrained` method in `PreTrainedModel`. - kwargs (`Dict[str, Any]`): + kwargs (`dict[str, Any]`): Additional parameters from which to initialize the configuration object. Returns: @@ -103,7 +105,7 @@ def from_dict(cls, config_dict, return_unused_kwargs=False, **kwargs): else: return config - def to_json_file(self, json_file_path: Union[str, os.PathLike]): + def to_json_file(self, json_file_path: str | os.PathLike): """ Save this instance to a JSON file. @@ -120,10 +122,10 @@ def to_json_file(self, json_file_path: Union[str, os.PathLike]): writer.write(json_string) - def to_dict(self) -> Dict[str, Any]: + def to_dict(self) -> dict[str, Any]: """ Serializes this instance to a Python dictionary. Returns: - `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. + `dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. """ return copy.deepcopy(self.__dict__) @@ -159,11 +161,11 @@ def update(self, **kwargs): returning all the unused kwargs. Args: - kwargs (`Dict[str, Any]`): + kwargs (`dict[str, Any]`): Dictionary of attributes to tentatively update this class. Returns: - `Dict[str, Any]`: Dictionary containing all the key-value pairs that were not used to update the instance. + `dict[str, Any]`: Dictionary containing all the key-value pairs that were not used to update the instance. """ to_remove = [] for key, value in kwargs.items(): @@ -202,7 +204,7 @@ class BitsAndBytesConfig(QuantizationConfigMixin): These outliers are often in the interval [-60, -6] or [6, 60]. Int8 quantization works well for values of magnitude ~5, but beyond that, there is a significant performance penalty. A good default threshold is 6, but a lower threshold might be needed for more unstable models (small models, fine-tuning). - llm_int8_skip_modules (`List[str]`, *optional*): + llm_int8_skip_modules (`list[str]`, *optional*): An explicit list of the modules that we do not want to convert in 8-bit. This is useful for models such as Jukebox that has several heads in different places and not necessarily at the last position. For example for `CausalLM` models, the last `lm_head` is typically kept in its original `dtype`. @@ -225,7 +227,7 @@ class BitsAndBytesConfig(QuantizationConfigMixin): quantized again. bnb_4bit_quant_storage (`torch.dtype` or str, *optional*, defaults to `torch.uint8`): This sets the storage type to pack the quanitzed 4-bit prarams. - kwargs (`Dict[str, Any]`, *optional*): + kwargs (`dict[str, Any]`, *optional*): Additional parameters from which to initialize the configuration object. """ @@ -376,10 +378,10 @@ def quantization_method(self): else: return None - def to_dict(self) -> Dict[str, Any]: + def to_dict(self) -> dict[str, Any]: """ Serializes this instance to a Python dictionary. Returns: - `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. + `dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. """ output = copy.deepcopy(self.__dict__) output["bnb_4bit_compute_dtype"] = str(output["bnb_4bit_compute_dtype"]).split(".")[1] @@ -393,13 +395,13 @@ def __repr__(self): config_dict = self.to_dict() return f"{self.__class__.__name__} {json.dumps(config_dict, indent=2, sort_keys=True)}\n" - def to_diff_dict(self) -> Dict[str, Any]: + def to_diff_dict(self) -> dict[str, Any]: """ Removes all attributes from config which correspond to the default config attributes for better readability and serializes to a Python dictionary. Returns: - `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance, + `dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance, """ config_dict = self.to_dict() @@ -427,7 +429,7 @@ class GGUFQuantizationConfig(QuantizationConfigMixin): """ - def __init__(self, compute_dtype: Optional["torch.dtype"] = None): + def __init__(self, compute_dtype: "torch.dtype" | None = None): self.quant_method = QuantizationMethod.GGUF self.compute_dtype = compute_dtype self.pre_quantized = True @@ -444,7 +446,7 @@ class TorchAoConfig(QuantizationConfigMixin): """This is a config class for torchao quantization/sparsity techniques. Args: - quant_type (Union[`str`, AOBaseConfig]): + quant_type (`str` | AOBaseConfig): The type of quantization we want to use, currently supporting: - **Integer quantization:** - Full function names: `int4_weight_only`, `int8_dynamic_activation_int4_weight`, @@ -457,7 +459,7 @@ class TorchAoConfig(QuantizationConfigMixin): - Shorthands: `float8wo`, `float8wo_e5m2`, `float8wo_e4m3`, `float8dq`, `float8dq_e4m3`, `float8_e4m3_tensor`, `float8_e4m3_row`, - - **Floating point X-bit quantization:** + - **Floating point X-bit quantization:** (in torchao <= 0.14.1, not supported in torchao >= 0.15.0) - Full function names: `fpx_weight_only` - Shorthands: `fpX_eAwB`, where `X` is the number of bits (between `1` to `7`), `A` is the number of exponent bits and `B` is the number of mantissa bits. The constraint of `X == A + B + 1` must @@ -467,10 +469,10 @@ class TorchAoConfig(QuantizationConfigMixin): - Full function names: `uintx_weight_only` - Shorthands: `uint1wo`, `uint2wo`, `uint3wo`, `uint4wo`, `uint5wo`, `uint6wo`, `uint7wo` - An AOBaseConfig instance: for more advanced configuration options. - modules_to_not_convert (`List[str]`, *optional*, default to `None`): + modules_to_not_convert (`list[str]`, *optional*, default to `None`): The list of modules to not quantize, useful for quantizing models that explicitly require to have some modules left in their original precision. - kwargs (`Dict[str, Any]`, *optional*): + kwargs (`dict[str, Any]`, *optional*): The keyword arguments for the chosen type of quantization, for example, int4_weight_only quantization supports two keyword arguments `group_size` and `inner_k_tiles` currently. More API examples and documentation of arguments can be found in @@ -498,8 +500,8 @@ class TorchAoConfig(QuantizationConfigMixin): def __init__( self, - quant_type: Union[str, "AOBaseConfig"], # noqa: F821 - modules_to_not_convert: Optional[List[str]] = None, + quant_type: str | "AOBaseConfig", # noqa: F821 + modules_to_not_convert: list[str] | None = None, **kwargs, ) -> None: self.quant_method = QuantizationMethod.TORCHAO @@ -531,12 +533,18 @@ def post_init(self): TORCHAO_QUANT_TYPE_METHODS = self._get_torchao_quant_type_to_method() if self.quant_type not in TORCHAO_QUANT_TYPE_METHODS.keys(): - is_floating_quant_type = self.quant_type.startswith("float") or self.quant_type.startswith("fp") - if is_floating_quant_type and not self._is_xpu_or_cuda_capability_atleast_8_9(): + is_floatx_quant_type = self.quant_type.startswith("fp") + is_float_quant_type = self.quant_type.startswith("float") or is_floatx_quant_type + if is_float_quant_type and not self._is_xpu_or_cuda_capability_atleast_8_9(): raise ValueError( f"Requested quantization type: {self.quant_type} is not supported on GPUs with CUDA capability <= 8.9. You " f"can check the CUDA capability of your GPU using `torch.cuda.get_device_capability()`." ) + elif is_floatx_quant_type and not is_torchao_version("<=", "0.14.1"): + raise ValueError( + f"Requested quantization type: {self.quant_type} is only supported in torchao <= 0.14.1. " + f"Please downgrade to torchao <= 0.14.1 to use this quantization type." + ) raise ValueError( f"Requested quantization type: {self.quant_type} is not supported or is an incorrect `quant_type` name. If you think the " @@ -617,12 +625,11 @@ def _get_torchao_quant_type_to_method(cls): """ if is_torchao_available(): - # TODO(aryan): Support autoquant and sparsify + # TODO(aryan): Support sparsify from torchao.quantization import ( float8_dynamic_activation_float8_weight, float8_static_activation_float8_weight, float8_weight_only, - fpx_weight_only, int4_weight_only, int8_dynamic_activation_int4_weight, int8_dynamic_activation_int8_weight, @@ -630,6 +637,8 @@ def _get_torchao_quant_type_to_method(cls): uintx_weight_only, ) + if is_torchao_version("<=", "0.14.1"): + from torchao.quantization import fpx_weight_only # TODO(aryan): Add a note on how to use PerAxis and PerGroup observers from torchao.quantization.observer import PerRow, PerTensor @@ -650,18 +659,21 @@ def generate_float8dq_types(dtype: torch.dtype): return types def generate_fpx_quantization_types(bits: int): - types = {} + if is_torchao_version("<=", "0.14.1"): + types = {} - for ebits in range(1, bits): - mbits = bits - ebits - 1 - types[f"fp{bits}_e{ebits}m{mbits}"] = partial(fpx_weight_only, ebits=ebits, mbits=mbits) + for ebits in range(1, bits): + mbits = bits - ebits - 1 + types[f"fp{bits}_e{ebits}m{mbits}"] = partial(fpx_weight_only, ebits=ebits, mbits=mbits) - non_sign_bits = bits - 1 - default_ebits = (non_sign_bits + 1) // 2 - default_mbits = non_sign_bits - default_ebits - types[f"fp{bits}"] = partial(fpx_weight_only, ebits=default_ebits, mbits=default_mbits) + non_sign_bits = bits - 1 + default_ebits = (non_sign_bits + 1) // 2 + default_mbits = non_sign_bits - default_ebits + types[f"fp{bits}"] = partial(fpx_weight_only, ebits=default_ebits, mbits=default_mbits) - return types + return types + else: + raise ValueError("Floating point X-bit quantization is not supported in torchao >= 0.15.0") INT4_QUANTIZATION_TYPES = { # int4 weight + bfloat16/float16 activation @@ -710,15 +722,15 @@ def generate_fpx_quantization_types(bits: int): **generate_float8dq_types(torch.float8_e4m3fn), # float8 weight + float8 activation (static) "float8_static_activation_float8_weight": float8_static_activation_float8_weight, - # For fpx, only x <= 8 is supported by default. Other dtypes can be explored by users directly - # fpx weight + bfloat16/float16 activation - **generate_fpx_quantization_types(3), - **generate_fpx_quantization_types(4), - **generate_fpx_quantization_types(5), - **generate_fpx_quantization_types(6), - **generate_fpx_quantization_types(7), } + if is_torchao_version("<=", "0.14.1"): + FLOATX_QUANTIZATION_TYPES.update(generate_fpx_quantization_types(3)) + FLOATX_QUANTIZATION_TYPES.update(generate_fpx_quantization_types(4)) + FLOATX_QUANTIZATION_TYPES.update(generate_fpx_quantization_types(5)) + FLOATX_QUANTIZATION_TYPES.update(generate_fpx_quantization_types(6)) + FLOATX_QUANTIZATION_TYPES.update(generate_fpx_quantization_types(7)) + UINTX_QUANTIZATION_DTYPES = { "uintx_weight_only": uintx_weight_only, "uint1wo": partial(uintx_weight_only, dtype=torch.uint1), @@ -829,7 +841,7 @@ class QuantoConfig(QuantizationConfigMixin): def __init__( self, weights_dtype: str = "int8", - modules_to_not_convert: Optional[List[str]] = None, + modules_to_not_convert: list[str] | None = None, **kwargs, ): self.quant_method = QuantizationMethod.QUANTO @@ -861,7 +873,7 @@ class NVIDIAModelOptConfig(QuantizationConfigMixin): - INT4 - NF4 - NVFP4 - modules_to_not_convert (`List[str]`, *optional*, default to `None`): + modules_to_not_convert (`list[str]`, *optional*, default to `None`): The list of modules to not quantize, useful for quantizing models that explicitly require to have some weight_only (`bool`, *optional*, default to `False`): If set to `True`, the quantization will be applied only to the weights of the model. @@ -881,7 +893,7 @@ class NVIDIAModelOptConfig(QuantizationConfigMixin): The modelopt config, useful for passing custom configs to modelopt. disable_conv_quantization (`bool`, *optional*, default to `False`): If set to `True`, the quantization will be disabled for convolutional layers. - kwargs (`Dict[str, Any]`, *optional*): + kwargs (`dict[str, Any]`, *optional*): Additional parameters which are to be used for calibration. """ @@ -900,15 +912,15 @@ class NVIDIAModelOptConfig(QuantizationConfigMixin): def __init__( self, quant_type: str, - modules_to_not_convert: Optional[List[str]] = None, + modules_to_not_convert: list[str] | None = None, weight_only: bool = True, - channel_quantize: Optional[int] = None, - block_quantize: Optional[int] = None, - scale_channel_quantize: Optional[int] = None, - scale_block_quantize: Optional[int] = None, + channel_quantize: int | None = None, + block_quantize: int | None = None, + scale_channel_quantize: int | None = None, + scale_block_quantize: int | None = None, algorithm: str = "max", - forward_loop: Optional[Callable] = None, - modelopt_config: Optional[dict] = None, + forward_loop: Callable | None = None, + modelopt_config: dict | None = None, disable_conv_quantization: bool = False, **kwargs, ) -> None: @@ -970,7 +982,7 @@ def _normalize_quant_type(self, quant_type: str) -> str: act_type = None self.quant_type = w_type + ("_" + act_type if act_type is not None else "") - def get_config_from_quant_type(self) -> Dict[str, Any]: + def get_config_from_quant_type(self) -> dict[str, Any]: """ Get the config from the quantization type. """ diff --git a/src/diffusers/quantizers/quanto/quanto_quantizer.py b/src/diffusers/quantizers/quanto/quanto_quantizer.py index c5f71f816fc3..a036dabfe6f4 100644 --- a/src/diffusers/quantizers/quanto/quanto_quantizer.py +++ b/src/diffusers/quantizers/quanto/quanto_quantizer.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Dict, List, Union +from typing import TYPE_CHECKING, Any from diffusers.utils.import_utils import is_optimum_quanto_version @@ -68,7 +68,7 @@ def check_if_quantized_param( model: "ModelMixin", param_value: "torch.Tensor", param_name: str, - state_dict: Dict[str, Any], + state_dict: dict[str, Any], **kwargs, ): # Quanto imports diffusers internally. This is here to prevent circular imports @@ -105,7 +105,7 @@ def create_quantized_param( module.freeze() module.weight.requires_grad = False - def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]: + def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]: max_memory = {key: val * 0.90 for key, val in max_memory.items()} return max_memory @@ -127,7 +127,7 @@ def update_torch_dtype(self, torch_dtype: "torch.dtype" = None) -> "torch.dtype" torch_dtype = torch.float32 return torch_dtype - def update_missing_keys(self, model, missing_keys: List[str], prefix: str) -> List[str]: + def update_missing_keys(self, model, missing_keys: list[str], prefix: str) -> list[str]: # Quanto imports diffusers internally. This is here to prevent circular imports from optimum.quanto import QModuleMixin @@ -147,7 +147,7 @@ def _process_model_before_weight_loading( self, model: "ModelMixin", device_map, - keep_in_fp32_modules: List[str] = [], + keep_in_fp32_modules: list[str] = [], **kwargs, ): self.modules_to_not_convert = self.quantization_config.modules_to_not_convert diff --git a/src/diffusers/quantizers/torchao/torchao_quantizer.py b/src/diffusers/quantizers/torchao/torchao_quantizer.py index 2334c7af8630..1679ed26a104 100644 --- a/src/diffusers/quantizers/torchao/torchao_quantizer.py +++ b/src/diffusers/quantizers/torchao/torchao_quantizer.py @@ -21,7 +21,7 @@ import re import types from fnmatch import fnmatch -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any from packaging import version @@ -36,6 +36,9 @@ from ..base import DiffusersQuantizer +logger = logging.get_logger(__name__) + + if TYPE_CHECKING: from ...models.modeling_utils import ModelMixin @@ -83,11 +86,19 @@ def _update_torch_safe_globals(): ] try: from torchao.dtypes import NF4Tensor - from torchao.dtypes.floatx.float8_layout import Float8AQTTensorImpl - from torchao.dtypes.uintx.uint4_layout import UInt4Tensor from torchao.dtypes.uintx.uintx_layout import UintxAQTTensorImpl, UintxTensor - safe_globals.extend([UintxTensor, UInt4Tensor, UintxAQTTensorImpl, Float8AQTTensorImpl, NF4Tensor]) + safe_globals.extend([UintxTensor, UintxAQTTensorImpl, NF4Tensor]) + + # note: is_torchao_version(">=", "0.16.0") does not work correctly + # with torchao nightly, so using a ">" check which does work correctly + if is_torchao_version(">", "0.15.0"): + pass + else: + from torchao.dtypes.floatx.float8_layout import Float8AQTTensorImpl + from torchao.dtypes.uintx.uint4_layout import UInt4Tensor + + safe_globals.extend([UInt4Tensor, Float8AQTTensorImpl]) except (ImportError, ModuleNotFoundError) as e: logger.warning( @@ -108,7 +119,7 @@ def _update_torch_safe_globals(): _update_torch_safe_globals() -def fuzzy_match_size(config_name: str) -> Optional[str]: +def fuzzy_match_size(config_name: str) -> str | None: """ Extract the size digit from strings like "4weight", "8weight". Returns the digit as an integer if found, otherwise None. @@ -123,9 +134,6 @@ def fuzzy_match_size(config_name: str) -> Optional[str]: return None -logger = logging.get_logger(__name__) - - def _quantization_type(weight): from torchao.dtypes import AffineQuantizedTensor from torchao.quantization.linear_activation_quantized_tensor import LinearActivationQuantizedTensor @@ -263,7 +271,7 @@ def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype": f"dtype you are using should be supported, please open an issue at https://github.com/huggingface/diffusers/issues." ) - def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]: + def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]: max_memory = {key: val * 0.9 for key, val in max_memory.items()} return max_memory @@ -272,7 +280,7 @@ def check_if_quantized_param( model: "ModelMixin", param_value: "torch.Tensor", param_name: str, - state_dict: Dict[str, Any], + state_dict: dict[str, Any], **kwargs, ) -> bool: param_device = kwargs.pop("param_device", None) @@ -293,8 +301,8 @@ def create_quantized_param( param_value: "torch.Tensor", param_name: str, target_device: "torch.device", - state_dict: Dict[str, Any], - unexpected_keys: List[str], + state_dict: dict[str, Any], + unexpected_keys: list[str], **kwargs, ): r""" @@ -336,7 +344,6 @@ def get_cuda_warm_up_factor(self): from torchao.core.config import AOBaseConfig quant_type = self.quantization_config.quant_type - # For autoquant case, it will be treated in the string implementation below in map_to_target_dtype if isinstance(quant_type, AOBaseConfig): # Extract size digit using fuzzy match on the class name config_name = quant_type.__class__.__name__ @@ -358,7 +365,7 @@ def _process_model_before_weight_loading( self, model: "ModelMixin", device_map, - keep_in_fp32_modules: List[str] = [], + keep_in_fp32_modules: list[str] = [], **kwargs, ): self.modules_to_not_convert = self.quantization_config.modules_to_not_convert diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py index 29052c1ba0cb..c7101d1b0401 100644 --- a/src/diffusers/schedulers/__init__.py +++ b/src/diffusers/schedulers/__init__.py @@ -61,11 +61,14 @@ _import_structure["scheduling_flow_match_euler_discrete"] = ["FlowMatchEulerDiscreteScheduler"] _import_structure["scheduling_flow_match_heun_discrete"] = ["FlowMatchHeunDiscreteScheduler"] _import_structure["scheduling_flow_match_lcm"] = ["FlowMatchLCMScheduler"] + _import_structure["scheduling_helios"] = ["HeliosScheduler"] + _import_structure["scheduling_helios_dmd"] = ["HeliosDMDScheduler"] _import_structure["scheduling_heun_discrete"] = ["HeunDiscreteScheduler"] _import_structure["scheduling_ipndm"] = ["IPNDMScheduler"] _import_structure["scheduling_k_dpm_2_ancestral_discrete"] = ["KDPM2AncestralDiscreteScheduler"] _import_structure["scheduling_k_dpm_2_discrete"] = ["KDPM2DiscreteScheduler"] _import_structure["scheduling_lcm"] = ["LCMScheduler"] + _import_structure["scheduling_ltx_euler_ancestral_rf"] = ["LTXEulerAncestralRFScheduler"] _import_structure["scheduling_pndm"] = ["PNDMScheduler"] _import_structure["scheduling_repaint"] = ["RePaintScheduler"] _import_structure["scheduling_sasolver"] = ["SASolverScheduler"] @@ -163,11 +166,14 @@ from .scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler from .scheduling_flow_match_heun_discrete import FlowMatchHeunDiscreteScheduler from .scheduling_flow_match_lcm import FlowMatchLCMScheduler + from .scheduling_helios import HeliosScheduler + from .scheduling_helios_dmd import HeliosDMDScheduler from .scheduling_heun_discrete import HeunDiscreteScheduler from .scheduling_ipndm import IPNDMScheduler from .scheduling_k_dpm_2_ancestral_discrete import KDPM2AncestralDiscreteScheduler from .scheduling_k_dpm_2_discrete import KDPM2DiscreteScheduler from .scheduling_lcm import LCMScheduler + from .scheduling_ltx_euler_ancestral_rf import LTXEulerAncestralRFScheduler from .scheduling_pndm import PNDMScheduler from .scheduling_repaint import RePaintScheduler from .scheduling_sasolver import SASolverScheduler diff --git a/src/diffusers/schedulers/deprecated/scheduling_karras_ve.py b/src/diffusers/schedulers/deprecated/scheduling_karras_ve.py index 9206ee80a6b6..b6b905b0f088 100644 --- a/src/diffusers/schedulers/deprecated/scheduling_karras_ve.py +++ b/src/diffusers/schedulers/deprecated/scheduling_karras_ve.py @@ -14,7 +14,6 @@ from dataclasses import dataclass -from typing import Optional, Tuple, Union import numpy as np import torch @@ -43,7 +42,7 @@ class KarrasVeOutput(BaseOutput): prev_sample: torch.Tensor derivative: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None + pred_original_sample: torch.Tensor | None = None class KarrasVeScheduler(SchedulerMixin, ConfigMixin): @@ -93,7 +92,7 @@ def __init__( self.timesteps: np.IntTensor = None self.schedule: torch.Tensor = None # sigma(t_i) - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: int = None) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. @@ -110,7 +109,7 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None """ return sample - def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None): + def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -133,8 +132,8 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic self.schedule = torch.tensor(schedule, dtype=torch.float32, device=device) def add_noise_to_input( - self, sample: torch.Tensor, sigma: float, generator: Optional[torch.Generator] = None - ) -> Tuple[torch.Tensor, float]: + self, sample: torch.Tensor, sigma: float, generator: torch.Generator | None = None + ) -> tuple[torch.Tensor, float]: """ Explicit Langevin-like "churn" step of adding noise to the sample according to a `gamma_i ≥ 0` to reach a higher noise level `sigma_hat = sigma_i + gamma_i*sigma_i`. @@ -165,7 +164,7 @@ def step( sigma_prev: float, sample_hat: torch.Tensor, return_dict: bool = True, - ) -> Union[KarrasVeOutput, Tuple]: + ) -> KarrasVeOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). @@ -206,7 +205,7 @@ def step_correct( sample_prev: torch.Tensor, derivative: torch.Tensor, return_dict: bool = True, - ) -> Union[KarrasVeOutput, Tuple]: + ) -> KarrasVeOutput | tuple: """ Corrects the predicted sample based on the `model_output` of the network. diff --git a/src/diffusers/schedulers/deprecated/scheduling_sde_vp.py b/src/diffusers/schedulers/deprecated/scheduling_sde_vp.py index 5088bdb49761..99fa408e70c0 100644 --- a/src/diffusers/schedulers/deprecated/scheduling_sde_vp.py +++ b/src/diffusers/schedulers/deprecated/scheduling_sde_vp.py @@ -15,7 +15,6 @@ # DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch import math -from typing import Union import torch @@ -48,7 +47,7 @@ def __init__(self, num_train_timesteps=2000, beta_min=0.1, beta_max=20, sampling self.discrete_sigmas = None self.timesteps = None - def set_timesteps(self, num_inference_steps, device: Union[str, torch.device] = None): + def set_timesteps(self, num_inference_steps, device: str | torch.device = None): """ Sets the continuous timesteps used for the diffusion chain (to be run before inference). diff --git a/src/diffusers/schedulers/scheduling_amused.py b/src/diffusers/schedulers/scheduling_amused.py index a0b8fbc862b0..ef49d8db91b1 100644 --- a/src/diffusers/schedulers/scheduling_amused.py +++ b/src/diffusers/schedulers/scheduling_amused.py @@ -1,6 +1,6 @@ import math from dataclasses import dataclass -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import torch @@ -9,7 +9,7 @@ from .scheduling_utils import SchedulerMixin -def gumbel_noise(t: torch.Tensor, generator: Optional[torch.Generator] = None) -> torch.Tensor: +def gumbel_noise(t: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor: """ Generate Gumbel noise for sampling. @@ -32,7 +32,7 @@ def mask_by_random_topk( mask_len: torch.Tensor, probs: torch.Tensor, temperature: float = 1.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, ) -> torch.Tensor: """ Mask tokens by selecting the top-k lowest confidence scores with temperature-based randomness. @@ -73,7 +73,7 @@ class AmusedSchedulerOutput(BaseOutput): """ prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None + pred_original_sample: torch.Generator | None = None class AmusedScheduler(SchedulerMixin, ConfigMixin): @@ -96,8 +96,8 @@ class AmusedScheduler(SchedulerMixin, ConfigMixin): order = 1 - temperatures: Optional[torch.Tensor] - timesteps: Optional[torch.Tensor] + temperatures: torch.Generator | None + timesteps: torch.Generator | None @register_to_config def __init__( @@ -111,23 +111,9 @@ def __init__( def set_timesteps( self, num_inference_steps: int, - temperature: Union[float, Tuple[float, float], List[float]] = (2, 0), - device: Optional[Union[str, torch.device]] = None, - ) -> None: - """ - Set the discrete timesteps used for the diffusion chain (to be run before inference). - - Args: - num_inference_steps (`int`): - The number of diffusion steps used when generating samples with a pre-trained model. - temperature (`Union[float, Tuple[float, float], List[float]]`, *optional*, defaults to `(2, 0)`): - Temperature parameter(s) for controlling the randomness of sampling. If a tuple or list is provided, - temperatures will be linearly interpolated between the first and second values across all timesteps. If - a single value is provided, temperatures will be linearly interpolated from that value to 0.01. - device (`str` or `torch.device`, *optional*): - The device to which the timesteps and temperatures should be moved to. If `None`, the timesteps are not - moved. - """ + temperature: int | tuple[int, int] | list[int] = (2, 0), + device: str | torch.device = None, + ): self.timesteps = torch.arange(num_inference_steps, device=device).flip(0) if isinstance(temperature, (tuple, list)): @@ -141,35 +127,9 @@ def step( timestep: int, sample: torch.LongTensor, starting_mask_ratio: float = 1.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, return_dict: bool = True, - ) -> Union[AmusedSchedulerOutput, Tuple[torch.Tensor, torch.Tensor]]: - """ - Predict the sample at the previous timestep by masking tokens based on confidence scores. - - Args: - model_output (`torch.Tensor`): - The direct output from the learned diffusion model. Typically of shape `(batch_size, num_tokens, - codebook_size)` or `(batch_size, codebook_size, height, width)` for 2D inputs. - timestep (`int`): - The current discrete timestep in the diffusion chain. - sample (`torch.LongTensor`): - A current instance of a sample created by the diffusion process. Contains token IDs, with masked - positions indicated by `mask_token_id`. - starting_mask_ratio (`float`, *optional*, defaults to 1.0): - A multiplier applied to the mask ratio schedule. Values less than 1.0 will result in fewer tokens being - masked at each step. - generator (`torch.Generator`, *optional*): - A random number generator for reproducible sampling. - return_dict (`bool`, *optional*, defaults to `True`): - Whether to return an [`~schedulers.scheduling_amused.AmusedSchedulerOutput`] or a plain tuple. - - Returns: - [`~schedulers.scheduling_amused.AmusedSchedulerOutput`] or `tuple`: - If `return_dict` is `True`, [`~schedulers.scheduling_amused.AmusedSchedulerOutput`] is returned, - otherwise a tuple is returned where the first element is the sample tensor (`prev_sample`) and the - second element is the predicted original sample tensor (`pred_original_sample`). - """ + ) -> AmusedSchedulerOutput | tuple: two_dim_input = sample.ndim == 3 and model_output.ndim == 4 if two_dim_input: @@ -234,7 +194,7 @@ def add_noise( self, sample: torch.LongTensor, timesteps: int, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, ) -> torch.LongTensor: """ Add noise to a sample by randomly masking tokens according to the masking schedule. diff --git a/src/diffusers/schedulers/scheduling_consistency_decoder.py b/src/diffusers/schedulers/scheduling_consistency_decoder.py index 767fa9157f59..7f5e0cad22b9 100644 --- a/src/diffusers/schedulers/scheduling_consistency_decoder.py +++ b/src/diffusers/schedulers/scheduling_consistency_decoder.py @@ -1,6 +1,6 @@ import math from dataclasses import dataclass -from typing import Literal, Optional, Tuple, Union +from typing import Literal import torch @@ -14,7 +14,7 @@ def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -28,8 +28,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -40,6 +40,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -71,6 +78,22 @@ class ConsistencyDecoderSchedulerOutput(BaseOutput): class ConsistencyDecoderScheduler(SchedulerMixin, ConfigMixin): + """ + A scheduler for the consistency decoder used in Stable Diffusion pipelines. + + This scheduler implements a two-step denoising process using consistency models for decoding latent representations + into images. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, *optional*, defaults to `1024`): + The number of diffusion steps to train the model. + sigma_data (`float`, *optional*, defaults to `0.5`): + The standard deviation of the data distribution. Used for computing the skip and output scaling factors. + """ + order = 1 @register_to_config @@ -78,7 +101,7 @@ def __init__( self, num_train_timesteps: int = 1024, sigma_data: float = 0.5, - ): + ) -> None: betas = betas_for_alpha_bar(num_train_timesteps) alphas = 1.0 - betas @@ -97,8 +120,8 @@ def __init__( def set_timesteps( self, - num_inference_steps: Optional[int] = None, - device: Union[str, torch.device] = None, + num_inference_steps: int | None = None, + device: str | torch.device = None, ): if num_inference_steps != 2: raise ValueError("Currently more than 2 inference steps are not supported.") @@ -111,10 +134,18 @@ def set_timesteps( self.c_in = self.c_in.to(device) @property - def init_noise_sigma(self): + def init_noise_sigma(self) -> torch.Tensor: + """ + Return the standard deviation of the initial noise distribution. + + Returns: + `torch.Tensor`: + The initial noise sigma value from the precomputed `sqrt_one_minus_alphas_cumprod` at the first + timestep. + """ return self.sqrt_one_minus_alphas_cumprod[self.timesteps[0]] - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. @@ -134,11 +165,11 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None def step( self, model_output: torch.Tensor, - timestep: Union[float, torch.Tensor], + timestep: float | torch.Tensor, sample: torch.Tensor, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, return_dict: bool = True, - ) -> Union[ConsistencyDecoderSchedulerOutput, Tuple]: + ) -> ConsistencyDecoderSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). @@ -146,20 +177,20 @@ def step( Args: model_output (`torch.Tensor`): The direct output from the learned diffusion model. - timestep (`float`): + timestep (`float` or `torch.Tensor`): The current timestep in the diffusion chain. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. generator (`torch.Generator`, *optional*): - A random number generator. + A random number generator for reproducibility. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a - [`~schedulers.scheduling_consistency_models.ConsistencyDecoderSchedulerOutput`] or `tuple`. + [`~schedulers.scheduling_consistency_decoder.ConsistencyDecoderSchedulerOutput`] or `tuple`. Returns: - [`~schedulers.scheduling_consistency_models.ConsistencyDecoderSchedulerOutput`] or `tuple`: - If return_dict is `True`, - [`~schedulers.scheduling_consistency_models.ConsistencyDecoderSchedulerOutput`] is returned, otherwise + [`~schedulers.scheduling_consistency_decoder.ConsistencyDecoderSchedulerOutput`] or `tuple`: + If `return_dict` is `True`, + [`~schedulers.scheduling_consistency_decoder.ConsistencyDecoderSchedulerOutput`] is returned, otherwise a tuple is returned where the first element is the sample tensor. """ x_0 = self.c_out[timestep] * model_output + self.c_skip[timestep] * sample diff --git a/src/diffusers/schedulers/scheduling_consistency_models.py b/src/diffusers/schedulers/scheduling_consistency_models.py index 386a43db0f9c..d6bdf73ff9ce 100644 --- a/src/diffusers/schedulers/scheduling_consistency_models.py +++ b/src/diffusers/schedulers/scheduling_consistency_models.py @@ -13,7 +13,6 @@ # limitations under the License. from dataclasses import dataclass -from typing import List, Optional, Tuple, Union import numpy as np import torch @@ -66,7 +65,7 @@ class CMStochasticIterativeScheduler(SchedulerMixin, ConfigMixin): [paper](https://huggingface.co/papers/2206.00364). Defaults to 7.0 from the original implementation. clip_denoised (`bool`, defaults to `True`): Whether to clip the denoised outputs to `(-1, 1)`. - timesteps (`List` or `np.ndarray` or `torch.Tensor`, *optional*): + timesteps (`list` or `np.ndarray` or `torch.Tensor`, *optional*): An explicit timestep schedule that can be optionally specified. The timesteps are expected to be in increasing order. """ @@ -83,7 +82,7 @@ def __init__( s_noise: float = 1.0, rho: float = 7.0, clip_denoised: bool = True, - ): + ) -> None: # standard deviation of the initial noise distribution self.init_noise_sigma = sigma_max @@ -102,21 +101,29 @@ def __init__( self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication @property - def step_index(self): + def step_index(self) -> int: """ The index counter for current timestep. It will increase 1 after each scheduler step. + + Returns: + `int` or `None`: + The current step index, or `None` if not yet initialized. """ return self._step_index @property - def begin_index(self): + def begin_index(self) -> int: """ The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + + Returns: + `int` or `None`: + The begin index, or `None` if not yet set. """ return self._begin_index # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index - def set_begin_index(self, begin_index: int = 0): + def set_begin_index(self, begin_index: int = 0) -> None: """ Sets the begin index for the scheduler. This function should be run from pipeline before the inference. @@ -126,7 +133,7 @@ def set_begin_index(self, begin_index: int = 0): """ self._begin_index = begin_index - def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: float | torch.Tensor) -> torch.Tensor: """ Scales the consistency model input by `(sigma**2 + sigma_data**2) ** 0.5`. @@ -151,7 +158,7 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.T self.is_scale_input_called = True return sample - def sigma_to_t(self, sigmas: Union[float, np.ndarray]): + def sigma_to_t(self, sigmas: float | np.ndarray): """ Gets scaled timesteps from the Karras sigmas for input to the consistency model. @@ -160,8 +167,8 @@ def sigma_to_t(self, sigmas: Union[float, np.ndarray]): A single Karras sigma or an array of Karras sigmas. Returns: - `float` or `np.ndarray`: - A scaled input timestep or scaled input timestep array. + `np.ndarray`: + A scaled input timestep array. """ if not isinstance(sigmas, np.ndarray): sigmas = np.array(sigmas, dtype=np.float64) @@ -172,19 +179,19 @@ def sigma_to_t(self, sigmas: Union[float, np.ndarray]): def set_timesteps( self, - num_inference_steps: Optional[int] = None, - device: Union[str, torch.device] = None, - timesteps: Optional[List[int]] = None, + num_inference_steps: int | None = None, + device: str | torch.device = None, + timesteps: list[int] | None = None, ): """ Sets the timesteps used for the diffusion chain (to be run before inference). Args: - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*): The number of diffusion steps used when generating samples with a pre-trained model. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default timestep spacing strategy of equal spacing between timesteps is used. If `timesteps` is passed, `num_inference_steps` must be `None`. @@ -244,9 +251,19 @@ def set_timesteps( self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication # Modified _convert_to_karras implementation that takes in ramp as argument - def _convert_to_karras(self, ramp): - """Constructs the noise schedule of Karras et al. (2022).""" + def _convert_to_karras(self, ramp: np.ndarray) -> np.ndarray: + """ + Construct the noise schedule as proposed in [Elucidating the Design Space of Diffusion-Based Generative + Models](https://huggingface.co/papers/2206.00364). + Args: + ramp (`np.ndarray`): + A ramp array of values between 0 and 1 used to interpolate between sigma_min and sigma_max. + + Returns: + `np.ndarray`: + The Karras sigma schedule array. + """ sigma_min: float = self.config.sigma_min sigma_max: float = self.config.sigma_max @@ -256,14 +273,25 @@ def _convert_to_karras(self, ramp): sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho return sigmas - def get_scalings(self, sigma): + def get_scalings(self, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """ + Computes the scaling factors for the consistency model output. + + Args: + sigma (`torch.Tensor`): + The current sigma value in the noise schedule. + + Returns: + `tuple[torch.Tensor, torch.Tensor]`: + A tuple containing `c_skip` (scaling for the input sample) and `c_out` (scaling for the model output). + """ sigma_data = self.config.sigma_data c_skip = sigma_data**2 / (sigma**2 + sigma_data**2) c_out = sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5 return c_skip, c_out - def get_scalings_for_boundary_condition(self, sigma): + def get_scalings_for_boundary_condition(self, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """ Gets the scalings used in the consistency model parameterization (from Appendix C of the [paper](https://huggingface.co/papers/2303.01469)) to enforce boundary condition. @@ -275,7 +303,7 @@ def get_scalings_for_boundary_condition(self, sigma): The current sigma in the Karras sigma schedule. Returns: - `tuple`: + `tuple[torch.Tensor, torch.Tensor]`: A two-element tuple where `c_skip` (which weights the current sample) is the first element and `c_out` (which weights the consistency model output) is the second element. """ @@ -288,7 +316,7 @@ def get_scalings_for_boundary_condition(self, sigma): # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep def index_for_timestep( - self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None + self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None ) -> int: """ Find the index of a given timestep in the timestep schedule. @@ -318,7 +346,7 @@ def index_for_timestep( return indices[pos].item() # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: + def _init_step_index(self, timestep: float | torch.Tensor) -> None: """ Initialize the step index for the scheduler based on the given timestep. @@ -336,11 +364,11 @@ def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: def step( self, model_output: torch.Tensor, - timestep: Union[float, torch.Tensor], + timestep: float | torch.Tensor, sample: torch.Tensor, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, return_dict: bool = True, - ) -> Union[CMStochasticIterativeSchedulerOutput, Tuple]: + ) -> CMStochasticIterativeSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). @@ -348,13 +376,13 @@ def step( Args: model_output (`torch.Tensor`): The direct output from the learned diffusion model. - timestep (`float`): + timestep (`float` or `torch.Tensor`): The current timestep in the diffusion chain. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. generator (`torch.Generator`, *optional*): A random number generator. - return_dict (`bool`, *optional*, defaults to `True`): + return_dict (`bool`, defaults to `True`): Whether or not to return a [`~schedulers.scheduling_consistency_models.CMStochasticIterativeSchedulerOutput`] or `tuple`. @@ -406,7 +434,10 @@ def step( # Noise is not used for onestep sampling. if len(self.timesteps) > 1: noise = randn_tensor( - model_output.shape, dtype=model_output.dtype, device=model_output.device, generator=generator + model_output.shape, + dtype=model_output.dtype, + device=model_output.device, + generator=generator, ) else: noise = torch.zeros_like(model_output) @@ -475,5 +506,12 @@ def add_noise( noisy_samples = original_samples + noise * sigma return noisy_samples - def __len__(self): + def __len__(self) -> int: + """ + Returns the number of training timesteps. + + Returns: + `int`: + The number of training timesteps configured for the scheduler. + """ return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py index 103cca81c6a5..18ee272ef619 100644 --- a/src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py +++ b/src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py @@ -15,7 +15,7 @@ # DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver and https://github.com/NVlabs/edm import math -from typing import List, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -36,27 +36,30 @@ class CosineDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): methods the library implements for all schedulers such as loading and saving. Args: - sigma_min (`float`, *optional*, defaults to 0.3): + sigma_min (`float`, defaults to `0.3`): Minimum noise magnitude in the sigma schedule. This was set to 0.3 in Stable Audio Open [1]. - sigma_max (`float`, *optional*, defaults to 500): + sigma_max (`float`, defaults to `500`): Maximum noise magnitude in the sigma schedule. This was set to 500 in Stable Audio Open [1]. - sigma_data (`float`, *optional*, defaults to 1.0): + sigma_data (`float`, defaults to `1.0`): The standard deviation of the data distribution. This is set to 1.0 in Stable Audio Open [1]. - sigma_schedule (`str`, *optional*, defaults to `exponential`): - Sigma schedule to compute the `sigmas`. By default, we the schedule introduced in the EDM paper - (https://huggingface.co/papers/2206.00364). Other acceptable value is "exponential". The exponential - schedule was incorporated in this model: https://huggingface.co/stabilityai/cosxl. - num_train_timesteps (`int`, defaults to 1000): + sigma_schedule (`str`, defaults to `"exponential"`): + Sigma schedule to compute the `sigmas`. Must be one of `"exponential"` or `"karras"`. The exponential + schedule was incorporated in [stabilityai/cosxl](https://huggingface.co/stabilityai/cosxl). The Karras + schedule is introduced in the [EDM](https://huggingface.co/papers/2206.00364) paper. + num_train_timesteps (`int`, defaults to `1000`): The number of diffusion steps to train the model. - solver_order (`int`, defaults to 2): + solver_order (`int`, defaults to `2`): The DPMSolver order which can be `1` or `2`. It is recommended to use `solver_order=2`. - prediction_type (`str`, defaults to `v_prediction`, *optional*): - Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process), - `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen + prediction_type (`str`, defaults to `"v_prediction"`): + Prediction type of the scheduler function. Must be one of `"epsilon"` (predicts the noise of the diffusion + process), `"sample"` (directly predicts the noisy sample), or `"v_prediction"` (see section 2.4 of [Imagen Video](https://huggingface.co/papers/2210.02303) paper). - solver_type (`str`, defaults to `midpoint`): - Solver type for the second-order solver; can be `midpoint` or `heun`. The solver type slightly affects the - sample quality, especially for a small number of steps. It is recommended to use `midpoint` solvers. + rho (`float`, defaults to `7.0`): + The parameter for calculating the Karras sigma schedule from the EDM + [paper](https://huggingface.co/papers/2206.00364). + solver_type (`str`, defaults to `"midpoint"`): + Solver type for the second-order solver. Must be one of `"midpoint"` or `"heun"`. The solver type slightly + affects the sample quality, especially for a small number of steps. It is recommended to use `"midpoint"`. lower_order_final (`bool`, defaults to `True`): Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10. @@ -65,8 +68,9 @@ class CosineDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): richness. This can stabilize the sampling of the SDE variant of DPMSolver for small number of inference steps, but sometimes may result in blurring. final_sigmas_type (`str`, defaults to `"zero"`): - The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final - sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0. + The final `sigma` value for the noise schedule during the sampling process. Must be one of `"zero"` or + `"sigma_min"`. If `"sigma_min"`, the final sigma is the same as the last sigma in the training schedule. If + `"zero"`, the final sigma is set to 0. """ _compatibles = [] @@ -78,16 +82,16 @@ def __init__( sigma_min: float = 0.3, sigma_max: float = 500, sigma_data: float = 1.0, - sigma_schedule: str = "exponential", + sigma_schedule: Literal["exponential", "karras"] = "exponential", num_train_timesteps: int = 1000, solver_order: int = 2, - prediction_type: str = "v_prediction", + prediction_type: Literal["epsilon", "sample", "v_prediction"] = "v_prediction", rho: float = 7.0, - solver_type: str = "midpoint", + solver_type: Literal["midpoint", "heun"] = "midpoint", lower_order_final: bool = True, euler_at_final: bool = False, - final_sigmas_type: Optional[str] = "zero", # "zero", "sigma_min" - ): + final_sigmas_type: Literal["zero", "sigma_min"] = "zero", + ) -> None: if solver_type not in ["midpoint", "heun"]: if solver_type in ["logrho", "bh1", "bh2"]: self.register_to_config(solver_type="midpoint") @@ -113,26 +117,40 @@ def __init__( self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication @property - def init_noise_sigma(self): - # standard deviation of the initial noise distribution + def init_noise_sigma(self) -> float: + """ + The standard deviation of the initial noise distribution. + + Returns: + `float`: + The initial noise sigma value computed as `sqrt(sigma_max^2 + 1)`. + """ return (self.config.sigma_max**2 + 1) ** 0.5 @property - def step_index(self): + def step_index(self) -> int: """ The index counter for current timestep. It will increase 1 after each scheduler step. + + Returns: + `int` or `None`: + The current step index, or `None` if not yet initialized. """ return self._step_index @property - def begin_index(self): + def begin_index(self) -> int: """ The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + + Returns: + `int` or `None`: + The begin index, or `None` if not yet set. """ return self._begin_index # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index - def set_begin_index(self, begin_index: int = 0): + def set_begin_index(self, begin_index: int = 0) -> None: """ Sets the begin index for the scheduler. This function should be run from pipeline before the inference. @@ -143,19 +161,63 @@ def set_begin_index(self, begin_index: int = 0): self._begin_index = begin_index # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.precondition_inputs - def precondition_inputs(self, sample, sigma): + def precondition_inputs(self, sample: torch.Tensor, sigma: float | torch.Tensor) -> torch.Tensor: + """ + Precondition the input sample by scaling it according to the EDM formulation. + + Args: + sample (`torch.Tensor`): + The input sample tensor to precondition. + sigma (`float` or `torch.Tensor`): + The current sigma (noise level) value. + + Returns: + `torch.Tensor`: + The scaled input sample. + """ c_in = self._get_conditioning_c_in(sigma) scaled_sample = sample * c_in return scaled_sample - def precondition_noise(self, sigma): + def precondition_noise(self, sigma: float | torch.Tensor) -> torch.Tensor: + """ + Precondition the noise level by computing a normalized timestep representation. + + Args: + sigma (`float` or `torch.Tensor`): + The sigma (noise level) value to precondition. + + Returns: + `torch.Tensor`: + The preconditioned noise value computed as `atan(sigma) / pi * 2`. + """ if not isinstance(sigma, torch.Tensor): sigma = torch.tensor([sigma]) return sigma.atan() / math.pi * 2 # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.precondition_outputs - def precondition_outputs(self, sample, model_output, sigma): + def precondition_outputs( + self, + sample: torch.Tensor, + model_output: torch.Tensor, + sigma: float | torch.Tensor, + ) -> torch.Tensor: + """ + Precondition the model outputs according to the EDM formulation. + + Args: + sample (`torch.Tensor`): + The input sample tensor. + model_output (`torch.Tensor`): + The direct output from the learned diffusion model. + sigma (`float` or `torch.Tensor`): + The current sigma (noise level) value. + + Returns: + `torch.Tensor`: + The denoised sample computed by combining the skip connection and output scaling. + """ sigma_data = self.config.sigma_data c_skip = sigma_data**2 / (sigma**2 + sigma_data**2) @@ -171,15 +233,15 @@ def precondition_outputs(self, sample, model_output, sigma): return denoised # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.scale_model_input - def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: float | torch.Tensor) -> torch.Tensor: """ - Ensures interchangeability with schedulers that need to scale the denoising model input depending on the - current timestep. Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the Euler algorithm. + Scale the denoising model input to match the Euler algorithm. Ensures interchangeability with schedulers that + need to scale the denoising model input depending on the current timestep. Args: sample (`torch.Tensor`): - The input sample. - timestep (`int`, *optional*): + The input sample tensor. + timestep (`float` or `torch.Tensor`): The current timestep in the diffusion chain. Returns: @@ -195,12 +257,12 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.T self.is_scale_input_called = True return sample - def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torch.device] = None): + def set_timesteps(self, num_inference_steps: int | None = None, device: str | torch.device | None = None) -> None: """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). Args: - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*): The number of diffusion steps used when generating samples with a pre-trained model. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. @@ -242,8 +304,27 @@ def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torc self.noise_sampler = None # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler._compute_karras_sigmas - def _compute_karras_sigmas(self, ramp, sigma_min=None, sigma_max=None) -> torch.Tensor: - """Constructs the noise schedule of Karras et al. (2022).""" + def _compute_karras_sigmas( + self, + ramp: torch.Tensor, + sigma_min: float | None = None, + sigma_max: float | None = None, + ) -> torch.Tensor: + """ + Construct the noise schedule of [Karras et al. (2022)](https://huggingface.co/papers/2206.00364). + + Args: + ramp (`torch.Tensor`): + A tensor of values in [0, 1] representing the interpolation positions. + sigma_min (`float`, *optional*): + Minimum sigma value. If `None`, uses `self.config.sigma_min`. + sigma_max (`float`, *optional*): + Maximum sigma value. If `None`, uses `self.config.sigma_max`. + + Returns: + `torch.Tensor`: + The computed Karras sigma schedule. + """ sigma_min = sigma_min or self.config.sigma_min sigma_max = sigma_max or self.config.sigma_max @@ -254,10 +335,27 @@ def _compute_karras_sigmas(self, ramp, sigma_min=None, sigma_max=None) -> torch. return sigmas # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler._compute_exponential_sigmas - def _compute_exponential_sigmas(self, ramp, sigma_min=None, sigma_max=None) -> torch.Tensor: - """Implementation closely follows k-diffusion. - + def _compute_exponential_sigmas( + self, + ramp: torch.Tensor, + sigma_min: float | None = None, + sigma_max: float | None = None, + ) -> torch.Tensor: + """ + Compute the exponential sigma schedule. Implementation closely follows k-diffusion: https://github.com/crowsonkb/k-diffusion/blob/6ab5146d4a5ef63901326489f31f1d8e7dd36b48/k_diffusion/sampling.py#L26 + + Args: + ramp (`torch.Tensor`): + A tensor of values representing the interpolation positions. + sigma_min (`float`, *optional*): + Minimum sigma value. If `None`, uses `self.config.sigma_min`. + sigma_max (`float`, *optional*): + Maximum sigma value. If `None`, uses `self.config.sigma_max`. + + Returns: + `torch.Tensor`: + The computed exponential sigma schedule. """ sigma_min = sigma_min or self.config.sigma_min sigma_max = sigma_max or self.config.sigma_max @@ -265,7 +363,7 @@ def _compute_exponential_sigmas(self, ramp, sigma_min=None, sigma_max=None) -> t return sigmas # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t - def _sigma_to_t(self, sigma, log_sigmas): + def _sigma_to_t(self, sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray: """ Convert sigma values to corresponding timestep values through interpolation. @@ -301,7 +399,19 @@ def _sigma_to_t(self, sigma, log_sigmas): t = t.reshape(sigma.shape) return t - def _sigma_to_alpha_sigma_t(self, sigma): + def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """ + Convert sigma to alpha and sigma_t values for the diffusion process. + + Args: + sigma (`torch.Tensor`): + The sigma (noise level) value. + + Returns: + `tuple[torch.Tensor, torch.Tensor]`: + A tuple containing `alpha_t` (always 1 since inputs are pre-scaled) and `sigma_t` (same as input + sigma). + """ alpha_t = torch.tensor(1) # Inputs are pre-scaled before going into unet, so alpha_t = 1 sigma_t = sigma @@ -339,7 +449,7 @@ def dpm_solver_first_order_update( self, model_output: torch.Tensor, sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, + noise: torch.Tensor | None = None, ) -> torch.Tensor: """ One step for the first-order DPMSolver (equivalent to DDIM). @@ -354,7 +464,10 @@ def dpm_solver_first_order_update( `torch.Tensor`: The sample tensor at the previous timestep. """ - sigma_t, sigma_s = self.sigmas[self.step_index + 1], self.sigmas[self.step_index] + sigma_t, sigma_s = ( + self.sigmas[self.step_index + 1], + self.sigmas[self.step_index], + ) alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) alpha_s, sigma_s = self._sigma_to_alpha_sigma_t(sigma_s) lambda_t = torch.log(alpha_t) - torch.log(sigma_t) @@ -372,15 +485,15 @@ def dpm_solver_first_order_update( def multistep_dpm_solver_second_order_update( self, - model_output_list: List[torch.Tensor], + model_output_list: list[torch.Tensor], sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, + noise: torch.Tensor | None = None, ) -> torch.Tensor: """ One step for the second-order multistep DPMSolver. Args: - model_output_list (`List[torch.Tensor]`): + model_output_list (`list[torch.Tensor]`): The direct outputs from learned diffusion model at current and latter timesteps. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. @@ -430,7 +543,9 @@ def multistep_dpm_solver_second_order_update( # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep def index_for_timestep( - self, timestep: Union[int, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None + self, + timestep: int | torch.Tensor, + schedule_timesteps: torch.Tensor | None = None, ) -> int: """ Find the index for a given timestep in the schedule. @@ -464,7 +579,7 @@ def index_for_timestep( return step_index # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index - def _init_step_index(self, timestep): + def _init_step_index(self, timestep: int | torch.Tensor) -> None: """ Initialize the step_index counter for the scheduler. @@ -483,11 +598,11 @@ def _init_step_index(self, timestep): def step( self, model_output: torch.Tensor, - timestep: Union[int, torch.Tensor], + timestep: int | torch.Tensor, sample: torch.Tensor, - generator=None, + generator: torch.Generator | None = None, return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: + ) -> SchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the multistep DPMSolver. @@ -495,20 +610,19 @@ def step( Args: model_output (`torch.Tensor`): The direct output from learned diffusion model. - timestep (`int`): + timestep (`int` or `torch.Tensor`): The current discrete timestep in the diffusion chain. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. generator (`torch.Generator`, *optional*): A random number generator. - return_dict (`bool`): + return_dict (`bool`, defaults to `True`): Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`. Returns: [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`: If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a tuple is returned where the first element is the sample tensor. - """ if self.num_inference_steps is None: raise ValueError( @@ -540,7 +654,10 @@ def step( [g.initial_seed() for g in generator] if isinstance(generator, list) else generator.initial_seed() ) self.noise_sampler = BrownianTreeNoiseSampler( - model_output, sigma_min=self.config.sigma_min, sigma_max=self.config.sigma_max, seed=seed + model_output, + sigma_min=self.config.sigma_min, + sigma_max=self.config.sigma_max, + seed=seed, ) noise = self.noise_sampler(self.sigmas[self.step_index], self.sigmas[self.step_index + 1]).to( model_output.device @@ -612,9 +729,27 @@ def add_noise( return noisy_samples # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler._get_conditioning_c_in - def _get_conditioning_c_in(self, sigma): + def _get_conditioning_c_in(self, sigma: float | torch.Tensor) -> float | torch.Tensor: + """ + Compute the input conditioning factor for the EDM formulation. + + Args: + sigma (`float` or `torch.Tensor`): + The current sigma (noise level) value. + + Returns: + `float` or `torch.Tensor`: + The input conditioning factor `c_in`. + """ c_in = 1 / ((sigma**2 + self.config.sigma_data**2) ** 0.5) return c_in - def __len__(self): + def __len__(self) -> int: + """ + Returns the number of training timesteps. + + Returns: + `int`: + The number of training timesteps configured for the scheduler. + """ return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_ddim.py b/src/diffusers/schedulers/scheduling_ddim.py index d7fe29a72ac9..d29bb5b2593c 100644 --- a/src/diffusers/schedulers/scheduling_ddim.py +++ b/src/diffusers/schedulers/scheduling_ddim.py @@ -17,7 +17,7 @@ import math from dataclasses import dataclass -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -44,14 +44,14 @@ class DDIMSchedulerOutput(BaseOutput): """ prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None + pred_original_sample: torch.Tensor | None = None # Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -65,8 +65,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -77,6 +77,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -189,8 +196,8 @@ def __init__( num_train_timesteps: int = 1000, beta_start: float = 0.0001, beta_end: float = 0.02, - beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, + beta_schedule: str = "linear", + trained_betas: np.ndarray | list[float] | None = None, clip_sample: bool = True, set_alpha_to_one: bool = True, steps_offset: int = 0, @@ -235,7 +242,7 @@ def __init__( self.num_inference_steps = None self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64)) - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: int = None) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. @@ -324,14 +331,14 @@ def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: return sample - def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None) -> None: + def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). Args: num_inference_steps (`int`): The number of diffusion steps used when generating samples with a pre-trained model. - device (`Union[str, torch.device]`, *optional*): + device (`str | torch.device`, *optional*): The device to use for the timesteps. Raises: @@ -381,10 +388,10 @@ def step( sample: torch.Tensor, eta: float = 0.0, use_clipped_model_output: bool = False, - generator: Optional[torch.Generator] = None, - variance_noise: Optional[torch.Tensor] = None, + generator: torch.Generator | None = None, + variance_noise: torch.Tensor | None = None, return_dict: bool = True, - ) -> Union[DDIMSchedulerOutput, Tuple]: + ) -> DDIMSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). diff --git a/src/diffusers/schedulers/scheduling_ddim_cogvideox.py b/src/diffusers/schedulers/scheduling_ddim_cogvideox.py index f2683d1304ec..11cbb8a6b9f7 100644 --- a/src/diffusers/schedulers/scheduling_ddim_cogvideox.py +++ b/src/diffusers/schedulers/scheduling_ddim_cogvideox.py @@ -18,7 +18,7 @@ import math from dataclasses import dataclass -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -44,14 +44,14 @@ class DDIMSchedulerOutput(BaseOutput): """ prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None + pred_original_sample: torch.Tensor | None = None # Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -65,8 +65,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -77,6 +77,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -93,14 +100,13 @@ def alpha_bar_fn(t): return torch.tensor(betas, dtype=torch.float32) -def rescale_zero_terminal_snr(alphas_cumprod): +def rescale_zero_terminal_snr(alphas_cumprod: torch.Tensor) -> torch.Tensor: """ - Rescales betas to have zero terminal SNR Based on https://huggingface.co/papers/2305.08891 (Algorithm 1) - + Rescales betas to have zero terminal SNR Based on (Algorithm 1)[https://huggingface.co/papers/2305.08891] Args: - betas (`torch.Tensor`): - the betas that the scheduler is being initialized with. + alphas_cumprod (`torch.Tensor`): + The alphas cumulative products that the scheduler is being initialized with. Returns: `torch.Tensor`: rescaled betas with zero terminal SNR @@ -135,11 +141,11 @@ class CogVideoXDDIMScheduler(SchedulerMixin, ConfigMixin): Args: num_train_timesteps (`int`, defaults to 1000): The number of diffusion steps to train the model. - beta_start (`float`, defaults to 0.0001): + beta_start (`float`, defaults to 0.00085): The starting `beta` value of inference. - beta_end (`float`, defaults to 0.02): + beta_end (`float`, defaults to 0.0120): The final `beta` value. - beta_schedule (`str`, defaults to `"linear"`): + beta_schedule (`str`, defaults to `"scaled_linear"`): The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from `linear`, `scaled_linear`, or `squaredcos_cap_v2`. trained_betas (`np.ndarray`, *optional*): @@ -172,6 +178,8 @@ class CogVideoXDDIMScheduler(SchedulerMixin, ConfigMixin): Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and dark samples instead of limiting it to samples with medium brightness. Loosely related to [`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506). + snr_shift_scale (`float`, defaults to 3.0): + Shift scale for SNR. """ _compatibles = [e.name for e in KarrasDiffusionSchedulers] @@ -183,15 +191,15 @@ def __init__( num_train_timesteps: int = 1000, beta_start: float = 0.00085, beta_end: float = 0.0120, - beta_schedule: str = "scaled_linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, + beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "scaled_linear", + trained_betas: np.ndarray | list[float] | None = None, clip_sample: bool = True, set_alpha_to_one: bool = True, steps_offset: int = 0, - prediction_type: str = "epsilon", + prediction_type: Literal["epsilon", "sample", "v_prediction"] = "epsilon", clip_sample_range: float = 1.0, sample_max_value: float = 1.0, - timestep_spacing: str = "leading", + timestep_spacing: Literal["linspace", "leading", "trailing"] = "leading", rescale_betas_zero_snr: bool = False, snr_shift_scale: float = 3.0, ): @@ -201,7 +209,15 @@ def __init__( self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32) elif beta_schedule == "scaled_linear": # this schedule is very specific to the latent diffusion model. - self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float64) ** 2 + self.betas = ( + torch.linspace( + beta_start**0.5, + beta_end**0.5, + num_train_timesteps, + dtype=torch.float64, + ) + ** 2 + ) elif beta_schedule == "squaredcos_cap_v2": # Glide cosine schedule self.betas = betas_for_alpha_bar(num_train_timesteps) @@ -231,7 +247,7 @@ def __init__( self.num_inference_steps = None self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64)) - def _get_variance(self, timestep, prev_timestep): + def _get_variance(self, timestep: int, prev_timestep: int) -> torch.Tensor: alpha_prod_t = self.alphas_cumprod[timestep] alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod beta_prod_t = 1 - alpha_prod_t @@ -241,7 +257,7 @@ def _get_variance(self, timestep, prev_timestep): return variance - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: int = None) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. @@ -258,7 +274,11 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None """ return sample - def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None): + def set_timesteps( + self, + num_inference_steps: int, + device: str | torch.device | None = None, + ) -> None: """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -310,10 +330,10 @@ def step( sample: torch.Tensor, eta: float = 0.0, use_clipped_model_output: bool = False, - generator=None, - variance_noise: Optional[torch.Tensor] = None, + generator: torch.Generator | None = None, + variance_noise: torch.Tensor | None = None, return_dict: bool = True, - ) -> Union[DDIMSchedulerOutput, Tuple]: + ) -> DDIMSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). @@ -321,7 +341,7 @@ def step( Args: model_output (`torch.Tensor`): The direct output from learned diffusion model. - timestep (`float`): + timestep (`int`): The current discrete timestep in the diffusion chain. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. @@ -480,5 +500,5 @@ def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: tor velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample return velocity - def __len__(self): + def __len__(self) -> int: return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_ddim_flax.py b/src/diffusers/schedulers/scheduling_ddim_flax.py index 802d8f79779d..3bb44caef2bf 100644 --- a/src/diffusers/schedulers/scheduling_ddim_flax.py +++ b/src/diffusers/schedulers/scheduling_ddim_flax.py @@ -16,12 +16,12 @@ # and https://github.com/hojonathanho/diffusion from dataclasses import dataclass -from typing import Optional, Tuple, Union import flax import jax.numpy as jnp from ..configuration_utils import ConfigMixin, register_to_config +from ..utils import logging from .scheduling_utils_flax import ( CommonSchedulerState, FlaxKarrasDiffusionSchedulers, @@ -32,6 +32,9 @@ ) +logger = logging.get_logger(__name__) + + @flax.struct.dataclass class DDIMSchedulerState: common: CommonSchedulerState @@ -40,7 +43,7 @@ class DDIMSchedulerState: # setable values init_noise_sigma: jnp.ndarray timesteps: jnp.ndarray - num_inference_steps: Optional[int] = None + num_inference_steps: int = None @classmethod def create( @@ -117,7 +120,7 @@ def __init__( beta_start: float = 0.0001, beta_end: float = 0.02, beta_schedule: str = "linear", - trained_betas: Optional[jnp.ndarray] = None, + trained_betas: jnp.ndarray | None = None, clip_sample: bool = True, clip_sample_range: float = 1.0, set_alpha_to_one: bool = True, @@ -125,9 +128,13 @@ def __init__( prediction_type: str = "epsilon", dtype: jnp.dtype = jnp.float32, ): + logger.warning( + "Flax classes are deprecated and will be removed in Diffusers v1.0.0. We " + "recommend migrating to PyTorch classes or pinning your version of Diffusers." + ) self.dtype = dtype - def create_state(self, common: Optional[CommonSchedulerState] = None) -> DDIMSchedulerState: + def create_state(self, common: CommonSchedulerState | None = None) -> DDIMSchedulerState: if common is None: common = CommonSchedulerState.create(self) @@ -152,7 +159,10 @@ def create_state(self, common: Optional[CommonSchedulerState] = None) -> DDIMSch ) def scale_model_input( - self, state: DDIMSchedulerState, sample: jnp.ndarray, timestep: Optional[int] = None + self, + state: DDIMSchedulerState, + sample: jnp.ndarray, + timestep: int | None = None, ) -> jnp.ndarray: """ Args: @@ -166,7 +176,7 @@ def scale_model_input( return sample def set_timesteps( - self, state: DDIMSchedulerState, num_inference_steps: int, shape: Tuple = () + self, state: DDIMSchedulerState, num_inference_steps: int, shape: tuple = () ) -> DDIMSchedulerState: """ Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference. @@ -190,7 +200,9 @@ def set_timesteps( def _get_variance(self, state: DDIMSchedulerState, timestep, prev_timestep): alpha_prod_t = state.common.alphas_cumprod[timestep] alpha_prod_t_prev = jnp.where( - prev_timestep >= 0, state.common.alphas_cumprod[prev_timestep], state.final_alpha_cumprod + prev_timestep >= 0, + state.common.alphas_cumprod[prev_timestep], + state.final_alpha_cumprod, ) beta_prod_t = 1 - alpha_prod_t beta_prod_t_prev = 1 - alpha_prod_t_prev @@ -207,7 +219,7 @@ def step( sample: jnp.ndarray, eta: float = 0.0, return_dict: bool = True, - ) -> Union[FlaxDDIMSchedulerOutput, Tuple]: + ) -> FlaxDDIMSchedulerOutput | tuple: """ Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion process from the learned model outputs (most often the predicted noise). diff --git a/src/diffusers/schedulers/scheduling_ddim_inverse.py b/src/diffusers/schedulers/scheduling_ddim_inverse.py index 8ae13ad49d10..4d130266841b 100644 --- a/src/diffusers/schedulers/scheduling_ddim_inverse.py +++ b/src/diffusers/schedulers/scheduling_ddim_inverse.py @@ -16,7 +16,7 @@ # and https://github.com/hojonathanho/diffusion import math from dataclasses import dataclass -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -42,14 +42,14 @@ class DDIMSchedulerOutput(BaseOutput): """ prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None + pred_original_sample: torch.Tensor | None = None # Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -63,8 +63,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -75,6 +75,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -92,7 +99,7 @@ def alpha_bar_fn(t): # Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr -def rescale_zero_terminal_snr(betas): +def rescale_zero_terminal_snr(betas: torch.Tensor) -> torch.Tensor: """ Rescales betas to have zero terminal SNR Based on https://huggingface.co/papers/2305.08891 (Algorithm 1) @@ -180,14 +187,14 @@ def __init__( num_train_timesteps: int = 1000, beta_start: float = 0.0001, beta_end: float = 0.02, - beta_schedule: str = "linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, + beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear", + trained_betas: np.ndarray | list[float] | None = None, clip_sample: bool = True, set_alpha_to_one: bool = True, steps_offset: int = 0, - prediction_type: str = "epsilon", + prediction_type: Literal["epsilon", "sample", "v_prediction"] = "epsilon", clip_sample_range: float = 1.0, - timestep_spacing: str = "leading", + timestep_spacing: Literal["leading", "trailing"] = "leading", rescale_betas_zero_snr: bool = False, **kwargs, ): @@ -203,7 +210,15 @@ def __init__( self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32) elif beta_schedule == "scaled_linear": # this schedule is very specific to the latent diffusion model. - self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2 + self.betas = ( + torch.linspace( + beta_start**0.5, + beta_end**0.5, + num_train_timesteps, + dtype=torch.float32, + ) + ** 2 + ) elif beta_schedule == "squaredcos_cap_v2": # Glide cosine schedule self.betas = betas_for_alpha_bar(num_train_timesteps) @@ -232,7 +247,7 @@ def __init__( self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps).copy().astype(np.int64)) # Copied from diffusers.schedulers.scheduling_ddim.DDIMScheduler.scale_model_input - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. @@ -249,7 +264,11 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None """ return sample - def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None): + def set_timesteps( + self, + num_inference_steps: int, + device: str | torch.device | None = None, + ) -> None: """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -293,7 +312,7 @@ def step( timestep: int, sample: torch.Tensor, return_dict: bool = True, - ) -> Union[DDIMSchedulerOutput, Tuple]: + ) -> DDIMSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). @@ -301,20 +320,10 @@ def step( Args: model_output (`torch.Tensor`): The direct output from learned diffusion model. - timestep (`float`): + timestep (`int`): The current discrete timestep in the diffusion chain. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. - eta (`float`): - The weight of noise for added noise in diffusion step. - use_clipped_model_output (`bool`, defaults to `False`): - If `True`, computes "corrected" `model_output` from the clipped predicted original sample. Necessary - because predicted original sample is clipped to [-1, 1] when `self.config.clip_sample` is `True`. If no - clipping has happened, "corrected" `model_output` would coincide with the one provided as input and - `use_clipped_model_output` has no effect. - variance_noise (`torch.Tensor`): - Alternative to generating noise with `generator` by directly providing the noise for the variance - itself. Useful for methods such as [`CycleDiffusion`]. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~schedulers.scheduling_ddim_inverse.DDIMInverseSchedulerOutput`] or `tuple`. @@ -328,7 +337,8 @@ def step( # 1. get previous step value (=t+1) prev_timestep = timestep timestep = min( - timestep - self.config.num_train_timesteps // self.num_inference_steps, self.config.num_train_timesteps - 1 + timestep - self.config.num_train_timesteps // self.num_inference_steps, + self.config.num_train_timesteps - 1, ) # 2. compute alphas, betas @@ -371,5 +381,5 @@ def step( return (prev_sample, pred_original_sample) return DDIMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) - def __len__(self): + def __len__(self) -> int: return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_ddim_parallel.py b/src/diffusers/schedulers/scheduling_ddim_parallel.py index 10873a082fee..a5bbac60275c 100644 --- a/src/diffusers/schedulers/scheduling_ddim_parallel.py +++ b/src/diffusers/schedulers/scheduling_ddim_parallel.py @@ -17,7 +17,7 @@ import math from dataclasses import dataclass -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -44,14 +44,14 @@ class DDIMParallelSchedulerOutput(BaseOutput): """ prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None + pred_original_sample: torch.Tensor | None = None # Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -65,8 +65,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -77,6 +77,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -94,7 +101,7 @@ def alpha_bar_fn(t): # Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr -def rescale_zero_terminal_snr(betas): +def rescale_zero_terminal_snr(betas: torch.Tensor) -> torch.Tensor: """ Rescales betas to have zero terminal SNR Based on https://huggingface.co/papers/2305.08891 (Algorithm 1) @@ -195,8 +202,8 @@ def __init__( num_train_timesteps: int = 1000, beta_start: float = 0.0001, beta_end: float = 0.02, - beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, + beta_schedule: str = "linear", + trained_betas: np.ndarray | list[float] | None = None, clip_sample: bool = True, set_alpha_to_one: bool = True, steps_offset: int = 0, @@ -242,7 +249,7 @@ def __init__( self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64)) # Copied from diffusers.schedulers.scheduling_ddim.DDIMScheduler.scale_model_input - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. @@ -259,7 +266,7 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None """ return sample - def _get_variance(self, timestep, prev_timestep=None): + def _get_variance(self, timestep: int, prev_timestep: int | None = None) -> torch.Tensor: if prev_timestep is None: prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps @@ -272,7 +279,7 @@ def _get_variance(self, timestep, prev_timestep=None): return variance - def _batch_get_variance(self, t, prev_t): + def _batch_get_variance(self, t: torch.Tensor, prev_t: torch.Tensor) -> torch.Tensor: alpha_prod_t = self.alphas_cumprod[t] alpha_prod_t_prev = self.alphas_cumprod[torch.clip(prev_t, min=0)] alpha_prod_t_prev[prev_t < 0] = torch.tensor(1.0) @@ -328,14 +335,14 @@ def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: return sample # Copied from diffusers.schedulers.scheduling_ddim.DDIMScheduler.set_timesteps - def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None): + def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None) -> None: """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). Args: num_inference_steps (`int`): The number of diffusion steps used when generating samples with a pre-trained model. - device (`Union[str, torch.device]`, *optional*): + device (`str | torch.device`, *optional*): The device to use for the timesteps. Raises: @@ -385,10 +392,10 @@ def step( sample: torch.Tensor, eta: float = 0.0, use_clipped_model_output: bool = False, - generator=None, - variance_noise: Optional[torch.Tensor] = None, + generator: torch.Generator | None = None, + variance_noise: torch.Tensor | None = None, return_dict: bool = True, - ) -> Union[DDIMParallelSchedulerOutput, Tuple]: + ) -> DDIMParallelSchedulerOutput | tuple: """ Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion process from the learned model outputs (most often the predicted noise). @@ -399,11 +406,13 @@ def step( sample (`torch.Tensor`): current instance of sample being created by diffusion process. eta (`float`): weight of noise for added noise in diffusion step. - use_clipped_model_output (`bool`): if `True`, compute "corrected" `model_output` from the clipped - predicted original sample. Necessary because predicted original sample is clipped to [-1, 1] when - `self.config.clip_sample` is `True`. If no clipping has happened, "corrected" `model_output` would - coincide with the one provided as input and `use_clipped_model_output` will have not effect. - generator: random number generator. + use_clipped_model_output (`bool`, defaults to `False`): + If `True`, compute "corrected" `model_output` from the clipped predicted original sample. This + correction is necessary because the predicted original sample is clipped to [-1, 1] when + `self.config.clip_sample` is `True`. If no clipping occurred, the "corrected" `model_output` matches + the input and `use_clipped_model_output` has no effect. + generator (`torch.Generator`, *optional*): + Random number generator. variance_noise (`torch.Tensor`): instead of generating noise for the variance using `generator`, we can directly provide the noise for the variance itself. This is useful for methods such as CycleDiffusion. (https://huggingface.co/papers/2210.05559) @@ -489,7 +498,10 @@ def step( if variance_noise is None: variance_noise = randn_tensor( - model_output.shape, generator=generator, device=model_output.device, dtype=model_output.dtype + model_output.shape, + generator=generator, + device=model_output.device, + dtype=model_output.dtype, ) variance = std_dev_t * variance_noise @@ -506,7 +518,7 @@ def step( def batch_step_no_noise( self, model_output: torch.Tensor, - timesteps: List[int], + timesteps: torch.Tensor, sample: torch.Tensor, eta: float = 0.0, use_clipped_model_output: bool = False, @@ -521,7 +533,7 @@ def batch_step_no_noise( Args: model_output (`torch.Tensor`): direct output from learned diffusion model. - timesteps (`List[int]`): + timesteps (`torch.Tensor`): current discrete timesteps in the diffusion chain. This is now a list of integers. sample (`torch.Tensor`): current instance of sample being created by diffusion process. @@ -689,5 +701,5 @@ def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: tor velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample return velocity - def __len__(self): + def __len__(self) -> int: return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_ddpm.py b/src/diffusers/schedulers/scheduling_ddpm.py index ded88b8e1e0a..972c46c6e930 100644 --- a/src/diffusers/schedulers/scheduling_ddpm.py +++ b/src/diffusers/schedulers/scheduling_ddpm.py @@ -16,7 +16,7 @@ import math from dataclasses import dataclass -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -42,13 +42,13 @@ class DDPMSchedulerOutput(BaseOutput): """ prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None + pred_original_sample: torch.Tensor | None = None def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -62,8 +62,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -74,6 +74,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -183,9 +190,14 @@ def __init__( beta_start: float = 0.0001, beta_end: float = 0.02, beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2", "sigmoid"] = "linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, + trained_betas: np.ndarray | list[float] | None = None, variance_type: Literal[ - "fixed_small", "fixed_small_log", "fixed_large", "fixed_large_log", "learned", "learned_range" + "fixed_small", + "fixed_small_log", + "fixed_large", + "fixed_large_log", + "learned", + "learned_range", ] = "fixed_small", clip_sample: bool = True, prediction_type: Literal["epsilon", "sample", "v_prediction"] = "epsilon", @@ -203,10 +215,20 @@ def __init__( self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32) elif beta_schedule == "scaled_linear": # this schedule is very specific to the latent diffusion model. - self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2 + self.betas = ( + torch.linspace( + beta_start**0.5, + beta_end**0.5, + num_train_timesteps, + dtype=torch.float32, + ) + ** 2 + ) elif beta_schedule == "squaredcos_cap_v2": # Glide cosine schedule self.betas = betas_for_alpha_bar(num_train_timesteps) + elif beta_schedule == "laplace": + self.betas = betas_for_alpha_bar(num_train_timesteps, alpha_transform_type="laplace") elif beta_schedule == "sigmoid": # GeoDiff sigmoid schedule betas = torch.linspace(-6, 6, num_train_timesteps) @@ -232,7 +254,7 @@ def __init__( self.variance_type = variance_type - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. @@ -251,20 +273,20 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None def set_timesteps( self, - num_inference_steps: Optional[int] = None, - device: Union[str, torch.device] = None, - timesteps: Optional[List[int]] = None, + num_inference_steps: int = None, + device: str | torch.device = None, + timesteps: list[int] | None = None, ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). Args: - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*): The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default timestep spacing strategy of equal spacing between timesteps is used. If `timesteps` is passed, `num_inference_steps` must be `None`. @@ -326,10 +348,16 @@ def set_timesteps( def _get_variance( self, t: int, - predicted_variance: Optional[torch.Tensor] = None, - variance_type: Optional[ - Literal["fixed_small", "fixed_small_log", "fixed_large", "fixed_large_log", "learned", "learned_range"] - ] = None, + predicted_variance: torch.Tensor | None = None, + variance_type: Literal[ + "fixed_small", + "fixed_small_log", + "fixed_large", + "fixed_large_log", + "learned", + "learned_range", + ] + | None = None, ) -> torch.Tensor: """ Compute the variance for a given timestep according to the specified variance type. @@ -435,9 +463,9 @@ def step( model_output: torch.Tensor, timestep: int, sample: torch.Tensor, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, return_dict: bool = True, - ) -> Union[DDPMSchedulerOutput, Tuple]: + ) -> DDPMSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). @@ -463,7 +491,10 @@ def step( prev_t = self.previous_timestep(t) - if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in ["learned", "learned_range"]: + if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in [ + "learned", + "learned_range", + ]: model_output, predicted_variance = torch.split(model_output, sample.shape[1], dim=1) else: predicted_variance = None @@ -512,7 +543,10 @@ def step( if t > 0: device = model_output.device variance_noise = randn_tensor( - model_output.shape, generator=generator, device=device, dtype=model_output.dtype + model_output.shape, + generator=generator, + device=device, + dtype=model_output.dtype, ) if self.variance_type == "fixed_small_log": variance = self._get_variance(t, predicted_variance=predicted_variance) * variance_noise @@ -611,7 +645,7 @@ def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: tor def __len__(self) -> int: return self.config.num_train_timesteps - def previous_timestep(self, timestep: int) -> int: + def previous_timestep(self, timestep: int) -> int | torch.Tensor: """ Compute the previous timestep in the diffusion chain. @@ -620,7 +654,7 @@ def previous_timestep(self, timestep: int) -> int: The current timestep. Returns: - `int`: + `int` or `torch.Tensor`: The previous timestep. """ if self.custom_timesteps or self.num_inference_steps: diff --git a/src/diffusers/schedulers/scheduling_ddpm_flax.py b/src/diffusers/schedulers/scheduling_ddpm_flax.py index a3264f54f572..7840ff729488 100644 --- a/src/diffusers/schedulers/scheduling_ddpm_flax.py +++ b/src/diffusers/schedulers/scheduling_ddpm_flax.py @@ -15,13 +15,13 @@ # DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim from dataclasses import dataclass -from typing import Optional, Tuple, Union import flax import jax import jax.numpy as jnp from ..configuration_utils import ConfigMixin, register_to_config +from ..utils import logging from .scheduling_utils_flax import ( CommonSchedulerState, FlaxKarrasDiffusionSchedulers, @@ -32,6 +32,9 @@ ) +logger = logging.get_logger(__name__) + + @flax.struct.dataclass class DDPMSchedulerState: common: CommonSchedulerState @@ -39,10 +42,15 @@ class DDPMSchedulerState: # setable values init_noise_sigma: jnp.ndarray timesteps: jnp.ndarray - num_inference_steps: Optional[int] = None + num_inference_steps: int = None @classmethod - def create(cls, common: CommonSchedulerState, init_noise_sigma: jnp.ndarray, timesteps: jnp.ndarray): + def create( + cls, + common: CommonSchedulerState, + init_noise_sigma: jnp.ndarray, + timesteps: jnp.ndarray, + ): return cls(common=common, init_noise_sigma=init_noise_sigma, timesteps=timesteps) @@ -99,15 +107,19 @@ def __init__( beta_start: float = 0.0001, beta_end: float = 0.02, beta_schedule: str = "linear", - trained_betas: Optional[jnp.ndarray] = None, + trained_betas: jnp.ndarray | None = None, variance_type: str = "fixed_small", clip_sample: bool = True, prediction_type: str = "epsilon", dtype: jnp.dtype = jnp.float32, ): + logger.warning( + "Flax classes are deprecated and will be removed in Diffusers v1.0.0. We " + "recommend migrating to PyTorch classes or pinning your version of Diffusers." + ) self.dtype = dtype - def create_state(self, common: Optional[CommonSchedulerState] = None) -> DDPMSchedulerState: + def create_state(self, common: CommonSchedulerState | None = None) -> DDPMSchedulerState: if common is None: common = CommonSchedulerState.create(self) @@ -123,7 +135,10 @@ def create_state(self, common: Optional[CommonSchedulerState] = None) -> DDPMSch ) def scale_model_input( - self, state: DDPMSchedulerState, sample: jnp.ndarray, timestep: Optional[int] = None + self, + state: DDPMSchedulerState, + sample: jnp.ndarray, + timestep: int | None = None, ) -> jnp.ndarray: """ Args: @@ -137,7 +152,7 @@ def scale_model_input( return sample def set_timesteps( - self, state: DDPMSchedulerState, num_inference_steps: int, shape: Tuple = () + self, state: DDPMSchedulerState, num_inference_steps: int, shape: tuple = () ) -> DDPMSchedulerState: """ Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference. @@ -198,9 +213,9 @@ def step( model_output: jnp.ndarray, timestep: int, sample: jnp.ndarray, - key: Optional[jax.Array] = None, + key: jax.Array | None = None, return_dict: bool = True, - ) -> Union[FlaxDDPMSchedulerOutput, Tuple]: + ) -> FlaxDDPMSchedulerOutput | tuple: """ Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion process from the learned model outputs (most often the predicted noise). diff --git a/src/diffusers/schedulers/scheduling_ddpm_parallel.py b/src/diffusers/schedulers/scheduling_ddpm_parallel.py index 941fc16be080..725958b01a44 100644 --- a/src/diffusers/schedulers/scheduling_ddpm_parallel.py +++ b/src/diffusers/schedulers/scheduling_ddpm_parallel.py @@ -16,7 +16,7 @@ import math from dataclasses import dataclass -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -43,14 +43,14 @@ class DDPMParallelSchedulerOutput(BaseOutput): """ prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None + pred_original_sample: torch.Tensor | None = None # Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -64,8 +64,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -76,6 +76,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -142,38 +149,41 @@ class DDPMParallelScheduler(SchedulerMixin, ConfigMixin): For more details, see the original paper: https://huggingface.co/papers/2006.11239 Args: - num_train_timesteps (`int`): number of diffusion steps used to train the model. - beta_start (`float`): the starting `beta` value of inference. - beta_end (`float`): the final `beta` value. - beta_schedule (`str`): - the beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + beta_start (`float`, defaults to 0.0001): + The starting `beta` value of inference. + beta_end (`float`, defaults to 0.02): + The final `beta` value. + beta_schedule (`str`, defaults to `"linear"`): + The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from `linear`, `scaled_linear`, `squaredcos_cap_v2` or `sigmoid`. - trained_betas (`np.ndarray`, optional): - option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc. - variance_type (`str`): - options to clip the variance used when adding noise to the denoised sample. Choose from `fixed_small`, + trained_betas (`np.ndarray`, *optional*): + Option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc. + variance_type (`str`, defaults to `"fixed_small"`): + Options to clip the variance used when adding noise to the denoised sample. Choose from `fixed_small`, `fixed_small_log`, `fixed_large`, `fixed_large_log`, `learned` or `learned_range`. - clip_sample (`bool`, default `True`): - option to clip predicted sample for numerical stability. - clip_sample_range (`float`, default `1.0`): - the maximum magnitude for sample clipping. Valid only when `clip_sample=True`. - prediction_type (`str`, default `epsilon`, optional): - prediction type of the scheduler function, one of `epsilon` (predicting the noise of the diffusion + clip_sample (`bool`, defaults to `True`): + Option to clip predicted sample for numerical stability. + prediction_type (`str`, defaults to `"epsilon"`): + Prediction type of the scheduler function, one of `epsilon` (predicting the noise of the diffusion process), `sample` (directly predicting the noisy sample`) or `v_prediction` (see section 2.4 https://huggingface.co/papers/2210.02303) - thresholding (`bool`, default `False`): - whether to use the "dynamic thresholding" method (introduced by Imagen, + thresholding (`bool`, defaults to `False`): + Whether to use the "dynamic thresholding" method (introduced by Imagen, https://huggingface.co/papers/2205.11487). Note that the thresholding method is unsuitable for latent-space diffusion models (such as stable-diffusion). - dynamic_thresholding_ratio (`float`, default `0.995`): - the ratio for the dynamic thresholding method. Default is `0.995`, the same as Imagen + dynamic_thresholding_ratio (`float`, defaults to 0.995): + The ratio for the dynamic thresholding method. Default is `0.995`, the same as Imagen (https://huggingface.co/papers/2205.11487). Valid only when `thresholding=True`. - sample_max_value (`float`, default `1.0`): - the threshold value for dynamic thresholding. Valid only when `thresholding=True`. - timestep_spacing (`str`, default `"leading"`): + clip_sample_range (`float`, defaults to 1.0): + The maximum magnitude for sample clipping. Valid only when `clip_sample=True`. + sample_max_value (`float`, defaults to 1.0): + The threshold value for dynamic thresholding. Valid only when `thresholding=True`. + timestep_spacing (`str`, defaults to `"leading"`): The way the timesteps should be scaled. Refer to Table 2. of [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. - steps_offset (`int`, default `0`): + steps_offset (`int`, defaults to 0): An offset added to the inference steps, as required by some model families. rescale_betas_zero_snr (`bool`, defaults to `False`): Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and @@ -193,9 +203,14 @@ def __init__( beta_start: float = 0.0001, beta_end: float = 0.02, beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2", "sigmoid"] = "linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, + trained_betas: np.ndarray | list[float] | None = None, variance_type: Literal[ - "fixed_small", "fixed_small_log", "fixed_large", "fixed_large_log", "learned", "learned_range" + "fixed_small", + "fixed_small_log", + "fixed_large", + "fixed_large_log", + "learned", + "learned_range", ] = "fixed_small", clip_sample: bool = True, prediction_type: Literal["epsilon", "sample", "v_prediction"] = "epsilon", @@ -213,10 +228,20 @@ def __init__( self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32) elif beta_schedule == "scaled_linear": # this schedule is very specific to the latent diffusion model. - self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2 + self.betas = ( + torch.linspace( + beta_start**0.5, + beta_end**0.5, + num_train_timesteps, + dtype=torch.float32, + ) + ** 2 + ) elif beta_schedule == "squaredcos_cap_v2": # Glide cosine schedule self.betas = betas_for_alpha_bar(num_train_timesteps) + elif beta_schedule == "laplace": + self.betas = betas_for_alpha_bar(num_train_timesteps, alpha_transform_type="laplace") elif beta_schedule == "sigmoid": # GeoDiff sigmoid schedule betas = torch.linspace(-6, 6, num_train_timesteps) @@ -243,7 +268,7 @@ def __init__( self.variance_type = variance_type # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.scale_model_input - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. @@ -263,20 +288,20 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.set_timesteps def set_timesteps( self, - num_inference_steps: Optional[int] = None, - device: Union[str, torch.device] = None, - timesteps: Optional[List[int]] = None, + num_inference_steps: int = None, + device: str | torch.device = None, + timesteps: list[int] | None = None, ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). Args: - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*): The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default timestep spacing strategy of equal spacing between timesteps is used. If `timesteps` is passed, `num_inference_steps` must be `None`. @@ -339,10 +364,16 @@ def set_timesteps( def _get_variance( self, t: int, - predicted_variance: Optional[torch.Tensor] = None, - variance_type: Optional[ - Literal["fixed_small", "fixed_small_log", "fixed_large", "fixed_large_log", "learned", "learned_range"] - ] = None, + predicted_variance: torch.Tensor | None = None, + variance_type: Literal[ + "fixed_small", + "fixed_small_log", + "fixed_large", + "fixed_large_log", + "learned", + "learned_range", + ] + | None = None, ) -> torch.Tensor: """ Compute the variance for a given timestep according to the specified variance type. @@ -449,9 +480,9 @@ def step( model_output: torch.Tensor, timestep: int, sample: torch.Tensor, - generator=None, + generator: torch.Generator | None = None, return_dict: bool = True, - ) -> Union[DDPMParallelSchedulerOutput, Tuple]: + ) -> DDPMParallelSchedulerOutput | tuple: """ Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion process from the learned model outputs (most often the predicted noise). @@ -461,7 +492,8 @@ def step( timestep (`int`): current discrete timestep in the diffusion chain. sample (`torch.Tensor`): current instance of sample being created by diffusion process. - generator: random number generator. + generator (`torch.Generator`, *optional*): + Random number generator. return_dict (`bool`): option for returning tuple rather than DDPMParallelSchedulerOutput class Returns: @@ -474,7 +506,10 @@ def step( prev_t = self.previous_timestep(t) - if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in ["learned", "learned_range"]: + if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in [ + "learned", + "learned_range", + ]: model_output, predicted_variance = torch.split(model_output, sample.shape[1], dim=1) else: predicted_variance = None @@ -523,7 +558,10 @@ def step( if t > 0: device = model_output.device variance_noise = randn_tensor( - model_output.shape, generator=generator, device=device, dtype=model_output.dtype + model_output.shape, + generator=generator, + device=device, + dtype=model_output.dtype, ) if self.variance_type == "fixed_small_log": variance = self._get_variance(t, predicted_variance=predicted_variance) * variance_noise @@ -546,7 +584,7 @@ def step( def batch_step_no_noise( self, model_output: torch.Tensor, - timesteps: List[int], + timesteps: torch.Tensor, sample: torch.Tensor, ) -> torch.Tensor: """ @@ -559,8 +597,8 @@ def batch_step_no_noise( Args: model_output (`torch.Tensor`): direct output from learned diffusion model. - timesteps (`List[int]`): - current discrete timesteps in the diffusion chain. This is now a list of integers. + timesteps (`torch.Tensor`): + Current discrete timesteps in the diffusion chain. This is a tensor of integers. sample (`torch.Tensor`): current instance of sample being created by diffusion process. @@ -574,7 +612,10 @@ def batch_step_no_noise( t = t.view(-1, *([1] * (model_output.ndim - 1))) prev_t = prev_t.view(-1, *([1] * (model_output.ndim - 1))) - if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in ["learned", "learned_range"]: + if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in [ + "learned", + "learned_range", + ]: model_output, predicted_variance = torch.split(model_output, sample.shape[1], dim=1) else: pass @@ -705,7 +746,7 @@ def __len__(self): return self.config.num_train_timesteps # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.previous_timestep - def previous_timestep(self, timestep): + def previous_timestep(self, timestep: int) -> int | torch.Tensor: """ Compute the previous timestep in the diffusion chain. @@ -714,7 +755,7 @@ def previous_timestep(self, timestep): The current timestep. Returns: - `int`: + `int` or `torch.Tensor`: The previous timestep. """ if self.custom_timesteps or self.num_inference_steps: diff --git a/src/diffusers/schedulers/scheduling_ddpm_wuerstchen.py b/src/diffusers/schedulers/scheduling_ddpm_wuerstchen.py index 71f08277ebd7..e3b87b789a03 100644 --- a/src/diffusers/schedulers/scheduling_ddpm_wuerstchen.py +++ b/src/diffusers/schedulers/scheduling_ddpm_wuerstchen.py @@ -17,7 +17,6 @@ import math from dataclasses import dataclass -from typing import List, Optional, Tuple, Union import torch @@ -125,7 +124,7 @@ def _alpha_cumprod(self, t, device): ) ** 2 / self._init_alpha_cumprod.to(device) return alpha_cumprod.clamp(0.0001, 0.9999) - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: int = None) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. @@ -142,14 +141,14 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None def set_timesteps( self, num_inference_steps: int = None, - timesteps: Optional[List[int]] = None, - device: Union[str, torch.device] = None, + timesteps: list[int] | None = None, + device: str | torch.device = None, ): """ Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference. Args: - num_inference_steps (`Dict[float, int]`): + num_inference_steps (`dict[float, int]`): the number of diffusion steps used when generating samples with a pre-trained model. If passed, then `timesteps` must be `None`. device (`str` or `torch.device`, optional): @@ -168,7 +167,7 @@ def step( sample: torch.Tensor, generator=None, return_dict: bool = True, - ) -> Union[DDPMWuerstchenSchedulerOutput, Tuple]: + ) -> DDPMWuerstchenSchedulerOutput | tuple: """ Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion process from the learned model outputs (most often the predicted noise). diff --git a/src/diffusers/schedulers/scheduling_deis_multistep.py b/src/diffusers/schedulers/scheduling_deis_multistep.py index b7d64fc00bae..c22d2b5b439d 100644 --- a/src/diffusers/schedulers/scheduling_deis_multistep.py +++ b/src/diffusers/schedulers/scheduling_deis_multistep.py @@ -16,7 +16,7 @@ # The codebase is modified based on https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py import math -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -34,7 +34,7 @@ def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -48,8 +48,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -60,6 +60,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -93,7 +100,7 @@ class DEISMultistepScheduler(SchedulerMixin, ConfigMixin): beta_schedule (`"linear"`, `"scaled_linear"`, or `"squaredcos_cap_v2"`, defaults to `"linear"`): The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from `linear`, `scaled_linear`, or `squaredcos_cap_v2`. - trained_betas (`np.ndarray` or `List[float]`, *optional*): + trained_betas (`np.ndarray` or `list[float]`, *optional*): Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. solver_order (`int`, defaults to `2`): The DEIS order which can be `1` or `2` or `3`. It is recommended to use `solver_order=2` for guided @@ -148,7 +155,7 @@ def __init__( beta_start: float = 0.0001, beta_end: float = 0.02, beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, + trained_betas: np.ndarray | list[float] | None = None, solver_order: int = 2, prediction_type: Literal["epsilon", "sample", "v_prediction", "flow_prediction"] = "epsilon", thresholding: bool = False, @@ -157,11 +164,11 @@ def __init__( algorithm_type: Literal["deis"] = "deis", solver_type: Literal["logrho"] = "logrho", lower_order_final: bool = True, - use_karras_sigmas: Optional[bool] = False, - use_exponential_sigmas: Optional[bool] = False, - use_beta_sigmas: Optional[bool] = False, - use_flow_sigmas: Optional[bool] = False, - flow_shift: Optional[float] = 1.0, + use_karras_sigmas: bool = False, + use_exponential_sigmas: bool = False, + use_beta_sigmas: bool = False, + use_flow_sigmas: bool = False, + flow_shift: float = 1.0, timestep_spacing: Literal["linspace", "leading", "trailing"] = "linspace", steps_offset: int = 0, use_dynamic_shifting: bool = False, @@ -238,14 +245,14 @@ def __init__( self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication @property - def step_index(self) -> Optional[int]: + def step_index(self) -> int: """ The index counter for current timestep. It will increase 1 after each scheduler step. """ return self._step_index @property - def begin_index(self) -> Optional[int]: + def begin_index(self) -> int: """ The index for the first timestep. It should be set from pipeline with `set_begin_index` method. """ @@ -262,12 +269,7 @@ def set_begin_index(self, begin_index: int = 0) -> None: """ self._begin_index = begin_index - def set_timesteps( - self, - num_inference_steps: int, - device: Union[str, torch.device] = None, - mu: Optional[float] = None, - ) -> None: + def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None, mu: float | None = None): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -433,7 +435,7 @@ def _sigma_to_t(self, sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray: return t # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._sigma_to_alpha_sigma_t - def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """ Convert sigma values to alpha_t and sigma_t values. @@ -442,7 +444,7 @@ def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, to The sigma value(s) to convert. Returns: - `Tuple[torch.Tensor, torch.Tensor]`: + `tuple[torch.Tensor, torch.Tensor]`: A tuple containing (alpha_t, sigma_t) values. """ if self.config.use_flow_sigmas: @@ -699,7 +701,7 @@ def deis_first_order_update( def multistep_deis_second_order_update( self, - model_output_list: List[torch.Tensor], + model_output_list: list[torch.Tensor], *args, sample: torch.Tensor = None, **kwargs, @@ -708,7 +710,7 @@ def multistep_deis_second_order_update( One step for the second-order multistep DEIS. Args: - model_output_list (`List[torch.Tensor]`): + model_output_list (`list[torch.Tensor]`): The direct outputs from learned diffusion model at current and latter timesteps. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. @@ -772,7 +774,7 @@ def ind_fn(t, b, c): def multistep_deis_third_order_update( self, - model_output_list: List[torch.Tensor], + model_output_list: list[torch.Tensor], *args, sample: torch.Tensor = None, **kwargs, @@ -781,7 +783,7 @@ def multistep_deis_third_order_update( One step for the third-order multistep DEIS. Args: - model_output_list (`List[torch.Tensor]`): + model_output_list (`list[torch.Tensor]`): The direct outputs from learned diffusion model at current and latter timesteps. sample (`torch.Tensor`): A current instance of a sample created by diffusion process. @@ -860,7 +862,9 @@ def ind_fn(t, b, c, d): # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep def index_for_timestep( - self, timestep: Union[int, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None + self, + timestep: int | torch.Tensor, + schedule_timesteps: torch.Tensor | None = None, ) -> int: """ Find the index for a given timestep in the schedule. @@ -894,7 +898,7 @@ def index_for_timestep( return step_index # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index - def _init_step_index(self, timestep: Union[int, torch.Tensor]) -> None: + def _init_step_index(self, timestep: int | torch.Tensor) -> None: """ Initialize the step_index counter for the scheduler. @@ -913,10 +917,10 @@ def _init_step_index(self, timestep: Union[int, torch.Tensor]) -> None: def step( self, model_output: torch.Tensor, - timestep: Union[int, torch.Tensor], + timestep: int | torch.Tensor, sample: torch.Tensor, return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: + ) -> SchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the multistep DEIS. diff --git a/src/diffusers/schedulers/scheduling_dpm_cogvideox.py b/src/diffusers/schedulers/scheduling_dpm_cogvideox.py index 0a9082208cf4..4a238a5408ae 100644 --- a/src/diffusers/schedulers/scheduling_dpm_cogvideox.py +++ b/src/diffusers/schedulers/scheduling_dpm_cogvideox.py @@ -18,7 +18,7 @@ import math from dataclasses import dataclass -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -45,14 +45,14 @@ class DDIMSchedulerOutput(BaseOutput): """ prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None + pred_original_sample: torch.Tensor | None = None # Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -66,8 +66,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -78,6 +78,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -98,7 +105,6 @@ def rescale_zero_terminal_snr(alphas_cumprod): """ Rescales betas to have zero terminal SNR Based on https://huggingface.co/papers/2305.08891 (Algorithm 1) - Args: betas (`torch.Tensor`): the betas that the scheduler is being initialized with. @@ -168,11 +174,14 @@ class CogVideoXDPMScheduler(SchedulerMixin, ConfigMixin): The threshold value for dynamic thresholding. Valid only when `thresholding=True`. timestep_spacing (`str`, defaults to `"leading"`): The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and - Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. + Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. Choose from + `leading`, `linspace` or `trailing`. rescale_betas_zero_snr (`bool`, defaults to `False`): Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and dark samples instead of limiting it to samples with medium brightness. Loosely related to [`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506). + snr_shift_scale (`float`, defaults to 3.0): + Shift scale for SNR. """ _compatibles = [e.name for e in KarrasDiffusionSchedulers] @@ -184,15 +193,15 @@ def __init__( num_train_timesteps: int = 1000, beta_start: float = 0.00085, beta_end: float = 0.0120, - beta_schedule: str = "scaled_linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, + beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "scaled_linear", + trained_betas: np.ndarray | list[float] | None = None, clip_sample: bool = True, set_alpha_to_one: bool = True, steps_offset: int = 0, - prediction_type: str = "epsilon", + prediction_type: Literal["epsilon", "sample", "v_prediction"] = "epsilon", clip_sample_range: float = 1.0, sample_max_value: float = 1.0, - timestep_spacing: str = "leading", + timestep_spacing: Literal["leading", "linspace", "trailing"] = "leading", rescale_betas_zero_snr: bool = False, snr_shift_scale: float = 3.0, ): @@ -202,7 +211,15 @@ def __init__( self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32) elif beta_schedule == "scaled_linear": # this schedule is very specific to the latent diffusion model. - self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float64) ** 2 + self.betas = ( + torch.linspace( + beta_start**0.5, + beta_end**0.5, + num_train_timesteps, + dtype=torch.float64, + ) + ** 2 + ) elif beta_schedule == "squaredcos_cap_v2": # Glide cosine schedule self.betas = betas_for_alpha_bar(num_train_timesteps) @@ -242,7 +259,7 @@ def _get_variance(self, timestep, prev_timestep): return variance - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. @@ -259,13 +276,20 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None """ return sample - def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None): + def set_timesteps( + self, + num_inference_steps: int, + device: str | torch.device | None = None, + ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). Args: num_inference_steps (`int`): The number of diffusion steps used when generating samples with a pre-trained model. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None` (the default), the timesteps are not + moved. """ if num_inference_steps > self.config.num_train_timesteps: @@ -304,7 +328,27 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic self.timesteps = torch.from_numpy(timesteps).to(device) - def get_variables(self, alpha_prod_t, alpha_prod_t_prev, alpha_prod_t_back=None): + def get_variables( + self, + alpha_prod_t: torch.Tensor, + alpha_prod_t_prev: torch.Tensor, + alpha_prod_t_back: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Compute the variables used for DPM-Solver++ (2M) referencing the original implementation. + + Args: + alpha_prod_t (`torch.Tensor`): + The cumulative product of alphas at the current timestep. + alpha_prod_t_prev (`torch.Tensor`): + The cumulative product of alphas at the previous timestep. + alpha_prod_t_back (`torch.Tensor`, *optional*): + The cumulative product of alphas at the timestep before the previous timestep. + + Returns: + `tuple`: + A tuple containing the variables `h`, `r`, `lamb`, `lamb_next`. + """ lamb = ((alpha_prod_t / (1 - alpha_prod_t)) ** 0.5).log() lamb_next = ((alpha_prod_t_prev / (1 - alpha_prod_t_prev)) ** 0.5).log() h = lamb_next - lamb @@ -317,7 +361,33 @@ def get_variables(self, alpha_prod_t, alpha_prod_t_prev, alpha_prod_t_back=None) else: return h, None, lamb, lamb_next - def get_mult(self, h, r, alpha_prod_t, alpha_prod_t_prev, alpha_prod_t_back): + def get_mult( + self, + h: torch.Tensor, + r: torch.Tensor, + alpha_prod_t: torch.Tensor, + alpha_prod_t_prev: torch.Tensor, + alpha_prod_t_back: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Compute the multipliers for the previous sample and the predicted original sample. + + Args: + h (`torch.Tensor`): + The log-SNR difference. + r (`torch.Tensor`): + The ratio of log-SNR differences. + alpha_prod_t (`torch.Tensor`): + The cumulative product of alphas at the current timestep. + alpha_prod_t_prev (`torch.Tensor`): + The cumulative product of alphas at the previous timestep. + alpha_prod_t_back (`torch.Tensor`, *optional*): + The cumulative product of alphas at the timestep before the previous timestep. + + Returns: + `tuple`: + A tuple containing the multipliers. + """ mult1 = ((1 - alpha_prod_t_prev) / (1 - alpha_prod_t)) ** 0.5 * (-h).exp() mult2 = (-2 * h).expm1() * alpha_prod_t_prev**0.5 @@ -337,10 +407,10 @@ def step( sample: torch.Tensor, eta: float = 0.0, use_clipped_model_output: bool = False, - generator=None, - variance_noise: Optional[torch.Tensor] = None, + generator: torch.Generator | None = None, + variance_noise: torch.Tensor | None = None, return_dict: bool = False, - ) -> Union[DDIMSchedulerOutput, Tuple]: + ) -> DDIMSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). @@ -348,8 +418,12 @@ def step( Args: model_output (`torch.Tensor`): The direct output from learned diffusion model. - timestep (`float`): + old_pred_original_sample (`torch.Tensor`): + The predicted original sample from the previous timestep. + timestep (`int`): The current discrete timestep in the diffusion chain. + timestep_back (`int`): + The timestep to look back to. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. eta (`float`): @@ -429,7 +503,12 @@ def step( return prev_sample, pred_original_sample else: denoised_d = mult[2] * pred_original_sample - mult[3] * old_pred_original_sample - noise = randn_tensor(sample.shape, generator=generator, device=sample.device, dtype=sample.dtype) + noise = randn_tensor( + sample.shape, + generator=generator, + device=sample.device, + dtype=sample.dtype, + ) x_advanced = mult[0] * sample - mult[1] * denoised_d + mult_noise * noise prev_sample = x_advanced @@ -517,5 +596,5 @@ def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: tor velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample return velocity - def __len__(self): + def __len__(self) -> int: return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py index e7ba0ba1f30e..9c15df4569ca 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py @@ -15,7 +15,7 @@ # DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver import math -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -34,7 +34,7 @@ def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -48,8 +48,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -60,6 +60,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -210,8 +217,8 @@ def __init__( num_train_timesteps: int = 1000, beta_start: float = 0.0001, beta_end: float = 0.02, - beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, + beta_schedule: str = "linear", + trained_betas: np.ndarray | list[float] | None = None, solver_order: int = 2, prediction_type: Literal["epsilon", "sample", "v_prediction", "flow_prediction"] = "epsilon", thresholding: bool = False, @@ -221,15 +228,15 @@ def __init__( solver_type: Literal["midpoint", "heun"] = "midpoint", lower_order_final: bool = True, euler_at_final: bool = False, - use_karras_sigmas: Optional[bool] = False, - use_exponential_sigmas: Optional[bool] = False, - use_beta_sigmas: Optional[bool] = False, - use_lu_lambdas: Optional[bool] = False, - use_flow_sigmas: Optional[bool] = False, - flow_shift: Optional[float] = 1.0, - final_sigmas_type: Optional[Literal["zero", "sigma_min"]] = "zero", + use_karras_sigmas: bool = False, + use_exponential_sigmas: bool = False, + use_beta_sigmas: bool = False, + use_lu_lambdas: bool = False, + use_flow_sigmas: bool = False, + flow_shift: float = 1.0, + final_sigmas_type: Literal["zero", "sigma_min"] = "zero", lambda_min_clipped: float = -float("inf"), - variance_type: Optional[Literal["learned", "learned_range"]] = None, + variance_type: Literal["learned", "learned_range"] | None = None, timestep_spacing: Literal["linspace", "leading", "trailing"] = "linspace", steps_offset: int = 0, rescale_betas_zero_snr: bool = False, @@ -238,13 +245,26 @@ def __init__( ): if self.config.use_beta_sigmas and not is_scipy_available(): raise ImportError("Make sure to install scipy if you want to use beta sigmas.") - if sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1: + if ( + sum( + [ + self.config.use_beta_sigmas, + self.config.use_exponential_sigmas, + self.config.use_karras_sigmas, + ] + ) + > 1 + ): raise ValueError( "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." ) if algorithm_type in ["dpmsolver", "sde-dpmsolver"]: deprecation_message = f"algorithm_type {algorithm_type} is deprecated and will be removed in a future version. Choose from `dpmsolver++` or `sde-dpmsolver++` instead" - deprecate("algorithm_types dpmsolver and sde-dpmsolver", "1.0.0", deprecation_message) + deprecate( + "algorithm_types dpmsolver and sde-dpmsolver", + "1.0.0", + deprecation_message, + ) if trained_betas is not None: self.betas = torch.tensor(trained_betas, dtype=torch.float32) @@ -252,7 +272,15 @@ def __init__( self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32) elif beta_schedule == "scaled_linear": # this schedule is very specific to the latent diffusion model. - self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2 + self.betas = ( + torch.linspace( + beta_start**0.5, + beta_end**0.5, + num_train_timesteps, + dtype=torch.float32, + ) + ** 2 + ) elif beta_schedule == "squaredcos_cap_v2": # Glide cosine schedule self.betas = betas_for_alpha_bar(num_train_timesteps) @@ -280,7 +308,12 @@ def __init__( self.init_noise_sigma = 1.0 # settings for DPM-Solver - if algorithm_type not in ["dpmsolver", "dpmsolver++", "sde-dpmsolver", "sde-dpmsolver++"]: + if algorithm_type not in [ + "dpmsolver", + "dpmsolver++", + "sde-dpmsolver", + "sde-dpmsolver++", + ]: if algorithm_type == "deis": self.register_to_config(algorithm_type="dpmsolver++") else: @@ -333,11 +366,11 @@ def set_begin_index(self, begin_index: int = 0): def set_timesteps( self, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - mu: Optional[float] = None, - timesteps: Optional[List[int]] = None, - ) -> None: + num_inference_steps: int = None, + device: str | torch.device = None, + mu: float | None = None, + timesteps: list[int] | None = None, + ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -346,10 +379,7 @@ def set_timesteps( The number of diffusion steps used when generating samples with a pre-trained model. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - mu (`float`, *optional*): - The mu parameter for dynamic shifting. If provided, requires `use_dynamic_shifting=True` and - `time_shift_type="exponential"`. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to support arbitrary timesteps schedule. If `None`, timesteps will be generated based on the `timestep_spacing` attribute. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`, and `timestep_spacing` attribute will be ignored. @@ -544,7 +574,7 @@ def _sigma_to_t(self, sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray: t = t.reshape(sigma.shape) return t - def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """ Convert sigma values to alpha_t and sigma_t values. @@ -553,7 +583,7 @@ def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, to The sigma value(s) to convert. Returns: - `Tuple[torch.Tensor, torch.Tensor]`: + `tuple[torch.Tensor, torch.Tensor]`: A tuple containing (alpha_t, sigma_t) values. """ if self.config.use_flow_sigmas: @@ -717,7 +747,7 @@ def convert_model_output( self, model_output: torch.Tensor, *args, - sample: torch.Tensor = None, + sample: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: """ @@ -731,7 +761,7 @@ def convert_model_output( Args: model_output (`torch.Tensor`): The direct output from the learned diffusion model. - sample (`torch.Tensor`): + sample (`torch.Tensor`, *optional*): A current instance of a sample created by the diffusion process. Returns: @@ -815,8 +845,8 @@ def dpm_solver_first_order_update( self, model_output: torch.Tensor, *args, - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, + sample: torch.Tensor | None = None, + noise: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: """ @@ -825,8 +855,10 @@ def dpm_solver_first_order_update( Args: model_output (`torch.Tensor`): The direct output from the learned diffusion model. - sample (`torch.Tensor`): + sample (`torch.Tensor`, *optional*): A current instance of a sample created by the diffusion process. + noise (`torch.Tensor`, *optional*): + The noise tensor. Returns: `torch.Tensor`: @@ -853,7 +885,10 @@ def dpm_solver_first_order_update( "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", ) - sigma_t, sigma_s = self.sigmas[self.step_index + 1], self.sigmas[self.step_index] + sigma_t, sigma_s = ( + self.sigmas[self.step_index + 1], + self.sigmas[self.step_index], + ) alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) alpha_s, sigma_s = self._sigma_to_alpha_sigma_t(sigma_s) lambda_t = torch.log(alpha_t) - torch.log(sigma_t) @@ -882,19 +917,19 @@ def dpm_solver_first_order_update( def multistep_dpm_solver_second_order_update( self, - model_output_list: List[torch.Tensor], + model_output_list: list[torch.Tensor], *args, - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, + sample: torch.Tensor | None = None, + noise: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: """ One step for the second-order multistep DPMSolver. Args: - model_output_list (`List[torch.Tensor]`): + model_output_list (`list[torch.Tensor]`): The direct outputs from learned diffusion model at current and latter timesteps. - sample (`torch.Tensor`): + sample (`torch.Tensor`, *optional*): A current instance of a sample created by the diffusion process. Returns: @@ -1005,20 +1040,22 @@ def multistep_dpm_solver_second_order_update( def multistep_dpm_solver_third_order_update( self, - model_output_list: List[torch.Tensor], + model_output_list: list[torch.Tensor], *args, - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, + sample: torch.Tensor | None = None, + noise: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: """ One step for the third-order multistep DPMSolver. Args: - model_output_list (`List[torch.Tensor]`): + model_output_list (`list[torch.Tensor]`): The direct outputs from learned diffusion model at current and latter timesteps. - sample (`torch.Tensor`): + sample (`torch.Tensor`, *optional*): A current instance of a sample created by diffusion process. + noise (`torch.Tensor`, *optional*): + The noise tensor. Returns: `torch.Tensor`: @@ -1099,7 +1136,9 @@ def multistep_dpm_solver_third_order_update( return x_t def index_for_timestep( - self, timestep: Union[int, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None + self, + timestep: int | torch.Tensor, + schedule_timesteps: torch.Tensor | None = None, ) -> int: """ Find the index for a given timestep in the schedule. @@ -1132,7 +1171,7 @@ def index_for_timestep( return step_index - def _init_step_index(self, timestep: Union[int, torch.Tensor]) -> None: + def _init_step_index(self, timestep: int | torch.Tensor) -> None: """ Initialize the step_index counter for the scheduler. @@ -1151,12 +1190,12 @@ def _init_step_index(self, timestep: Union[int, torch.Tensor]) -> None: def step( self, model_output: torch.Tensor, - timestep: Union[int, torch.Tensor], + timestep: int | torch.Tensor, sample: torch.Tensor, - generator: Optional[torch.Generator] = None, - variance_noise: Optional[torch.Tensor] = None, + generator: torch.Generator | None = None, + variance_noise: torch.Tensor | None = None, return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: + ) -> SchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the multistep DPMSolver. @@ -1209,7 +1248,10 @@ def step( sample = sample.to(torch.float32) if self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"] and variance_noise is None: noise = randn_tensor( - model_output.shape, generator=generator, device=model_output.device, dtype=torch.float32 + model_output.shape, + generator=generator, + device=model_output.device, + dtype=torch.float32, ) elif self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"]: noise = variance_noise.to(device=model_output.device, dtype=torch.float32) diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py index 71b9960bf2ff..dd579e84d609 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py @@ -15,13 +15,13 @@ # DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver from dataclasses import dataclass -from typing import List, Optional, Tuple, Union import flax import jax import jax.numpy as jnp from ..configuration_utils import ConfigMixin, register_to_config +from ..utils import logging from .scheduling_utils_flax import ( CommonSchedulerState, FlaxKarrasDiffusionSchedulers, @@ -31,6 +31,9 @@ ) +logger = logging.get_logger(__name__) + + @flax.struct.dataclass class DPMSolverMultistepSchedulerState: common: CommonSchedulerState @@ -41,13 +44,13 @@ class DPMSolverMultistepSchedulerState: # setable values init_noise_sigma: jnp.ndarray timesteps: jnp.ndarray - num_inference_steps: Optional[int] = None + num_inference_steps: int = None # running values - model_outputs: Optional[jnp.ndarray] = None - lower_order_nums: Optional[jnp.int32] = None - prev_timestep: Optional[jnp.int32] = None - cur_sample: Optional[jnp.ndarray] = None + model_outputs: jnp.ndarray | None = None + lower_order_nums: jnp.int32 | None = None + prev_timestep: jnp.int32 | None = None + cur_sample: jnp.ndarray | None = None @classmethod def create( @@ -159,7 +162,7 @@ def __init__( beta_start: float = 0.0001, beta_end: float = 0.02, beta_schedule: str = "linear", - trained_betas: Optional[jnp.ndarray] = None, + trained_betas: jnp.ndarray | None = None, solver_order: int = 2, prediction_type: str = "epsilon", thresholding: bool = False, @@ -171,9 +174,13 @@ def __init__( timestep_spacing: str = "linspace", dtype: jnp.dtype = jnp.float32, ): + logger.warning( + "Flax classes are deprecated and will be removed in Diffusers v1.0.0. We " + "recommend migrating to PyTorch classes or pinning your version of Diffusers." + ) self.dtype = dtype - def create_state(self, common: Optional[CommonSchedulerState] = None) -> DPMSolverMultistepSchedulerState: + def create_state(self, common: CommonSchedulerState | None = None) -> DPMSolverMultistepSchedulerState: if common is None: common = CommonSchedulerState.create(self) @@ -203,7 +210,10 @@ def create_state(self, common: Optional[CommonSchedulerState] = None) -> DPMSolv ) def set_timesteps( - self, state: DPMSolverMultistepSchedulerState, num_inference_steps: int, shape: Tuple + self, + state: DPMSolverMultistepSchedulerState, + num_inference_steps: int, + shape: tuple, ) -> DPMSolverMultistepSchedulerState: """ Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference. @@ -213,7 +223,7 @@ def set_timesteps( the `FlaxDPMSolverMultistepScheduler` state data class instance. num_inference_steps (`int`): the number of diffusion steps used when generating samples with a pre-trained model. - shape (`Tuple`): + shape (`tuple`): the shape of the samples to be generated. """ last_timestep = self.config.num_train_timesteps @@ -301,10 +311,13 @@ def convert_model_output( if self.config.thresholding: # Dynamic thresholding in https://huggingface.co/papers/2205.11487 dynamic_max_val = jnp.percentile( - jnp.abs(x0_pred), self.config.dynamic_thresholding_ratio, axis=tuple(range(1, x0_pred.ndim)) + jnp.abs(x0_pred), + self.config.dynamic_thresholding_ratio, + axis=tuple(range(1, x0_pred.ndim)), ) dynamic_max_val = jnp.maximum( - dynamic_max_val, self.config.sample_max_value * jnp.ones_like(dynamic_max_val) + dynamic_max_val, + self.config.sample_max_value * jnp.ones_like(dynamic_max_val), ) x0_pred = jnp.clip(x0_pred, -dynamic_max_val, dynamic_max_val) / dynamic_max_val return x0_pred @@ -365,7 +378,7 @@ def multistep_dpm_solver_second_order_update( self, state: DPMSolverMultistepSchedulerState, model_output_list: jnp.ndarray, - timestep_list: List[int], + timestep_list: list[int], prev_timestep: int, sample: jnp.ndarray, ) -> jnp.ndarray: @@ -373,7 +386,7 @@ def multistep_dpm_solver_second_order_update( One step for the second-order multistep DPM-Solver. Args: - model_output_list (`List[jnp.ndarray]`): + model_output_list (`list[jnp.ndarray]`): direct outputs from learned diffusion model at current and latter timesteps. timestep (`int`): current and latter discrete timestep in the diffusion chain. prev_timestep (`int`): previous discrete timestep in the diffusion chain. @@ -385,7 +398,11 @@ def multistep_dpm_solver_second_order_update( """ t, s0, s1 = prev_timestep, timestep_list[-1], timestep_list[-2] m0, m1 = model_output_list[-1], model_output_list[-2] - lambda_t, lambda_s0, lambda_s1 = state.lambda_t[t], state.lambda_t[s0], state.lambda_t[s1] + lambda_t, lambda_s0, lambda_s1 = ( + state.lambda_t[t], + state.lambda_t[s0], + state.lambda_t[s1], + ) alpha_t, alpha_s0 = state.alpha_t[t], state.alpha_t[s0] sigma_t, sigma_s0 = state.sigma_t[t], state.sigma_t[s0] h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1 @@ -425,7 +442,7 @@ def multistep_dpm_solver_third_order_update( self, state: DPMSolverMultistepSchedulerState, model_output_list: jnp.ndarray, - timestep_list: List[int], + timestep_list: list[int], prev_timestep: int, sample: jnp.ndarray, ) -> jnp.ndarray: @@ -433,7 +450,7 @@ def multistep_dpm_solver_third_order_update( One step for the third-order multistep DPM-Solver. Args: - model_output_list (`List[jnp.ndarray]`): + model_output_list (`list[jnp.ndarray]`): direct outputs from learned diffusion model at current and latter timesteps. timestep (`int`): current and latter discrete timestep in the diffusion chain. prev_timestep (`int`): previous discrete timestep in the diffusion chain. @@ -443,7 +460,12 @@ def multistep_dpm_solver_third_order_update( Returns: `jnp.ndarray`: the sample tensor at the previous timestep. """ - t, s0, s1, s2 = prev_timestep, timestep_list[-1], timestep_list[-2], timestep_list[-3] + t, s0, s1, s2 = ( + prev_timestep, + timestep_list[-1], + timestep_list[-2], + timestep_list[-3], + ) m0, m1, m2 = model_output_list[-1], model_output_list[-2], model_output_list[-3] lambda_t, lambda_s0, lambda_s1, lambda_s2 = ( state.lambda_t[t], @@ -484,7 +506,7 @@ def step( timestep: int, sample: jnp.ndarray, return_dict: bool = True, - ) -> Union[FlaxDPMSolverMultistepSchedulerOutput, Tuple]: + ) -> FlaxDPMSolverMultistepSchedulerOutput | tuple: """ Predict the sample at the previous timestep by DPM-Solver. Core function to propagate the diffusion process from the learned model outputs (most often the predicted noise). @@ -615,7 +637,10 @@ def step_3(state: DPMSolverMultistepSchedulerState) -> jnp.ndarray: return FlaxDPMSolverMultistepSchedulerOutput(prev_sample=prev_sample, state=state) def scale_model_input( - self, state: DPMSolverMultistepSchedulerState, sample: jnp.ndarray, timestep: Optional[int] = None + self, + state: DPMSolverMultistepSchedulerState, + sample: jnp.ndarray, + timestep: int | None = None, ) -> jnp.ndarray: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py b/src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py index 6696b0375f9f..3386514f64e9 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py @@ -15,7 +15,7 @@ # DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver import math -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -34,7 +34,7 @@ def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -48,8 +48,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -60,6 +60,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -134,6 +141,10 @@ class DPMSolverMultistepInverseScheduler(SchedulerMixin, ConfigMixin): use_beta_sigmas (`bool`, *optional*, defaults to `False`): Whether to use beta sigmas for step sizes in the noise schedule during the sampling process. Refer to [Beta Sampling is All You Need](https://huggingface.co/papers/2407.12173) for more information. + use_flow_sigmas (`bool`, *optional*, defaults to `False`): + Whether to use flow sigmas for step sizes in the noise schedule during the sampling process. + flow_shift (`float`, *optional*, defaults to 1.0): + The flow shift factor. Valid only when `use_flow_sigmas=True`. lambda_min_clipped (`float`, defaults to `-inf`): Clipping threshold for the minimum value of `lambda(t)` for numerical stability. This is critical for the cosine (`squaredcos_cap_v2`) noise schedule. @@ -156,36 +167,49 @@ def __init__( num_train_timesteps: int = 1000, beta_start: float = 0.0001, beta_end: float = 0.02, - beta_schedule: str = "linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, + beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear", + trained_betas: np.ndarray | list[float] | None = None, solver_order: int = 2, - prediction_type: str = "epsilon", + prediction_type: Literal["epsilon", "sample", "v_prediction", "flow_prediction"] = "epsilon", thresholding: bool = False, dynamic_thresholding_ratio: float = 0.995, sample_max_value: float = 1.0, - algorithm_type: str = "dpmsolver++", - solver_type: str = "midpoint", + algorithm_type: Literal["dpmsolver", "dpmsolver++", "sde-dpmsolver", "sde-dpmsolver++"] = "dpmsolver++", + solver_type: Literal["midpoint", "heun"] = "midpoint", lower_order_final: bool = True, euler_at_final: bool = False, - use_karras_sigmas: Optional[bool] = False, - use_exponential_sigmas: Optional[bool] = False, - use_beta_sigmas: Optional[bool] = False, - use_flow_sigmas: Optional[bool] = False, - flow_shift: Optional[float] = 1.0, + use_karras_sigmas: bool = False, + use_exponential_sigmas: bool = False, + use_beta_sigmas: bool = False, + use_flow_sigmas: bool = False, + flow_shift: float = 1.0, lambda_min_clipped: float = -float("inf"), - variance_type: Optional[str] = None, - timestep_spacing: str = "linspace", + variance_type: Literal["learned", "learned_range"] | None = None, + timestep_spacing: Literal["linspace", "leading", "trailing"] = "linspace", steps_offset: int = 0, ): if self.config.use_beta_sigmas and not is_scipy_available(): raise ImportError("Make sure to install scipy if you want to use beta sigmas.") - if sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1: + if ( + sum( + [ + self.config.use_beta_sigmas, + self.config.use_exponential_sigmas, + self.config.use_karras_sigmas, + ] + ) + > 1 + ): raise ValueError( "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." ) if algorithm_type in ["dpmsolver", "sde-dpmsolver"]: deprecation_message = f"algorithm_type {algorithm_type} is deprecated and will be removed in a future version. Choose from `dpmsolver++` or `sde-dpmsolver++` instead" - deprecate("algorithm_types dpmsolver and sde-dpmsolver", "1.0.0", deprecation_message) + deprecate( + "algorithm_types dpmsolver and sde-dpmsolver", + "1.0.0", + deprecation_message, + ) if trained_betas is not None: self.betas = torch.tensor(trained_betas, dtype=torch.float32) @@ -193,7 +217,15 @@ def __init__( self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32) elif beta_schedule == "scaled_linear": # this schedule is very specific to the latent diffusion model. - self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2 + self.betas = ( + torch.linspace( + beta_start**0.5, + beta_end**0.5, + num_train_timesteps, + dtype=torch.float32, + ) + ** 2 + ) elif beta_schedule == "squaredcos_cap_v2": # Glide cosine schedule self.betas = betas_for_alpha_bar(num_train_timesteps) @@ -212,7 +244,12 @@ def __init__( self.init_noise_sigma = 1.0 # settings for DPM-Solver - if algorithm_type not in ["dpmsolver", "dpmsolver++", "sde-dpmsolver", "sde-dpmsolver++"]: + if algorithm_type not in [ + "dpmsolver", + "dpmsolver++", + "sde-dpmsolver", + "sde-dpmsolver++", + ]: if algorithm_type == "deis": self.register_to_config(algorithm_type="dpmsolver++") else: @@ -243,7 +280,11 @@ def step_index(self): """ return self._step_index - def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torch.device] = None): + def set_timesteps( + self, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -375,7 +416,7 @@ def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: return sample # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t - def _sigma_to_t(self, sigma, log_sigmas): + def _sigma_to_t(self, sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray: """ Convert sigma values to corresponding timestep values through interpolation. @@ -412,7 +453,7 @@ def _sigma_to_t(self, sigma, log_sigmas): return t # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._sigma_to_alpha_sigma_t - def _sigma_to_alpha_sigma_t(self, sigma): + def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """ Convert sigma values to alpha_t and sigma_t values. @@ -421,7 +462,7 @@ def _sigma_to_alpha_sigma_t(self, sigma): The sigma value(s) to convert. Returns: - `Tuple[torch.Tensor, torch.Tensor]`: + `tuple[torch.Tensor, torch.Tensor]`: A tuple containing (alpha_t, sigma_t) values. """ if self.config.use_flow_sigmas: @@ -434,7 +475,7 @@ def _sigma_to_alpha_sigma_t(self, sigma): return alpha_t, sigma_t # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras - def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor: + def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: """ Construct the noise schedule as proposed in [Elucidating the Design Space of Diffusion-Based Generative Models](https://huggingface.co/papers/2206.00364). @@ -560,7 +601,7 @@ def convert_model_output( self, model_output: torch.Tensor, *args, - sample: torch.Tensor = None, + sample: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: """ @@ -574,7 +615,7 @@ def convert_model_output( Args: model_output (`torch.Tensor`): The direct output from the learned diffusion model. - sample (`torch.Tensor`): + sample (`torch.Tensor`, *optional*): A current instance of a sample created by the diffusion process. Returns: @@ -659,8 +700,8 @@ def dpm_solver_first_order_update( self, model_output: torch.Tensor, *args, - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, + sample: torch.Tensor | None = None, + noise: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: """ @@ -669,8 +710,10 @@ def dpm_solver_first_order_update( Args: model_output (`torch.Tensor`): The direct output from the learned diffusion model. - sample (`torch.Tensor`): + sample (`torch.Tensor`, *optional*): A current instance of a sample created by the diffusion process. + noise (`torch.Tensor`, *optional*): + The noise tensor. Returns: `torch.Tensor`: @@ -697,7 +740,10 @@ def dpm_solver_first_order_update( "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", ) - sigma_t, sigma_s = self.sigmas[self.step_index + 1], self.sigmas[self.step_index] + sigma_t, sigma_s = ( + self.sigmas[self.step_index + 1], + self.sigmas[self.step_index], + ) alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) alpha_s, sigma_s = self._sigma_to_alpha_sigma_t(sigma_s) lambda_t = torch.log(alpha_t) - torch.log(sigma_t) @@ -727,19 +773,19 @@ def dpm_solver_first_order_update( # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.multistep_dpm_solver_second_order_update def multistep_dpm_solver_second_order_update( self, - model_output_list: List[torch.Tensor], + model_output_list: list[torch.Tensor], *args, - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, + sample: torch.Tensor | None = None, + noise: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: """ One step for the second-order multistep DPMSolver. Args: - model_output_list (`List[torch.Tensor]`): + model_output_list (`list[torch.Tensor]`): The direct outputs from learned diffusion model at current and latter timesteps. - sample (`torch.Tensor`): + sample (`torch.Tensor`, *optional*): A current instance of a sample created by the diffusion process. Returns: @@ -851,20 +897,22 @@ def multistep_dpm_solver_second_order_update( # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.multistep_dpm_solver_third_order_update def multistep_dpm_solver_third_order_update( self, - model_output_list: List[torch.Tensor], + model_output_list: list[torch.Tensor], *args, - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, + sample: torch.Tensor | None = None, + noise: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: """ One step for the third-order multistep DPMSolver. Args: - model_output_list (`List[torch.Tensor]`): + model_output_list (`list[torch.Tensor]`): The direct outputs from learned diffusion model at current and latter timesteps. - sample (`torch.Tensor`): + sample (`torch.Tensor`, *optional*): A current instance of a sample created by diffusion process. + noise (`torch.Tensor`, *optional*): + The noise tensor. Returns: `torch.Tensor`: @@ -944,7 +992,7 @@ def multistep_dpm_solver_third_order_update( ) return x_t - def _init_step_index(self, timestep): + def _init_step_index(self, timestep: int | torch.Tensor): if isinstance(timestep, torch.Tensor): timestep = timestep.to(self.timesteps.device) @@ -966,12 +1014,12 @@ def _init_step_index(self, timestep): def step( self, model_output: torch.Tensor, - timestep: Union[int, torch.Tensor], + timestep: int | torch.Tensor, sample: torch.Tensor, - generator=None, - variance_noise: Optional[torch.Tensor] = None, + generator: torch.Generator | None = None, + variance_noise: torch.Tensor | None = None, return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: + ) -> SchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the multistep DPMSolver. @@ -1020,7 +1068,10 @@ def step( if self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"] and variance_noise is None: noise = randn_tensor( - model_output.shape, generator=generator, device=model_output.device, dtype=model_output.dtype + model_output.shape, + generator=generator, + device=model_output.device, + dtype=model_output.dtype, ) elif self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"]: noise = variance_noise @@ -1067,6 +1118,21 @@ def add_noise( noise: torch.Tensor, timesteps: torch.IntTensor, ) -> torch.Tensor: + """ + Add noise to the clean `original_samples` using the scheduler's equivalent function. + + Args: + original_samples (`torch.Tensor`): + The original samples to add noise to. + noise (`torch.Tensor`): + The noise tensor. + timesteps (`torch.IntTensor`): + The timesteps at which to add noise. + + Returns: + `torch.Tensor`: + The noisy samples. + """ # Make sure sigmas and timesteps have the same device and dtype as original_samples sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype) if original_samples.device.type == "mps" and torch.is_floating_point(timesteps): @@ -1096,5 +1162,5 @@ def add_noise( noisy_samples = alpha_t * original_samples + sigma_t * noise return noisy_samples - def __len__(self): + def __len__(self) -> int: return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_sde.py b/src/diffusers/schedulers/scheduling_dpmsolver_sde.py index 81c9e4134f57..b288d9f61b14 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_sde.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_sde.py @@ -14,7 +14,7 @@ import math from dataclasses import dataclass -from typing import List, Literal, Optional, Tuple, Union +from typing import Callable, List, Literal, Optional, Tuple, Union import numpy as np import torch @@ -45,13 +45,20 @@ class DPMSolverSDESchedulerOutput(BaseOutput): """ prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None + pred_original_sample: torch.Tensor | None = None class BatchedBrownianTree: """A wrapper around torchsde.BrownianTree that enables batches of entropy.""" - def __init__(self, x, t0, t1, seed=None, **kwargs): + def __init__( + self, + x: torch.Tensor, + t0: float, + t1: float, + seed: Optional[Union[int, List[int]]] = None, + **kwargs, + ): t0, t1, self.sign = self.sort(t0, t1) w0 = kwargs.get("w0", torch.zeros_like(x)) if seed is None: @@ -79,10 +86,23 @@ def __init__(self, x, t0, t1, seed=None, **kwargs): ] @staticmethod - def sort(a, b): - return (a, b, 1) if a < b else (b, a, -1) + def sort(a: float, b: float) -> Tuple[float, float, float]: + """ + Sorts two float values and returns them along with a sign indicating if they were swapped. + + Args: + a (`float`): + The first value. + b (`float`): + The second value. - def __call__(self, t0, t1): + Returns: + `Tuple[float, float, float]`: + A tuple containing the sorted values (min, max) and a sign (1.0 if a < b, -1.0 otherwise). + """ + return (a, b, 1.0) if a < b else (b, a, -1.0) + + def __call__(self, t0: float, t1: float) -> torch.Tensor: t0, t1, sign = self.sort(t0, t1) w = torch.stack([tree(t0, t1) for tree in self.trees]) * (self.sign * sign) return w if self.batched else w[0] @@ -92,23 +112,29 @@ class BrownianTreeNoiseSampler: """A noise sampler backed by a torchsde.BrownianTree. Args: - x (Tensor): The tensor whose shape, device and dtype to use to generate - random samples. - sigma_min (float): The low end of the valid interval. - sigma_max (float): The high end of the valid interval. - seed (int or List[int]): The random seed. If a list of seeds is + x (`torch.Tensor`): The tensor whose shape, device and dtype is used to generate random samples. + sigma_min (`float`): The low end of the valid interval. + sigma_max (`float`): The high end of the valid interval. + seed (`int` or `List[int]`): The random seed. If a list of seeds is supplied instead of a single integer, then the noise sampler will use one BrownianTree per batch item, each with its own seed. - transform (callable): A function that maps sigma to the sampler's + transform (`callable`): A function that maps sigma to the sampler's internal timestep. """ - def __init__(self, x, sigma_min, sigma_max, seed=None, transform=lambda x: x): + def __init__( + self, + x: torch.Tensor, + sigma_min: float, + sigma_max: float, + seed: Optional[Union[int, List[int]]] = None, + transform: Callable[[float], float] = lambda x: x, + ): self.transform = transform t0, t1 = self.transform(torch.as_tensor(sigma_min)), self.transform(torch.as_tensor(sigma_max)) self.tree = BatchedBrownianTree(x, t0, t1, seed) - def __call__(self, sigma, sigma_next): + def __call__(self, sigma: float, sigma_next: float) -> torch.Tensor: t0, t1 = self.transform(torch.as_tensor(sigma)), self.transform(torch.as_tensor(sigma_next)) return self.tree(t0, t1) / (t1 - t0).abs().sqrt() @@ -117,7 +143,7 @@ def __call__(self, sigma, sigma_next): def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -131,8 +157,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -143,6 +169,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -209,19 +242,28 @@ def __init__( num_train_timesteps: int = 1000, beta_start: float = 0.00085, # sensible defaults beta_end: float = 0.012, - beta_schedule: str = "linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, - prediction_type: str = "epsilon", - use_karras_sigmas: Optional[bool] = False, - use_exponential_sigmas: Optional[bool] = False, - use_beta_sigmas: Optional[bool] = False, - noise_sampler_seed: Optional[int] = None, - timestep_spacing: str = "linspace", + beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear", + trained_betas: np.ndarray | list[float] | None = None, + prediction_type: Literal["epsilon", "sample", "v_prediction"] = "epsilon", + use_karras_sigmas: bool = False, + use_exponential_sigmas: bool = False, + use_beta_sigmas: bool = False, + noise_sampler_seed: int | None = None, + timestep_spacing: Literal["linspace", "leading", "trailing"] = "linspace", steps_offset: int = 0, ): if self.config.use_beta_sigmas and not is_scipy_available(): raise ImportError("Make sure to install scipy if you want to use beta sigmas.") - if sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1: + if ( + sum( + [ + self.config.use_beta_sigmas, + self.config.use_exponential_sigmas, + self.config.use_karras_sigmas, + ] + ) + > 1 + ): raise ValueError( "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." ) @@ -231,7 +273,15 @@ def __init__( self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32) elif beta_schedule == "scaled_linear": # this schedule is very specific to the latent diffusion model. - self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2 + self.betas = ( + torch.linspace( + beta_start**0.5, + beta_end**0.5, + num_train_timesteps, + dtype=torch.float32, + ) + ** 2 + ) elif beta_schedule == "squaredcos_cap_v2": # Glide cosine schedule self.betas = betas_for_alpha_bar(num_train_timesteps) @@ -252,7 +302,7 @@ def __init__( # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep def index_for_timestep( - self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None + self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None ) -> int: """ Find the index of a given timestep in the timestep schedule. @@ -282,7 +332,7 @@ def index_for_timestep( return indices[pos].item() # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: + def _init_step_index(self, timestep: float | torch.Tensor) -> None: """ Initialize the step index for the scheduler based on the given timestep. @@ -298,7 +348,7 @@ def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: self._step_index = self._begin_index @property - def init_noise_sigma(self): + def init_noise_sigma(self) -> torch.Tensor: # standard deviation of the initial noise distribution if self.config.timestep_spacing in ["linspace", "trailing"]: return self.sigmas.max() @@ -306,21 +356,21 @@ def init_noise_sigma(self): return (self.sigmas.max() ** 2 + 1) ** 0.5 @property - def step_index(self): + def step_index(self) -> Union[int, None]: """ The index counter for current timestep. It will increase 1 after each scheduler step. """ return self._step_index @property - def begin_index(self): + def begin_index(self) -> Union[int, None]: """ The index for the first timestep. It should be set from pipeline with `set_begin_index` method. """ return self._begin_index # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index - def set_begin_index(self, begin_index: int = 0): + def set_begin_index(self, begin_index: int = 0) -> None: """ Sets the begin index for the scheduler. This function should be run from pipeline before the inference. @@ -333,7 +383,7 @@ def set_begin_index(self, begin_index: int = 0): def scale_model_input( self, sample: torch.Tensor, - timestep: Union[float, torch.Tensor], + timestep: float | torch.Tensor, ) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the @@ -360,9 +410,9 @@ def scale_model_input( def set_timesteps( self, num_inference_steps: int, - device: Union[str, torch.device] = None, - num_train_timesteps: Optional[int] = None, - ): + device: str | torch.device = None, + num_train_timesteps: int | None = None, + ) -> None: """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -371,6 +421,8 @@ def set_timesteps( The number of diffusion steps used when generating samples with a pre-trained model. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + num_train_timesteps (`int`, *optional*): + The number of train timesteps. If `None`, uses `self.config.num_train_timesteps`. """ self.num_inference_steps = num_inference_steps @@ -436,7 +488,7 @@ def set_timesteps( self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication self.noise_sampler = None - def _second_order_timesteps(self, sigmas, log_sigmas): + def _second_order_timesteps(self, sigmas: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray: def sigma_fn(_t): return np.exp(-_t) @@ -452,7 +504,7 @@ def t_fn(_sigma): return timesteps # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t - def _sigma_to_t(self, sigma, log_sigmas): + def _sigma_to_t(self, sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray: """ Convert sigma values to corresponding timestep values through interpolation. @@ -597,27 +649,27 @@ def _convert_to_beta( return sigmas @property - def state_in_first_order(self): + def state_in_first_order(self) -> bool: return self.sample is None def step( self, - model_output: Union[torch.Tensor, np.ndarray], - timestep: Union[float, torch.Tensor], - sample: Union[torch.Tensor, np.ndarray], + model_output: torch.Tensor, + timestep: float | torch.Tensor, + sample: torch.Tensor, return_dict: bool = True, s_noise: float = 1.0, - ) -> Union[DPMSolverSDESchedulerOutput, Tuple]: + ) -> DPMSolverSDESchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). Args: - model_output (`torch.Tensor` or `np.ndarray`): + model_output (`torch.Tensor`): The direct output from learned diffusion model. timestep (`float` or `torch.Tensor`): The current discrete timestep in the diffusion chain. - sample (`torch.Tensor` or `np.ndarray`): + sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. return_dict (`bool`): Whether or not to return a [`~schedulers.scheduling_dpmsolver_sde.DPMSolverSDESchedulerOutput`] or @@ -636,7 +688,9 @@ def step( # Create a noise sampler if it hasn't been created yet if self.noise_sampler is None: min_sigma, max_sigma = self.sigmas[self.sigmas > 0].min(), self.sigmas.max() - self.noise_sampler = BrownianTreeNoiseSampler(sample, min_sigma, max_sigma, self.noise_sampler_seed) + self.noise_sampler = BrownianTreeNoiseSampler( + sample, min_sigma.item(), max_sigma.item(), self.noise_sampler_seed + ) # Define functions to compute sigma and t from each other def sigma_fn(_t: torch.Tensor) -> torch.Tensor: @@ -687,7 +741,10 @@ def t_fn(_sigma: torch.Tensor) -> torch.Tensor: sigma_from = sigma_fn(t) sigma_to = sigma_fn(t_next) - sigma_up = min(sigma_to, (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5) + sigma_up = min( + sigma_to, + (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5, + ) sigma_down = (sigma_to**2 - sigma_up**2) ** 0.5 ancestral_t = t_fn(sigma_down) prev_sample = (sigma_fn(ancestral_t) / sigma_fn(t)) * sample - ( @@ -764,5 +821,5 @@ def add_noise( noisy_samples = original_samples + noise * sigma return noisy_samples - def __len__(self): + def __len__(self) -> int: return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py b/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py index 4916e1abb549..002f2b844a52 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py @@ -15,7 +15,7 @@ # DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver import math -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -36,7 +36,7 @@ def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -50,8 +50,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -62,6 +62,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -95,7 +102,7 @@ class DPMSolverSinglestepScheduler(SchedulerMixin, ConfigMixin): beta_schedule (`"linear"`, `"scaled_linear"`, or `"squaredcos_cap_v2"`, defaults to `"linear"`): The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from `linear`, `scaled_linear`, or `squaredcos_cap_v2`. - trained_betas (`np.ndarray` or `List[float]`, *optional*): + trained_betas (`np.ndarray` or `list[float]`, *optional*): Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. solver_order (`int`, defaults to `2`): The DPMSolver order which can be `1` or `2` or `3`. It is recommended to use `solver_order=2` for guided @@ -161,7 +168,7 @@ def __init__( beta_start: float = 0.0001, beta_end: float = 0.02, beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, + trained_betas: np.ndarray | list[float] | None = None, solver_order: int = 2, prediction_type: Literal["epsilon", "sample", "v_prediction", "flow_prediction"] = "epsilon", thresholding: bool = False, @@ -170,14 +177,14 @@ def __init__( algorithm_type: Literal["dpmsolver", "dpmsolver++", "sde-dpmsolver++"] = "dpmsolver++", solver_type: Literal["midpoint", "heun"] = "midpoint", lower_order_final: bool = False, - use_karras_sigmas: Optional[bool] = False, - use_exponential_sigmas: Optional[bool] = False, - use_beta_sigmas: Optional[bool] = False, - use_flow_sigmas: Optional[bool] = False, - flow_shift: Optional[float] = 1.0, - final_sigmas_type: Optional[Literal["zero", "sigma_min"]] = "zero", + use_karras_sigmas: bool = False, + use_exponential_sigmas: bool = False, + use_beta_sigmas: bool = False, + use_flow_sigmas: bool = False, + flow_shift: float = 1.0, + final_sigmas_type: Literal["zero", "sigma_min"] = "zero", lambda_min_clipped: float = -float("inf"), - variance_type: Optional[Literal["learned", "learned_range"]] = None, + variance_type: Literal["learned", "learned_range"] | None = None, use_dynamic_shifting: bool = False, time_shift_type: Literal["exponential"] = "exponential", ) -> None: @@ -243,7 +250,7 @@ def __init__( self._begin_index = None self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication - def get_order_list(self, num_inference_steps: int) -> List[int]: + def get_order_list(self, num_inference_steps: int) -> list[int]: """ Computes the solver order at each time step. @@ -252,7 +259,7 @@ def get_order_list(self, num_inference_steps: int) -> List[int]: The number of diffusion steps used when generating samples with a pre-trained model. Returns: - `List[int]`: + `list[int]`: The list of solver orders for each timestep. """ steps = num_inference_steps @@ -288,7 +295,7 @@ def get_order_list(self, num_inference_steps: int) -> List[int]: return orders @property - def step_index(self) -> Optional[int]: + def step_index(self) -> int: """ The index counter for current timestep. It will increase 1 after each scheduler step. @@ -299,7 +306,7 @@ def step_index(self) -> Optional[int]: return self._step_index @property - def begin_index(self) -> Optional[int]: + def begin_index(self) -> int: """ The index for the first timestep. It should be set from pipeline with `set_begin_index` method. @@ -322,11 +329,11 @@ def set_begin_index(self, begin_index: int = 0) -> None: def set_timesteps( self, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - mu: Optional[float] = None, - timesteps: Optional[List[int]] = None, - ) -> None: + num_inference_steps: int = None, + device: str | torch.device = None, + mu: float | None = None, + timesteps: list[int] | None = None, + ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -335,9 +342,7 @@ def set_timesteps( The number of diffusion steps used when generating samples with a pre-trained model. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - mu (`float`, *optional*): - The mu parameter for dynamic shifting. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default timestep spacing strategy of equal spacing between timesteps schedule is used. If `timesteps` is passed, `num_inference_steps` must be `None`. @@ -512,7 +517,7 @@ def _sigma_to_t(self, sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray: return t # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._sigma_to_alpha_sigma_t - def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """ Convert sigma values to alpha_t and sigma_t values. @@ -521,7 +526,7 @@ def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, to The sigma value(s) to convert. Returns: - `Tuple[torch.Tensor, torch.Tensor]`: + `tuple[torch.Tensor, torch.Tensor]`: A tuple containing (alpha_t, sigma_t) values. """ if self.config.use_flow_sigmas: @@ -659,7 +664,7 @@ def convert_model_output( self, model_output: torch.Tensor, *args, - sample: Optional[torch.Tensor] = None, + sample: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: """ @@ -755,8 +760,8 @@ def dpm_solver_first_order_update( self, model_output: torch.Tensor, *args, - sample: Optional[torch.Tensor] = None, - noise: Optional[torch.Tensor] = None, + sample: torch.Tensor | None = None, + noise: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: """ @@ -817,10 +822,10 @@ def dpm_solver_first_order_update( def singlestep_dpm_solver_second_order_update( self, - model_output_list: List[torch.Tensor], + model_output_list: list[torch.Tensor], *args, - sample: Optional[torch.Tensor] = None, - noise: Optional[torch.Tensor] = None, + sample: torch.Tensor | None = None, + noise: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: """ @@ -828,7 +833,7 @@ def singlestep_dpm_solver_second_order_update( time `timestep_list[-2]`. Args: - model_output_list (`List[torch.Tensor]`): + model_output_list (`list[torch.Tensor]`): The direct outputs from learned diffusion model at current and latter timesteps. timestep (`int`): The current and latter discrete timestep in the diffusion chain. @@ -928,10 +933,10 @@ def singlestep_dpm_solver_second_order_update( def singlestep_dpm_solver_third_order_update( self, - model_output_list: List[torch.Tensor], + model_output_list: list[torch.Tensor], *args, - sample: Optional[torch.Tensor] = None, - noise: Optional[torch.Tensor] = None, + sample: torch.Tensor | None = None, + noise: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: """ @@ -939,7 +944,7 @@ def singlestep_dpm_solver_third_order_update( time `timestep_list[-3]`. Args: - model_output_list (`List[torch.Tensor]`): + model_output_list (`list[torch.Tensor]`): The direct outputs from learned diffusion model at current and latter timesteps. timestep (`int`): The current and latter discrete timestep in the diffusion chain. @@ -1050,18 +1055,18 @@ def singlestep_dpm_solver_third_order_update( def singlestep_dpm_solver_update( self, - model_output_list: List[torch.Tensor], + model_output_list: list[torch.Tensor], *args, - sample: Optional[torch.Tensor] = None, - order: Optional[int] = None, - noise: Optional[torch.Tensor] = None, + sample: torch.Tensor | None = None, + order: int = None, + noise: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: """ One step for the singlestep DPMSolver. Args: - model_output_list (`List[torch.Tensor]`): + model_output_list (`list[torch.Tensor]`): The direct outputs from learned diffusion model at current and latter timesteps. timestep (`int`): The current and latter discrete timestep in the diffusion chain. @@ -1113,7 +1118,9 @@ def singlestep_dpm_solver_update( # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep def index_for_timestep( - self, timestep: Union[int, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None + self, + timestep: int | torch.Tensor, + schedule_timesteps: torch.Tensor | None = None, ) -> int: """ Find the index for a given timestep in the schedule. @@ -1147,7 +1154,7 @@ def index_for_timestep( return step_index # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index - def _init_step_index(self, timestep: Union[int, torch.Tensor]) -> None: + def _init_step_index(self, timestep: int | torch.Tensor) -> None: """ Initialize the step_index counter for the scheduler. @@ -1166,11 +1173,11 @@ def _init_step_index(self, timestep: Union[int, torch.Tensor]) -> None: def step( self, model_output: torch.Tensor, - timestep: Union[int, torch.Tensor], + timestep: int | torch.Tensor, sample: torch.Tensor, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: + ) -> SchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the singlestep DPMSolver. diff --git a/src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py b/src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py index d4e8ca5e8b18..20ad17cfdd96 100644 --- a/src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py +++ b/src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py @@ -15,7 +15,7 @@ # DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver and https://github.com/NVlabs/edm import math -from typing import List, Optional, Tuple, Union +from typing import List, Literal import numpy as np import torch @@ -51,13 +51,15 @@ class EDMDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): schedule was incorporated in this model: https://huggingface.co/stabilityai/cosxl. num_train_timesteps (`int`, defaults to 1000): The number of diffusion steps to train the model. - solver_order (`int`, defaults to 2): - The DPMSolver order which can be `1` or `2` or `3`. It is recommended to use `solver_order=2` for guided - sampling, and `solver_order=3` for unconditional sampling. prediction_type (`str`, defaults to `epsilon`, *optional*): Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process), `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen Video](https://huggingface.co/papers/2210.02303) paper). + rho (`float`, *optional*, defaults to 7.0): + The rho parameter in the Karras sigma schedule. This was set to 7.0 in the EDM paper [1]. + solver_order (`int`, defaults to 2): + The DPMSolver order which can be `1` or `2` or `3`. It is recommended to use `solver_order=2` for guided + sampling, and `solver_order=3` for unconditional sampling. thresholding (`bool`, defaults to `False`): Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such as Stable Diffusion. @@ -94,19 +96,19 @@ def __init__( sigma_min: float = 0.002, sigma_max: float = 80.0, sigma_data: float = 0.5, - sigma_schedule: str = "karras", + sigma_schedule: Literal["karras", "exponential"] = "karras", num_train_timesteps: int = 1000, - prediction_type: str = "epsilon", + prediction_type: Literal["epsilon", "sample", "v_prediction"] = "epsilon", rho: float = 7.0, solver_order: int = 2, thresholding: bool = False, dynamic_thresholding_ratio: float = 0.995, sample_max_value: float = 1.0, - algorithm_type: str = "dpmsolver++", - solver_type: str = "midpoint", + algorithm_type: Literal["dpmsolver++", "sde-dpmsolver++"] = "dpmsolver++", + solver_type: Literal["midpoint", "heun"] = "midpoint", lower_order_final: bool = True, euler_at_final: bool = False, - final_sigmas_type: Optional[str] = "zero", # "zero", "sigma_min" + final_sigmas_type: Literal["zero", "sigma_min"] = "zero", # "zero", "sigma_min" ): # settings for DPM-Solver if algorithm_type not in ["dpmsolver++", "sde-dpmsolver++"]: @@ -145,19 +147,19 @@ def __init__( self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication @property - def init_noise_sigma(self): + def init_noise_sigma(self) -> float: # standard deviation of the initial noise distribution return (self.config.sigma_max**2 + 1) ** 0.5 @property - def step_index(self): + def step_index(self) -> int: """ The index counter for current timestep. It will increase 1 after each scheduler step. """ return self._step_index @property - def begin_index(self): + def begin_index(self) -> int: """ The index for the first timestep. It should be set from pipeline with `set_begin_index` method. """ @@ -175,13 +177,37 @@ def set_begin_index(self, begin_index: int = 0): self._begin_index = begin_index # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.precondition_inputs - def precondition_inputs(self, sample, sigma): + def precondition_inputs(self, sample: torch.Tensor, sigma: float | torch.Tensor) -> torch.Tensor: + """ + Precondition the input sample by scaling it according to the EDM formulation. + + Args: + sample (`torch.Tensor`): + The input sample tensor to precondition. + sigma (`float` or `torch.Tensor`): + The current sigma (noise level) value. + + Returns: + `torch.Tensor`: + The scaled input sample. + """ c_in = self._get_conditioning_c_in(sigma) scaled_sample = sample * c_in return scaled_sample # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.precondition_noise - def precondition_noise(self, sigma): + def precondition_noise(self, sigma: float | torch.Tensor) -> torch.Tensor: + """ + Precondition the noise level by applying a logarithmic transformation. + + Args: + sigma (`float` or `torch.Tensor`): + The sigma (noise level) value to precondition. + + Returns: + `torch.Tensor`: + The preconditioned noise value computed as `0.25 * log(sigma)`. + """ if not isinstance(sigma, torch.Tensor): sigma = torch.tensor([sigma]) @@ -190,7 +216,27 @@ def precondition_noise(self, sigma): return c_noise # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.precondition_outputs - def precondition_outputs(self, sample, model_output, sigma): + def precondition_outputs( + self, + sample: torch.Tensor, + model_output: torch.Tensor, + sigma: float | torch.Tensor, + ) -> torch.Tensor: + """ + Precondition the model outputs according to the EDM formulation. + + Args: + sample (`torch.Tensor`): + The input sample tensor. + model_output (`torch.Tensor`): + The direct output from the learned diffusion model. + sigma (`float` or `torch.Tensor`): + The current sigma (noise level) value. + + Returns: + `torch.Tensor`: + The denoised sample computed by combining the skip connection and output scaling. + """ sigma_data = self.config.sigma_data c_skip = sigma_data**2 / (sigma**2 + sigma_data**2) @@ -206,15 +252,15 @@ def precondition_outputs(self, sample, model_output, sigma): return denoised # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.scale_model_input - def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: float | torch.Tensor) -> torch.Tensor: """ - Ensures interchangeability with schedulers that need to scale the denoising model input depending on the - current timestep. Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the Euler algorithm. + Scale the denoising model input to match the Euler algorithm. Ensures interchangeability with schedulers that + need to scale the denoising model input depending on the current timestep. Args: sample (`torch.Tensor`): - The input sample. - timestep (`int`, *optional*): + The input sample tensor. + timestep (`float` or `torch.Tensor`): The current timestep in the diffusion chain. Returns: @@ -230,7 +276,11 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.T self.is_scale_input_called = True return sample - def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torch.device] = None): + def set_timesteps( + self, + num_inference_steps: int = None, + device: str | torch.device | None = None, + ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -274,8 +324,27 @@ def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torc self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler._compute_karras_sigmas - def _compute_karras_sigmas(self, ramp, sigma_min=None, sigma_max=None) -> torch.Tensor: - """Constructs the noise schedule of Karras et al. (2022).""" + def _compute_karras_sigmas( + self, + ramp: torch.Tensor, + sigma_min: float | None = None, + sigma_max: float | None = None, + ) -> torch.Tensor: + """ + Construct the noise schedule of [Karras et al. (2022)](https://huggingface.co/papers/2206.00364). + + Args: + ramp (`torch.Tensor`): + A tensor of values in [0, 1] representing the interpolation positions. + sigma_min (`float`, *optional*): + Minimum sigma value. If `None`, uses `self.config.sigma_min`. + sigma_max (`float`, *optional*): + Maximum sigma value. If `None`, uses `self.config.sigma_max`. + + Returns: + `torch.Tensor`: + The computed Karras sigma schedule. + """ sigma_min = sigma_min or self.config.sigma_min sigma_max = sigma_max or self.config.sigma_max @@ -286,10 +355,27 @@ def _compute_karras_sigmas(self, ramp, sigma_min=None, sigma_max=None) -> torch. return sigmas # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler._compute_exponential_sigmas - def _compute_exponential_sigmas(self, ramp, sigma_min=None, sigma_max=None) -> torch.Tensor: - """Implementation closely follows k-diffusion. - + def _compute_exponential_sigmas( + self, + ramp: torch.Tensor, + sigma_min: float | None = None, + sigma_max: float | None = None, + ) -> torch.Tensor: + """ + Compute the exponential sigma schedule. Implementation closely follows k-diffusion: https://github.com/crowsonkb/k-diffusion/blob/6ab5146d4a5ef63901326489f31f1d8e7dd36b48/k_diffusion/sampling.py#L26 + + Args: + ramp (`torch.Tensor`): + A tensor of values representing the interpolation positions. + sigma_min (`float`, *optional*): + Minimum sigma value. If `None`, uses `self.config.sigma_min`. + sigma_max (`float`, *optional*): + Maximum sigma value. If `None`, uses `self.config.sigma_max`. + + Returns: + `torch.Tensor`: + The computed exponential sigma schedule. """ sigma_min = sigma_min or self.config.sigma_min sigma_max = sigma_max or self.config.sigma_max @@ -380,13 +466,12 @@ def _sigma_to_t(self, sigma, log_sigmas): def _sigma_to_alpha_sigma_t(self, sigma): alpha_t = torch.tensor(1) # Inputs are pre-scaled before going into unet, so alpha_t = 1 sigma_t = sigma - return alpha_t, sigma_t def convert_model_output( self, model_output: torch.Tensor, - sample: torch.Tensor = None, + sample: torch.Tensor, ) -> torch.Tensor: """ Convert the model output to the corresponding type the DPMSolver/DPMSolver++ algorithm needs. DPM-Solver is @@ -417,8 +502,8 @@ def convert_model_output( def dpm_solver_first_order_update( self, model_output: torch.Tensor, - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, + sample: torch.Tensor, + noise: torch.Tensor | None = None, ) -> torch.Tensor: """ One step for the first-order DPMSolver (equivalent to DDIM). @@ -428,12 +513,17 @@ def dpm_solver_first_order_update( The direct output from the learned diffusion model. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. + noise (`torch.Tensor`, *optional*): + The noise tensor to add to the original samples. Returns: `torch.Tensor`: The sample tensor at the previous timestep. """ - sigma_t, sigma_s = self.sigmas[self.step_index + 1], self.sigmas[self.step_index] + sigma_t, sigma_s = ( + self.sigmas[self.step_index + 1], + self.sigmas[self.step_index], + ) alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) alpha_s, sigma_s = self._sigma_to_alpha_sigma_t(sigma_s) lambda_t = torch.log(alpha_t) - torch.log(sigma_t) @@ -455,17 +545,19 @@ def dpm_solver_first_order_update( def multistep_dpm_solver_second_order_update( self, model_output_list: List[torch.Tensor], - sample: torch.Tensor = None, - noise: Optional[torch.Tensor] = None, + sample: torch.Tensor, + noise: torch.Tensor | None = None, ) -> torch.Tensor: """ One step for the second-order multistep DPMSolver. Args: - model_output_list (`List[torch.Tensor]`): + model_output_list (`list[torch.Tensor]`): The direct outputs from learned diffusion model at current and latter timesteps. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. + noise (`torch.Tensor`, *optional*): + The noise tensor to add to the original samples. Returns: `torch.Tensor`: @@ -525,14 +617,14 @@ def multistep_dpm_solver_second_order_update( def multistep_dpm_solver_third_order_update( self, - model_output_list: List[torch.Tensor], - sample: torch.Tensor = None, + model_output_list: list[torch.Tensor], + sample: torch.Tensor, ) -> torch.Tensor: """ One step for the third-order multistep DPMSolver. Args: - model_output_list (`List[torch.Tensor]`): + model_output_list (`list[torch.Tensor]`): The direct outputs from learned diffusion model at current and latter timesteps. sample (`torch.Tensor`): A current instance of a sample created by diffusion process. @@ -579,7 +671,9 @@ def multistep_dpm_solver_third_order_update( # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep def index_for_timestep( - self, timestep: Union[int, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None + self, + timestep: int | torch.Tensor, + schedule_timesteps: torch.Tensor | None = None, ) -> int: """ Find the index for a given timestep in the schedule. @@ -613,7 +707,7 @@ def index_for_timestep( return step_index # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index - def _init_step_index(self, timestep): + def _init_step_index(self, timestep: int | torch.Tensor) -> None: """ Initialize the step_index counter for the scheduler. @@ -632,11 +726,11 @@ def _init_step_index(self, timestep): def step( self, model_output: torch.Tensor, - timestep: Union[int, torch.Tensor], + timestep: int | torch.Tensor, sample: torch.Tensor, - generator=None, + generator: torch.Generator | None = None, return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: + ) -> SchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the multistep DPMSolver. @@ -684,7 +778,10 @@ def step( if self.config.algorithm_type == "sde-dpmsolver++": noise = randn_tensor( - model_output.shape, generator=generator, device=model_output.device, dtype=model_output.dtype + model_output.shape, + generator=generator, + device=model_output.device, + dtype=model_output.dtype, ) else: noise = None @@ -757,9 +854,20 @@ def add_noise( return noisy_samples # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler._get_conditioning_c_in - def _get_conditioning_c_in(self, sigma): + def _get_conditioning_c_in(self, sigma: float | torch.Tensor) -> float | torch.Tensor: + """ + Compute the input conditioning factor for the EDM formulation. + + Args: + sigma (`float` or `torch.Tensor`): + The current sigma (noise level) value. + + Returns: + `float` or `torch.Tensor`: + The input conditioning factor `c_in`. + """ c_in = 1 / ((sigma**2 + self.config.sigma_data**2) ** 0.5) return c_in - def __len__(self): + def __len__(self) -> int: return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_edm_euler.py b/src/diffusers/schedulers/scheduling_edm_euler.py index 2ed05d396514..dfb70ce68b46 100644 --- a/src/diffusers/schedulers/scheduling_edm_euler.py +++ b/src/diffusers/schedulers/scheduling_edm_euler.py @@ -14,7 +14,7 @@ import math from dataclasses import dataclass -from typing import List, Optional, Tuple, Union +from typing import Literal import torch @@ -43,7 +43,7 @@ class EDMEulerSchedulerOutput(BaseOutput): """ prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None + pred_original_sample: torch.Tensor | None = None class EDMEulerScheduler(SchedulerMixin, ConfigMixin): @@ -57,29 +57,28 @@ class EDMEulerScheduler(SchedulerMixin, ConfigMixin): methods the library implements for all schedulers such as loading and saving. Args: - sigma_min (`float`, *optional*, defaults to 0.002): + sigma_min (`float`, *optional*, defaults to `0.002`): Minimum noise magnitude in the sigma schedule. This was set to 0.002 in the EDM paper [1]; a reasonable range is [0, 10]. - sigma_max (`float`, *optional*, defaults to 80.0): + sigma_max (`float`, *optional*, defaults to `80.0`): Maximum noise magnitude in the sigma schedule. This was set to 80.0 in the EDM paper [1]; a reasonable range is [0.2, 80.0]. - sigma_data (`float`, *optional*, defaults to 0.5): + sigma_data (`float`, *optional*, defaults to `0.5`): The standard deviation of the data distribution. This is set to 0.5 in the EDM paper [1]. - sigma_schedule (`str`, *optional*, defaults to `karras`): - Sigma schedule to compute the `sigmas`. By default, we the schedule introduced in the EDM paper - (https://huggingface.co/papers/2206.00364). Other acceptable value is "exponential". The exponential - schedule was incorporated in this model: https://huggingface.co/stabilityai/cosxl. - num_train_timesteps (`int`, defaults to 1000): + sigma_schedule (`Literal["karras", "exponential"]`, *optional*, defaults to `"karras"`): + Sigma schedule to compute the `sigmas`. By default, we use the schedule introduced in the EDM paper + (https://huggingface.co/papers/2206.00364). The `"exponential"` schedule was incorporated in this model: + https://huggingface.co/stabilityai/cosxl. + num_train_timesteps (`int`, *optional*, defaults to `1000`): The number of diffusion steps to train the model. - prediction_type (`str`, defaults to `epsilon`, *optional*): - Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process), - `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen - Video](https://huggingface.co/papers/2210.02303) paper). - rho (`float`, *optional*, defaults to 7.0): + prediction_type (`Literal["epsilon", "v_prediction"]`, *optional*, defaults to `"epsilon"`): + Prediction type of the scheduler function. `"epsilon"` predicts the noise of the diffusion process, and + `"v_prediction"` (see section 2.4 of [Imagen Video](https://huggingface.co/papers/2210.02303) paper). + rho (`float`, *optional*, defaults to `7.0`): The rho parameter used for calculating the Karras sigma schedule, which is set to 7.0 in the EDM paper [1]. - final_sigmas_type (`str`, defaults to `"zero"`): + final_sigmas_type (`Literal["zero", "sigma_min"]`, *optional*, defaults to `"zero"`): The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final - sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0. + sigma is the same as the last sigma in the training schedule. If `"zero"`, the final sigma is set to 0. """ _compatibles = [] @@ -91,12 +90,12 @@ def __init__( sigma_min: float = 0.002, sigma_max: float = 80.0, sigma_data: float = 0.5, - sigma_schedule: str = "karras", + sigma_schedule: Literal["karras", "exponential"] = "karras", num_train_timesteps: int = 1000, - prediction_type: str = "epsilon", + prediction_type: Literal["epsilon", "v_prediction"] = "epsilon", rho: float = 7.0, - final_sigmas_type: str = "zero", # can be "zero" or "sigma_min" - ): + final_sigmas_type: Literal["zero", "sigma_min"] = "zero", + ) -> None: if sigma_schedule not in ["karras", "exponential"]: raise ValueError(f"Wrong value for provided for `{sigma_schedule=}`.`") @@ -131,26 +130,41 @@ def __init__( self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication @property - def init_noise_sigma(self): - # standard deviation of the initial noise distribution + def init_noise_sigma(self) -> float: + """ + Return the standard deviation of the initial noise distribution. + + Returns: + `float`: + The initial noise sigma value computed as `(sigma_max**2 + 1) ** 0.5`. + """ return (self.config.sigma_max**2 + 1) ** 0.5 @property - def step_index(self): + def step_index(self) -> int: """ - The index counter for current timestep. It will increase 1 after each scheduler step. + Return the index counter for the current timestep. The index will increase by 1 after each scheduler step. + + Returns: + `int` or `None`: + The current step index, or `None` if not yet initialized. """ return self._step_index @property - def begin_index(self): + def begin_index(self) -> int: """ - The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + Return the index for the first timestep. This should be set from the pipeline with the `set_begin_index` + method. + + Returns: + `int` or `None`: + The begin index, or `None` if not yet set. """ return self._begin_index # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index - def set_begin_index(self, begin_index: int = 0): + def set_begin_index(self, begin_index: int = 0) -> None: """ Sets the begin index for the scheduler. This function should be run from pipeline before the inference. @@ -160,12 +174,36 @@ def set_begin_index(self, begin_index: int = 0): """ self._begin_index = begin_index - def precondition_inputs(self, sample, sigma): + def precondition_inputs(self, sample: torch.Tensor, sigma: float | torch.Tensor) -> torch.Tensor: + """ + Precondition the input sample by scaling it according to the EDM formulation. + + Args: + sample (`torch.Tensor`): + The input sample tensor to precondition. + sigma (`float` or `torch.Tensor`): + The current sigma (noise level) value. + + Returns: + `torch.Tensor`: + The scaled input sample. + """ c_in = self._get_conditioning_c_in(sigma) scaled_sample = sample * c_in return scaled_sample - def precondition_noise(self, sigma): + def precondition_noise(self, sigma: float | torch.Tensor) -> torch.Tensor: + """ + Precondition the noise level by applying a logarithmic transformation. + + Args: + sigma (`float` or `torch.Tensor`): + The sigma (noise level) value to precondition. + + Returns: + `torch.Tensor`: + The preconditioned noise value computed as `0.25 * log(sigma)`. + """ if not isinstance(sigma, torch.Tensor): sigma = torch.tensor([sigma]) @@ -173,7 +211,27 @@ def precondition_noise(self, sigma): return c_noise - def precondition_outputs(self, sample, model_output, sigma): + def precondition_outputs( + self, + sample: torch.Tensor, + model_output: torch.Tensor, + sigma: float | torch.Tensor, + ) -> torch.Tensor: + """ + Precondition the model outputs according to the EDM formulation. + + Args: + sample (`torch.Tensor`): + The input sample tensor. + model_output (`torch.Tensor`): + The direct output from the learned diffusion model. + sigma (`float` or `torch.Tensor`): + The current sigma (noise level) value. + + Returns: + `torch.Tensor`: + The denoised sample computed by combining the skip connection and output scaling. + """ sigma_data = self.config.sigma_data c_skip = sigma_data**2 / (sigma**2 + sigma_data**2) @@ -188,15 +246,15 @@ def precondition_outputs(self, sample, model_output, sigma): return denoised - def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: float | torch.Tensor) -> torch.Tensor: """ - Ensures interchangeability with schedulers that need to scale the denoising model input depending on the - current timestep. Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the Euler algorithm. + Scale the denoising model input to match the Euler algorithm. Ensures interchangeability with schedulers that + need to scale the denoising model input depending on the current timestep. Args: sample (`torch.Tensor`): - The input sample. - timestep (`int`, *optional*): + The input sample tensor. + timestep (`float` or `torch.Tensor`): The current timestep in the diffusion chain. Returns: @@ -215,18 +273,18 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.T def set_timesteps( self, num_inference_steps: int = None, - device: Union[str, torch.device] = None, - sigmas: Optional[Union[torch.Tensor, List[float]]] = None, + device: str | torch.device = None, + sigmas: torch.Tensor | list[float] | None = None, ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). Args: - num_inference_steps (`int`): + num_inference_steps (`int`, *optional*): The number of diffusion steps used when generating samples with a pre-trained model. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - sigmas (`Union[torch.Tensor, List[float]]`, *optional*): + sigmas (`torch.Tensor | list[float]`, *optional*): Custom sigmas to use for the denoising process. If not defined, the default behavior when `num_inference_steps` is passed will be used. """ @@ -262,8 +320,27 @@ def set_timesteps( self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication # Taken from https://github.com/crowsonkb/k-diffusion/blob/686dbad0f39640ea25c8a8c6a6e56bb40eacefa2/k_diffusion/sampling.py#L17 - def _compute_karras_sigmas(self, ramp, sigma_min=None, sigma_max=None) -> torch.Tensor: - """Constructs the noise schedule of Karras et al. (2022).""" + def _compute_karras_sigmas( + self, + ramp: torch.Tensor, + sigma_min: float | None = None, + sigma_max: float | None = None, + ) -> torch.Tensor: + """ + Construct the noise schedule of [Karras et al. (2022)](https://huggingface.co/papers/2206.00364). + + Args: + ramp (`torch.Tensor`): + A tensor of values in [0, 1] representing the interpolation positions. + sigma_min (`float`, *optional*): + Minimum sigma value. If `None`, uses `self.config.sigma_min`. + sigma_max (`float`, *optional*): + Maximum sigma value. If `None`, uses `self.config.sigma_max`. + + Returns: + `torch.Tensor`: + The computed Karras sigma schedule. + """ sigma_min = sigma_min or self.config.sigma_min sigma_max = sigma_max or self.config.sigma_max @@ -273,10 +350,27 @@ def _compute_karras_sigmas(self, ramp, sigma_min=None, sigma_max=None) -> torch. sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho return sigmas - def _compute_exponential_sigmas(self, ramp, sigma_min=None, sigma_max=None) -> torch.Tensor: - """Implementation closely follows k-diffusion. - + def _compute_exponential_sigmas( + self, + ramp: torch.Tensor, + sigma_min: float | None = None, + sigma_max: float | None = None, + ) -> torch.Tensor: + """ + Compute the exponential sigma schedule. Implementation closely follows k-diffusion: https://github.com/crowsonkb/k-diffusion/blob/6ab5146d4a5ef63901326489f31f1d8e7dd36b48/k_diffusion/sampling.py#L26 + + Args: + ramp (`torch.Tensor`): + A tensor of values representing the interpolation positions. + sigma_min (`float`, *optional*): + Minimum sigma value. If `None`, uses `self.config.sigma_min`. + sigma_max (`float`, *optional*): + Maximum sigma value. If `None`, uses `self.config.sigma_max`. + + Returns: + `torch.Tensor`: + The computed exponential sigma schedule. """ sigma_min = sigma_min or self.config.sigma_min sigma_max = sigma_max or self.config.sigma_max @@ -285,7 +379,7 @@ def _compute_exponential_sigmas(self, ramp, sigma_min=None, sigma_max=None) -> t # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep def index_for_timestep( - self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None + self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None ) -> int: """ Find the index of a given timestep in the timestep schedule. @@ -315,7 +409,7 @@ def index_for_timestep( return indices[pos].item() # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: + def _init_step_index(self, timestep: float | torch.Tensor) -> None: """ Initialize the step index for the scheduler based on the given timestep. @@ -333,41 +427,47 @@ def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: def step( self, model_output: torch.Tensor, - timestep: Union[float, torch.Tensor], + timestep: float | torch.Tensor, sample: torch.Tensor, s_churn: float = 0.0, s_tmin: float = 0.0, s_tmax: float = float("inf"), s_noise: float = 1.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, return_dict: bool = True, - pred_original_sample: Optional[torch.Tensor] = None, - ) -> Union[EDMEulerSchedulerOutput, Tuple]: + pred_original_sample: torch.Tensor | None = None, + ) -> EDMEulerSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). Args: model_output (`torch.Tensor`): - The direct output from learned diffusion model. - timestep (`float`): + The direct output from the learned diffusion model. + timestep (`float` or `torch.Tensor`): The current discrete timestep in the diffusion chain. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. - s_churn (`float`): - s_tmin (`float`): - s_tmax (`float`): - s_noise (`float`, defaults to 1.0): + s_churn (`float`, *optional*, defaults to `0.0`): + The amount of stochasticity to add at each step. Higher values add more noise. + s_tmin (`float`, *optional*, defaults to `0.0`): + The minimum sigma threshold below which no noise is added. + s_tmax (`float`, *optional*, defaults to `float("inf")`): + The maximum sigma threshold above which no noise is added. + s_noise (`float`, *optional*, defaults to `1.0`): Scaling factor for noise added to the sample. generator (`torch.Generator`, *optional*): - A random number generator. - return_dict (`bool`): - Whether or not to return a [`~schedulers.scheduling_euler_discrete.EDMEulerSchedulerOutput`] or tuple. + A random number generator for reproducibility. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return an [`~schedulers.scheduling_edm_euler.EDMEulerSchedulerOutput`] or tuple. + pred_original_sample (`torch.Tensor`, *optional*): + The predicted denoised sample from a previous step. If provided, skips recomputation. Returns: - [`~schedulers.scheduling_euler_discrete.EDMEulerSchedulerOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EDMEulerSchedulerOutput`] is - returned, otherwise a tuple is returned where the first element is the sample tensor. + [`~schedulers.scheduling_edm_euler.EDMEulerSchedulerOutput`] or `tuple`: + If `return_dict` is `True`, an [`~schedulers.scheduling_edm_euler.EDMEulerSchedulerOutput`] is + returned, otherwise a tuple is returned where the first element is the previous sample tensor and the + second element is the predicted original sample tensor. """ if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)): @@ -399,7 +499,10 @@ def step( if gamma > 0: noise = randn_tensor( - model_output.shape, dtype=model_output.dtype, device=model_output.device, generator=generator + model_output.shape, + dtype=model_output.dtype, + device=model_output.device, + generator=generator, ) eps = noise * s_noise sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5 @@ -478,9 +581,20 @@ def add_noise( noisy_samples = original_samples + noise * sigma return noisy_samples - def _get_conditioning_c_in(self, sigma): + def _get_conditioning_c_in(self, sigma: float | torch.Tensor) -> float | torch.Tensor: + """ + Compute the input conditioning factor for the EDM formulation. + + Args: + sigma (`float` or `torch.Tensor`): + The current sigma (noise level) value. + + Returns: + `float` or `torch.Tensor`: + The input conditioning factor `c_in`. + """ c_in = 1 / ((sigma**2 + self.config.sigma_data**2) ** 0.5) return c_in - def __len__(self): + def __len__(self) -> int: return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py index 97fd84db5621..8ad6abd90668 100644 --- a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py @@ -14,7 +14,7 @@ import math from dataclasses import dataclass -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -44,14 +44,14 @@ class EulerAncestralDiscreteSchedulerOutput(BaseOutput): """ prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None + pred_original_sample: torch.Tensor | None = None # Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -65,8 +65,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -77,6 +77,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -173,10 +180,10 @@ def __init__( num_train_timesteps: int = 1000, beta_start: float = 0.0001, beta_end: float = 0.02, - beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, - prediction_type: Literal["epsilon", "sample", "v_prediction"] = "epsilon", - timestep_spacing: Literal["linspace", "leading", "trailing"] = "linspace", + beta_schedule: str = "linear", + trained_betas: np.ndarray | list[float] | None = None, + prediction_type: str = "epsilon", + timestep_spacing: str = "linspace", steps_offset: int = 0, rescale_betas_zero_snr: bool = False, ) -> None: @@ -227,14 +234,14 @@ def init_noise_sigma(self) -> torch.Tensor: return (self.sigmas.max() ** 2 + 1) ** 0.5 @property - def step_index(self) -> Optional[int]: + def step_index(self) -> int: """ The index counter for current timestep. It will increase 1 after each scheduler step. """ return self._step_index @property - def begin_index(self) -> Optional[int]: + def begin_index(self) -> int: """ The index for the first timestep. It should be set from pipeline with `set_begin_index` method. """ @@ -251,7 +258,7 @@ def set_begin_index(self, begin_index: int = 0) -> None: """ self._begin_index = begin_index - def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: float | torch.Tensor) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the Euler algorithm. @@ -275,7 +282,7 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.T self.is_scale_input_called = True return sample - def set_timesteps(self, num_inference_steps: int, device: Optional[Union[str, torch.device]] = None) -> None: + def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -321,7 +328,7 @@ def set_timesteps(self, num_inference_steps: int, device: Optional[Union[str, to # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep def index_for_timestep( - self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None + self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None ) -> int: """ Find the index of a given timestep in the timestep schedule. @@ -351,7 +358,7 @@ def index_for_timestep( return indices[pos].item() # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: + def _init_step_index(self, timestep: float | torch.Tensor) -> None: """ Initialize the step index for the scheduler based on the given timestep. @@ -369,11 +376,11 @@ def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: def step( self, model_output: torch.Tensor, - timestep: Union[float, torch.Tensor], + timestep: float | torch.Tensor, sample: torch.Tensor, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, return_dict: bool = True, - ) -> Union[EulerAncestralDiscreteSchedulerOutput, Tuple]: + ) -> EulerAncestralDiscreteSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py index a55a76626cec..484f2ca58e1b 100644 --- a/src/diffusers/schedulers/scheduling_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_discrete.py @@ -14,7 +14,7 @@ import math from dataclasses import dataclass -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -47,14 +47,14 @@ class EulerDiscreteSchedulerOutput(BaseOutput): """ prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None + pred_original_sample: torch.Tensor | None = None # Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -68,8 +68,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -80,6 +80,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -198,15 +205,15 @@ def __init__( num_train_timesteps: int = 1000, beta_start: float = 0.0001, beta_end: float = 0.02, - beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, - prediction_type: Literal["epsilon", "sample", "v_prediction"] = "epsilon", - interpolation_type: Literal["linear", "log_linear"] = "linear", - use_karras_sigmas: Optional[bool] = False, - use_exponential_sigmas: Optional[bool] = False, - use_beta_sigmas: Optional[bool] = False, - sigma_min: Optional[float] = None, - sigma_max: Optional[float] = None, + beta_schedule: str = "linear", + trained_betas: np.ndarray | list[float] | None = None, + prediction_type: str = "epsilon", + interpolation_type: str = "linear", + use_karras_sigmas: bool | None = False, + use_exponential_sigmas: bool | None = False, + use_beta_sigmas: bool | None = False, + sigma_min: float | None = None, + sigma_max: float | None = None, timestep_spacing: Literal["linspace", "leading", "trailing"] = "linspace", timestep_type: Literal["discrete", "continuous"] = "discrete", steps_offset: int = 0, @@ -268,7 +275,7 @@ def __init__( self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication @property - def init_noise_sigma(self) -> Union[float, torch.Tensor]: + def init_noise_sigma(self) -> float | torch.Tensor: """ The standard deviation of the initial noise distribution. @@ -284,7 +291,7 @@ def init_noise_sigma(self) -> Union[float, torch.Tensor]: return (max_sigma**2 + 1) ** 0.5 @property - def step_index(self) -> Optional[int]: + def step_index(self) -> int: """ The index counter for current timestep. It will increase by 1 after each scheduler step. @@ -295,7 +302,7 @@ def step_index(self) -> Optional[int]: return self._step_index @property - def begin_index(self) -> Optional[int]: + def begin_index(self) -> int: """ The index for the first timestep. It should be set from pipeline with `set_begin_index` method. @@ -316,7 +323,7 @@ def set_begin_index(self, begin_index: int = 0) -> None: """ self._begin_index = begin_index - def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: float | torch.Tensor) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the Euler algorithm. @@ -342,11 +349,11 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.T def set_timesteps( self, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - sigmas: Optional[List[float]] = None, - ) -> None: + num_inference_steps: int = None, + device: str | torch.device = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -356,14 +363,15 @@ def set_timesteps( `timesteps` or `sigmas` must be provided. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to support arbitrary timesteps schedule. If `None`, timesteps will be generated based on the `timestep_spacing` attribute. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`, and `timestep_spacing` attribute will be ignored. - sigmas (`List[float]`, *optional*): - Custom sigmas used to support arbitrary timesteps schedule. If `None`, timesteps and sigmas will be - generated based on the relevant scheduler attributes. If `sigmas` is passed, `num_inference_steps` and - `timesteps` must be `None`, and the timesteps will be generated based on the custom sigmas schedule. + sigmas (`list[float]`, *optional*): + Custom sigmas used to support arbitrary timesteps schedule schedule. If `None`, timesteps and sigmas + will be generated based on the relevant scheduler attributes. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`, and the timesteps will be generated based on the + custom sigmas schedule. """ if timesteps is not None and sigmas is not None: @@ -631,7 +639,7 @@ def _convert_to_beta( return sigmas def index_for_timestep( - self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None + self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None ) -> int: """ Find the index of a given timestep in the timestep schedule. @@ -660,7 +668,7 @@ def index_for_timestep( return indices[pos].item() - def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: + def _init_step_index(self, timestep: float | torch.Tensor) -> None: """ Initialize the step index for the scheduler based on the given timestep. @@ -678,15 +686,15 @@ def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: def step( self, model_output: torch.Tensor, - timestep: Union[float, torch.Tensor], + timestep: float | torch.Tensor, sample: torch.Tensor, s_churn: float = 0.0, s_tmin: float = 0.0, s_tmax: float = float("inf"), s_noise: float = 1.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, return_dict: bool = True, - ) -> Union[EulerDiscreteSchedulerOutput, Tuple]: + ) -> EulerDiscreteSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). diff --git a/src/diffusers/schedulers/scheduling_euler_discrete_flax.py b/src/diffusers/schedulers/scheduling_euler_discrete_flax.py index 09341c909d2e..fc555dd21395 100644 --- a/src/diffusers/schedulers/scheduling_euler_discrete_flax.py +++ b/src/diffusers/schedulers/scheduling_euler_discrete_flax.py @@ -13,12 +13,12 @@ # limitations under the License. from dataclasses import dataclass -from typing import Optional, Tuple, Union import flax import jax.numpy as jnp from ..configuration_utils import ConfigMixin, register_to_config +from ..utils import logging from .scheduling_utils_flax import ( CommonSchedulerState, FlaxKarrasDiffusionSchedulers, @@ -28,6 +28,9 @@ ) +logger = logging.get_logger(__name__) + + @flax.struct.dataclass class EulerDiscreteSchedulerState: common: CommonSchedulerState @@ -36,13 +39,22 @@ class EulerDiscreteSchedulerState: init_noise_sigma: jnp.ndarray timesteps: jnp.ndarray sigmas: jnp.ndarray - num_inference_steps: Optional[int] = None + num_inference_steps: int = None @classmethod def create( - cls, common: CommonSchedulerState, init_noise_sigma: jnp.ndarray, timesteps: jnp.ndarray, sigmas: jnp.ndarray + cls, + common: CommonSchedulerState, + init_noise_sigma: jnp.ndarray, + timesteps: jnp.ndarray, + sigmas: jnp.ndarray, ): - return cls(common=common, init_noise_sigma=init_noise_sigma, timesteps=timesteps, sigmas=sigmas) + return cls( + common=common, + init_noise_sigma=init_noise_sigma, + timesteps=timesteps, + sigmas=sigmas, + ) @dataclass @@ -94,14 +106,18 @@ def __init__( beta_start: float = 0.0001, beta_end: float = 0.02, beta_schedule: str = "linear", - trained_betas: Optional[jnp.ndarray] = None, + trained_betas: jnp.ndarray | None = None, prediction_type: str = "epsilon", timestep_spacing: str = "linspace", dtype: jnp.dtype = jnp.float32, ): + logger.warning( + "Flax classes are deprecated and will be removed in Diffusers v1.0.0. We " + "recommend migrating to PyTorch classes or pinning your version of Diffusers." + ) self.dtype = dtype - def create_state(self, common: Optional[CommonSchedulerState] = None) -> EulerDiscreteSchedulerState: + def create_state(self, common: CommonSchedulerState | None = None) -> EulerDiscreteSchedulerState: if common is None: common = CommonSchedulerState.create(self) @@ -146,7 +162,10 @@ def scale_model_input(self, state: EulerDiscreteSchedulerState, sample: jnp.ndar return sample def set_timesteps( - self, state: EulerDiscreteSchedulerState, num_inference_steps: int, shape: Tuple = () + self, + state: EulerDiscreteSchedulerState, + num_inference_steps: int, + shape: tuple = (), ) -> EulerDiscreteSchedulerState: """ Sets the timesteps used for the diffusion chain. Supporting function to be run before inference. @@ -159,7 +178,12 @@ def set_timesteps( """ if self.config.timestep_spacing == "linspace": - timesteps = jnp.linspace(self.config.num_train_timesteps - 1, 0, num_inference_steps, dtype=self.dtype) + timesteps = jnp.linspace( + self.config.num_train_timesteps - 1, + 0, + num_inference_steps, + dtype=self.dtype, + ) elif self.config.timestep_spacing == "leading": step_ratio = self.config.num_train_timesteps // num_inference_steps timesteps = (jnp.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(float) @@ -193,7 +217,7 @@ def step( timestep: int, sample: jnp.ndarray, return_dict: bool = True, - ) -> Union[FlaxEulerDiscreteSchedulerOutput, Tuple]: + ) -> FlaxEulerDiscreteSchedulerOutput | tuple: """ Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion process from the learned model outputs (most often the predicted noise). diff --git a/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py b/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py index 9fd61d9e18d1..1021abf0f6f6 100644 --- a/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py @@ -14,7 +14,7 @@ import math from dataclasses import dataclass -from typing import List, Optional, Tuple, Union +from typing import Literal, Optional, Union import numpy as np import torch @@ -93,21 +93,30 @@ def __init__( num_train_timesteps: int = 1000, shift: float = 1.0, use_dynamic_shifting: bool = False, - base_shift: Optional[float] = 0.5, - max_shift: Optional[float] = 1.15, - base_image_seq_len: Optional[int] = 256, - max_image_seq_len: Optional[int] = 4096, + base_shift: float | None = 0.5, + max_shift: float | None = 1.15, + base_image_seq_len: int = 256, + max_image_seq_len: int = 4096, invert_sigmas: bool = False, - shift_terminal: Optional[float] = None, - use_karras_sigmas: Optional[bool] = False, - use_exponential_sigmas: Optional[bool] = False, - use_beta_sigmas: Optional[bool] = False, - time_shift_type: str = "exponential", + shift_terminal: float = None, + use_karras_sigmas: bool = False, + use_exponential_sigmas: bool = False, + use_beta_sigmas: bool = False, + time_shift_type: Literal["exponential", "linear"] = "exponential", stochastic_sampling: bool = False, ): if self.config.use_beta_sigmas and not is_scipy_available(): raise ImportError("Make sure to install scipy if you want to use beta sigmas.") - if sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1: + if ( + sum( + [ + self.config.use_beta_sigmas, + self.config.use_exponential_sigmas, + self.config.use_karras_sigmas, + ] + ) + > 1 + ): raise ValueError( "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." ) @@ -166,13 +175,20 @@ def set_begin_index(self, begin_index: int = 0): self._begin_index = begin_index def set_shift(self, shift: float): + """ + Sets the shift value for the scheduler. + + Args: + shift (`float`): + The shift value to be set. + """ self._shift = shift def scale_noise( self, sample: torch.FloatTensor, - timestep: Union[float, torch.FloatTensor], - noise: Optional[torch.FloatTensor] = None, + timestep: float | torch.FloatTensor, + noise: torch.FloatTensor | None = None, ) -> torch.FloatTensor: """ Forward process in flow-matching @@ -180,8 +196,10 @@ def scale_noise( Args: sample (`torch.FloatTensor`): The input sample. - timestep (`int`, *optional*): + timestep (`torch.FloatTensor`): The current timestep in the diffusion chain. + noise (`torch.FloatTensor`): + The noise tensor. Returns: `torch.FloatTensor`: @@ -216,10 +234,25 @@ def scale_noise( return sample - def _sigma_to_t(self, sigma): + def _sigma_to_t(self, sigma) -> float: return sigma * self.config.num_train_timesteps - def time_shift(self, mu: float, sigma: float, t: torch.Tensor): + def time_shift(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor: + """ + Apply time shifting to the sigmas. + + Args: + mu (`float`): + The mu parameter for the time shift. + sigma (`float`): + The sigma parameter for the time shift. + t (`torch.Tensor`): + The input timesteps. + + Returns: + `torch.Tensor`: + The time-shifted timesteps. + """ if self.config.time_shift_type == "exponential": return self._time_shift_exponential(mu, sigma, t) elif self.config.time_shift_type == "linear": @@ -248,11 +281,11 @@ def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor: def set_timesteps( self, - num_inference_steps: Optional[int] = None, - device: Union[str, torch.device] = None, - sigmas: Optional[List[float]] = None, - mu: Optional[float] = None, - timesteps: Optional[List[float]] = None, + num_inference_steps: int | None = None, + device: str | torch.device = None, + sigmas: list[float] | None = None, + mu: float | None = None, + timesteps: list[float] | None = None, ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -262,13 +295,13 @@ def set_timesteps( The number of diffusion steps used when generating samples with a pre-trained model. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom values for sigmas to be used for each diffusion step. If `None`, the sigmas are computed automatically. mu (`float`, *optional*): Determines the amount of shifting applied to sigmas when performing resolution-dependent timestep shifting. - timesteps (`List[float]`, *optional*): + timesteps (`list[float]`, *optional*): Custom values for timesteps to be used for each diffusion step. If `None`, the timesteps are computed automatically. """ @@ -300,7 +333,9 @@ def set_timesteps( if sigmas is None: if timesteps is None: timesteps = np.linspace( - self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps + self._sigma_to_t(self.sigma_max), + self._sigma_to_t(self.sigma_min), + num_inference_steps, ) sigmas = timesteps / self.config.num_train_timesteps else: @@ -348,7 +383,24 @@ def set_timesteps( self._step_index = None self._begin_index = None - def index_for_timestep(self, timestep, schedule_timesteps=None): + def index_for_timestep( + self, + timestep: Union[float, torch.FloatTensor], + schedule_timesteps: Optional[torch.FloatTensor] = None, + ) -> int: + """ + Get the index for the given timestep. + + Args: + timestep (`float` or `torch.FloatTensor`): + The timestep to find the index for. + schedule_timesteps (`torch.FloatTensor`, *optional*): + The schedule timesteps to validate against. If `None`, the scheduler's timesteps are used. + + Returns: + `int`: + The index of the timestep. + """ if schedule_timesteps is None: schedule_timesteps = self.timesteps @@ -362,7 +414,7 @@ def index_for_timestep(self, timestep, schedule_timesteps=None): return indices[pos].item() - def _init_step_index(self, timestep): + def _init_step_index(self, timestep: Union[float, torch.FloatTensor]) -> None: if self.begin_index is None: if isinstance(timestep, torch.Tensor): timestep = timestep.to(self.timesteps.device) @@ -373,16 +425,16 @@ def _init_step_index(self, timestep): def step( self, model_output: torch.FloatTensor, - timestep: Union[float, torch.FloatTensor], + timestep: float | torch.FloatTensor, sample: torch.FloatTensor, s_churn: float = 0.0, s_tmin: float = 0.0, s_tmax: float = float("inf"), s_noise: float = 1.0, - generator: Optional[torch.Generator] = None, - per_token_timesteps: Optional[torch.Tensor] = None, + generator: torch.Generator | None = None, + per_token_timesteps: torch.Tensor | None = None, return_dict: bool = True, - ) -> Union[FlowMatchEulerDiscreteSchedulerOutput, Tuple]: + ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). @@ -403,7 +455,7 @@ def step( A random number generator. per_token_timesteps (`torch.Tensor`, *optional*): The timesteps for each token in the sample. - return_dict (`bool`): + return_dict (`bool`, defaults to `True`): Whether or not to return a [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or tuple. @@ -472,7 +524,7 @@ def step( return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample) # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras - def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor: + def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: """ Construct the noise schedule as proposed in [Elucidating the Design Space of Diffusion-Based Generative Models](https://huggingface.co/papers/2206.00364). @@ -593,11 +645,11 @@ def _convert_to_beta( ) return sigmas - def _time_shift_exponential(self, mu, sigma, t): + def _time_shift_exponential(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor: return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) - def _time_shift_linear(self, mu, sigma, t): + def _time_shift_linear(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor: return mu / (mu + (1 / t - 1) ** sigma) - def __len__(self): + def __len__(self) -> int: return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_flow_match_heun_discrete.py b/src/diffusers/schedulers/scheduling_flow_match_heun_discrete.py index 6febee444c5a..49892e043abd 100644 --- a/src/diffusers/schedulers/scheduling_flow_match_heun_discrete.py +++ b/src/diffusers/schedulers/scheduling_flow_match_heun_discrete.py @@ -13,7 +13,6 @@ # limitations under the License. from dataclasses import dataclass -from typing import Optional, Tuple, Union import numpy as np import torch @@ -51,9 +50,6 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin): Args: num_train_timesteps (`int`, defaults to 1000): The number of diffusion steps to train the model. - timestep_spacing (`str`, defaults to `"linspace"`): - The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and - Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. shift (`float`, defaults to 1.0): The shift value for the timestep schedule. """ @@ -110,8 +106,8 @@ def set_begin_index(self, begin_index: int = 0): def scale_noise( self, sample: torch.FloatTensor, - timestep: Union[float, torch.FloatTensor], - noise: Optional[torch.FloatTensor] = None, + timestep: float | torch.FloatTensor, + noise: torch.FloatTensor, ) -> torch.FloatTensor: """ Forward process in flow-matching @@ -119,8 +115,10 @@ def scale_noise( Args: sample (`torch.FloatTensor`): The input sample. - timestep (`int`, *optional*): + timestep (`float` or `torch.FloatTensor`): The current timestep in the diffusion chain. + noise (`torch.FloatTensor`): + The noise tensor. Returns: `torch.FloatTensor`: @@ -130,14 +128,19 @@ def scale_noise( self._init_step_index(timestep) sigma = self.sigmas[self.step_index] + sample = sigma * noise + (1.0 - sigma) * sample return sample - def _sigma_to_t(self, sigma): + def _sigma_to_t(self, sigma: float) -> float: return sigma * self.config.num_train_timesteps - def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None): + def set_timesteps( + self, + num_inference_steps: int, + device: str | torch.device = None, + ) -> None: """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -150,7 +153,9 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic self.num_inference_steps = num_inference_steps timesteps = np.linspace( - self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps + self._sigma_to_t(self.sigma_max), + self._sigma_to_t(self.sigma_min), + num_inference_steps, ) sigmas = timesteps / self.config.num_train_timesteps @@ -171,7 +176,24 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic self._step_index = None self._begin_index = None - def index_for_timestep(self, timestep, schedule_timesteps=None): + def index_for_timestep( + self, + timestep: float | torch.FloatTensor, + schedule_timesteps: torch.FloatTensor | None = None, + ) -> int: + """ + Find the index of a given timestep in the timestep schedule. + + Args: + timestep (`float` or `torch.FloatTensor`): + The timestep value to find in the schedule. + schedule_timesteps (`torch.FloatTensor`, *optional*): + The timestep schedule to search in. If `None`, uses `self.timesteps`. + + Returns: + `int`: + The index of the timestep in the schedule. + """ if schedule_timesteps is None: schedule_timesteps = self.timesteps @@ -185,7 +207,7 @@ def index_for_timestep(self, timestep, schedule_timesteps=None): return indices[pos].item() - def _init_step_index(self, timestep): + def _init_step_index(self, timestep: float | torch.FloatTensor) -> None: if self.begin_index is None: if isinstance(timestep, torch.Tensor): timestep = timestep.to(self.timesteps.device) @@ -194,21 +216,24 @@ def _init_step_index(self, timestep): self._step_index = self._begin_index @property - def state_in_first_order(self): + def state_in_first_order(self) -> bool: + """ + Returns whether the scheduler is in the first-order state. + """ return self.dt is None def step( self, model_output: torch.FloatTensor, - timestep: Union[float, torch.FloatTensor], + timestep: float | torch.FloatTensor, sample: torch.FloatTensor, s_churn: float = 0.0, s_tmin: float = 0.0, s_tmax: float = float("inf"), s_noise: float = 1.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, return_dict: bool = True, - ) -> Union[FlowMatchHeunDiscreteSchedulerOutput, Tuple]: + ) -> FlowMatchHeunDiscreteSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). @@ -216,13 +241,19 @@ def step( Args: model_output (`torch.FloatTensor`): The direct output from learned diffusion model. - timestep (`float`): + timestep (`float` or `torch.FloatTensor`): The current discrete timestep in the diffusion chain. sample (`torch.FloatTensor`): A current instance of a sample created by the diffusion process. s_churn (`float`): - s_tmin (`float`): - s_tmax (`float`): + Stochasticity parameter that controls the amount of noise added during sampling. Higher values increase + randomness. + s_tmin (`float`): + Minimum timestep threshold for applying stochasticity. Only timesteps above this value will have noise + added. + s_tmax (`float`): + Maximum timestep threshold for applying stochasticity. Only timesteps below this value will have noise + added. s_noise (`float`, defaults to 1.0): Scaling factor for noise added to the sample. generator (`torch.Generator`, *optional*): @@ -271,7 +302,10 @@ def step( if gamma > 0: noise = randn_tensor( - model_output.shape, dtype=model_output.dtype, device=model_output.device, generator=generator + model_output.shape, + dtype=model_output.dtype, + device=model_output.device, + generator=generator, ) eps = noise * s_noise sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5 @@ -317,5 +351,5 @@ def step( return FlowMatchHeunDiscreteSchedulerOutput(prev_sample=prev_sample) - def __len__(self): + def __len__(self) -> int: return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_flow_match_lcm.py b/src/diffusers/schedulers/scheduling_flow_match_lcm.py index 25186d1fe969..97d4ebbc8e42 100644 --- a/src/diffusers/schedulers/scheduling_flow_match_lcm.py +++ b/src/diffusers/schedulers/scheduling_flow_match_lcm.py @@ -14,7 +14,7 @@ import math from dataclasses import dataclass -from typing import List, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -42,7 +42,7 @@ class FlowMatchLCMSchedulerOutput(BaseOutput): denoising loop. """ - prev_sample: torch.FloatTensor + prev_sample: torch.Tensor class FlowMatchLCMScheduler(SchedulerMixin, ConfigMixin): @@ -80,11 +80,11 @@ class FlowMatchLCMScheduler(SchedulerMixin, ConfigMixin): use_beta_sigmas (`bool`, defaults to False): Whether to use beta sigmas for step sizes in the noise schedule during sampling. time_shift_type (`str`, defaults to "exponential"): - The type of dynamic resolution-dependent timestep shifting to apply. Either "exponential" or "linear". - scale_factors ('list', defaults to None) + The type of dynamic resolution-dependent timestep shifting to apply. + scale_factors (`list[float]`, *optional*, defaults to `None`): It defines how to scale the latents at which predictions are made. - upscale_mode ('str', defaults to 'bicubic') - Upscaling method, applied if scale-wise generation is considered + upscale_mode (`str`, *optional*, defaults to "bicubic"): + Upscaling method, applied if scale-wise generation is considered. """ _compatibles = [] @@ -96,22 +96,39 @@ def __init__( num_train_timesteps: int = 1000, shift: float = 1.0, use_dynamic_shifting: bool = False, - base_shift: Optional[float] = 0.5, - max_shift: Optional[float] = 1.15, - base_image_seq_len: Optional[int] = 256, - max_image_seq_len: Optional[int] = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, + base_image_seq_len: int = 256, + max_image_seq_len: int = 4096, invert_sigmas: bool = False, - shift_terminal: Optional[float] = None, - use_karras_sigmas: Optional[bool] = False, - use_exponential_sigmas: Optional[bool] = False, - use_beta_sigmas: Optional[bool] = False, - time_shift_type: str = "exponential", - scale_factors: Optional[List[float]] = None, - upscale_mode: Optional[str] = "bicubic", + shift_terminal: float | None = None, + use_karras_sigmas: bool | None = False, + use_exponential_sigmas: bool | None = False, + use_beta_sigmas: bool | None = False, + time_shift_type: Literal["exponential", "linear"] = "exponential", + scale_factors: list[float] | None = None, + upscale_mode: Literal[ + "nearest", + "linear", + "bilinear", + "bicubic", + "trilinear", + "area", + "nearest-exact", + ] = "bicubic", ): if self.config.use_beta_sigmas and not is_scipy_available(): raise ImportError("Make sure to install scipy if you want to use beta sigmas.") - if sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1: + if ( + sum( + [ + self.config.use_beta_sigmas, + self.config.use_exponential_sigmas, + self.config.use_karras_sigmas, + ] + ) + > 1 + ): raise ValueError( "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." ) @@ -163,7 +180,7 @@ def begin_index(self): return self._begin_index # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index - def set_begin_index(self, begin_index: int = 0): + def set_begin_index(self, begin_index: int = 0) -> None: """ Sets the begin index for the scheduler. This function should be run from pipeline before the inference. @@ -173,18 +190,18 @@ def set_begin_index(self, begin_index: int = 0): """ self._begin_index = begin_index - def set_shift(self, shift: float): + def set_shift(self, shift: float) -> None: self._shift = shift - def set_scale_factors(self, scale_factors: list, upscale_mode): + def set_scale_factors(self, scale_factors: list[float], upscale_mode: str) -> None: """ Sets scale factors for a scale-wise generation regime. Args: - scale_factors (`list`): - The scale factors for each step + scale_factors (`list[float]`): + The scale factors for each step. upscale_mode (`str`): - Upscaling method + Upscaling method. """ self._scale_factors = scale_factors self._upscale_mode = upscale_mode @@ -192,8 +209,8 @@ def set_scale_factors(self, scale_factors: list, upscale_mode): def scale_noise( self, sample: torch.FloatTensor, - timestep: Union[float, torch.FloatTensor], - noise: Optional[torch.FloatTensor] = None, + timestep: float | torch.FloatTensor, + noise: torch.FloatTensor | None = None, ) -> torch.FloatTensor: """ Forward process in flow-matching @@ -201,8 +218,10 @@ def scale_noise( Args: sample (`torch.FloatTensor`): The input sample. - timestep (`int`, *optional*): + timestep (`torch.FloatTensor`): The current timestep in the diffusion chain. + noise (`torch.FloatTensor`): + The noise tensor. Returns: `torch.FloatTensor`: @@ -237,16 +256,18 @@ def scale_noise( return sample - def _sigma_to_t(self, sigma): + def _sigma_to_t(self, sigma: float | torch.FloatTensor) -> float | torch.FloatTensor: return sigma * self.config.num_train_timesteps - def time_shift(self, mu: float, sigma: float, t: torch.Tensor): + def time_shift( + self, mu: float, sigma: float, t: float | np.ndarray | torch.Tensor + ) -> float | np.ndarray | torch.Tensor: if self.config.time_shift_type == "exponential": return self._time_shift_exponential(mu, sigma, t) elif self.config.time_shift_type == "linear": return self._time_shift_linear(mu, sigma, t) - def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor: + def stretch_shift_to_terminal(self, t: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor: r""" Stretches and shifts the timestep schedule to ensure it terminates at the configured `shift_terminal` config value. @@ -255,12 +276,13 @@ def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor: https://github.com/Lightricks/LTX-Video/blob/a01a171f8fe3d99dce2728d60a73fecf4d4238ae/ltx_video/schedulers/rf.py#L51 Args: - t (`torch.Tensor`): - A tensor of timesteps to be stretched and shifted. + t (`torch.Tensor` or `np.ndarray`): + A tensor or numpy array of timesteps to be stretched and shifted. Returns: - `torch.Tensor`: - A tensor of adjusted timesteps such that the final value equals `self.config.shift_terminal`. + `torch.Tensor` or `np.ndarray`: + A tensor or numpy array of adjusted timesteps such that the final value equals + `self.config.shift_terminal`. """ one_minus_z = 1 - t scale_factor = one_minus_z[-1] / (1 - self.config.shift_terminal) @@ -269,12 +291,12 @@ def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor: def set_timesteps( self, - num_inference_steps: Optional[int] = None, - device: Union[str, torch.device] = None, - sigmas: Optional[List[float]] = None, - mu: Optional[float] = None, - timesteps: Optional[List[float]] = None, - ): + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + sigmas: list[float] | None = None, + mu: float | None = None, + timesteps: list[float] | None = None, + ) -> None: """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -283,13 +305,13 @@ def set_timesteps( The number of diffusion steps used when generating samples with a pre-trained model. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - sigmas (`List[float]`, *optional*): + sigmas (`list[float]`, *optional*): Custom values for sigmas to be used for each diffusion step. If `None`, the sigmas are computed automatically. mu (`float`, *optional*): Determines the amount of shifting applied to sigmas when performing resolution-dependent timestep shifting. - timesteps (`List[float]`, *optional*): + timesteps (`list[float]`, *optional*): Custom values for timesteps to be used for each diffusion step. If `None`, the timesteps are computed automatically. """ @@ -316,43 +338,45 @@ def set_timesteps( is_timesteps_provided = timesteps is not None if is_timesteps_provided: - timesteps = np.array(timesteps).astype(np.float32) + timesteps = np.array(timesteps).astype(np.float32) # type: ignore if sigmas is None: if timesteps is None: - timesteps = np.linspace( - self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps + timesteps = np.linspace( # type: ignore + self._sigma_to_t(self.sigma_max), + self._sigma_to_t(self.sigma_min), + num_inference_steps, ) - sigmas = timesteps / self.config.num_train_timesteps + sigmas = timesteps / self.config.num_train_timesteps # type: ignore else: - sigmas = np.array(sigmas).astype(np.float32) + sigmas = np.array(sigmas).astype(np.float32) # type: ignore num_inference_steps = len(sigmas) # 2. Perform timestep shifting. Either no shifting is applied, or resolution-dependent shifting of # "exponential" or "linear" type is applied if self.config.use_dynamic_shifting: - sigmas = self.time_shift(mu, 1.0, sigmas) + sigmas = self.time_shift(mu, 1.0, sigmas) # type: ignore else: - sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas) + sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas) # type: ignore # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value if self.config.shift_terminal: - sigmas = self.stretch_shift_to_terminal(sigmas) + sigmas = self.stretch_shift_to_terminal(sigmas) # type: ignore # 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules if self.config.use_karras_sigmas: - sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) + sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) # type: ignore elif self.config.use_exponential_sigmas: - sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) + sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) # type: ignore elif self.config.use_beta_sigmas: - sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) + sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) # type: ignore # 5. Convert sigmas and timesteps to tensors and move to specified device - sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device) + sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device) # type: ignore if not is_timesteps_provided: - timesteps = sigmas * self.config.num_train_timesteps + timesteps = sigmas * self.config.num_train_timesteps # type: ignore else: - timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32, device=device) + timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32, device=device) # type: ignore # 6. Append the terminal sigma value. # If a model requires inverted sigma schedule for denoising but timesteps without inversion, the @@ -369,7 +393,11 @@ def set_timesteps( self._step_index = None self._begin_index = None - def index_for_timestep(self, timestep, schedule_timesteps=None): + def index_for_timestep( + self, + timestep: float | torch.Tensor, + schedule_timesteps: torch.Tensor | None = None, + ) -> int: if schedule_timesteps is None: schedule_timesteps = self.timesteps @@ -381,9 +409,9 @@ def index_for_timestep(self, timestep, schedule_timesteps=None): # case we start in the middle of the denoising schedule (e.g. for image-to-image) pos = 1 if len(indices) > 1 else 0 - return indices[pos].item() + return int(indices[pos].item()) - def _init_step_index(self, timestep): + def _init_step_index(self, timestep: float | torch.Tensor) -> None: if self.begin_index is None: if isinstance(timestep, torch.Tensor): timestep = timestep.to(self.timesteps.device) @@ -394,11 +422,11 @@ def _init_step_index(self, timestep): def step( self, model_output: torch.FloatTensor, - timestep: Union[float, torch.FloatTensor], + timestep: float | torch.FloatTensor, sample: torch.FloatTensor, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, return_dict: bool = True, - ) -> Union[FlowMatchLCMSchedulerOutput, Tuple]: + ) -> FlowMatchLCMSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). @@ -458,7 +486,12 @@ def step( size = [round(self._scale_factors[self._step_index] * size) for size in self._init_size] x0_pred = torch.nn.functional.interpolate(x0_pred, size=size, mode=self._upscale_mode) - noise = randn_tensor(x0_pred.shape, generator=generator, device=x0_pred.device, dtype=x0_pred.dtype) + noise = randn_tensor( + x0_pred.shape, + generator=generator, + device=x0_pred.device, + dtype=x0_pred.dtype, + ) prev_sample = (1 - sigma_next) * x0_pred + sigma_next * noise # upon completion increase step index by one @@ -472,7 +505,7 @@ def step( return FlowMatchLCMSchedulerOutput(prev_sample=prev_sample) # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras - def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor: + def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: """ Construct the noise schedule as proposed in [Elucidating the Design Space of Diffusion-Based Generative Models](https://huggingface.co/papers/2206.00364). @@ -593,11 +626,15 @@ def _convert_to_beta( ) return sigmas - def _time_shift_exponential(self, mu, sigma, t): + def _time_shift_exponential( + self, mu: float, sigma: float, t: float | np.ndarray | torch.Tensor + ) -> float | np.ndarray | torch.Tensor: return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) - def _time_shift_linear(self, mu, sigma, t): + def _time_shift_linear( + self, mu: float, sigma: float, t: float | np.ndarray | torch.Tensor + ) -> float | np.ndarray | torch.Tensor: return mu / (mu + (1 / t - 1) ** sigma) - def __len__(self): + def __len__(self) -> int: return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_helios.py b/src/diffusers/schedulers/scheduling_helios.py new file mode 100644 index 000000000000..ed35245c9db3 --- /dev/null +++ b/src/diffusers/schedulers/scheduling_helios.py @@ -0,0 +1,867 @@ +# Copyright 2025 The Helios Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from dataclasses import dataclass +from typing import Literal + +import numpy as np +import torch + +from ..configuration_utils import ConfigMixin, register_to_config +from ..schedulers.scheduling_utils import SchedulerMixin +from ..utils import BaseOutput, deprecate + + +@dataclass +class HeliosSchedulerOutput(BaseOutput): + prev_sample: torch.FloatTensor + model_outputs: torch.FloatTensor | None = None + last_sample: torch.FloatTensor | None = None + this_order: int | None = None + + +class HeliosScheduler(SchedulerMixin, ConfigMixin): + _compatibles = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 1.0, # Following Stable diffusion 3, + stages: int = 3, + stage_range: list = [0, 1 / 3, 2 / 3, 1], + gamma: float = 1 / 3, + # For UniPC + thresholding: bool = False, + prediction_type: str = "flow_prediction", + solver_order: int = 2, + predict_x0: bool = True, + solver_type: str = "bh2", + lower_order_final: bool = True, + disable_corrector: list[int] = [], + solver_p: SchedulerMixin = None, + use_flow_sigmas: bool = True, + scheduler_type: str = "unipc", # ["euler", "unipc"] + use_dynamic_shifting: bool = False, + time_shift_type: Literal["exponential", "linear"] = "exponential", + ): + self.timestep_ratios = {} # The timestep ratio for each stage + self.timesteps_per_stage = {} # The detailed timesteps per stage (fix max and min per stage) + self.sigmas_per_stage = {} # always uniform [1000, 0] + self.start_sigmas = {} # for start point / upsample renoise + self.end_sigmas = {} # for end point + self.ori_start_sigmas = {} + + # self.init_sigmas() + self.init_sigmas_for_each_stage() + self.sigma_min = self.sigmas[-1].item() + self.sigma_max = self.sigmas[0].item() + self.gamma = gamma + + if solver_type not in ["bh1", "bh2"]: + if solver_type in ["midpoint", "heun", "logrho"]: + self.register_to_config(solver_type="bh2") + else: + raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}") + + self.predict_x0 = predict_x0 + self.model_outputs = [None] * solver_order + self.timestep_list = [None] * solver_order + self.lower_order_nums = 0 + self.disable_corrector = disable_corrector + self.solver_p = solver_p + self.last_sample = None + self._step_index = None + self._begin_index = None + + def init_sigmas(self): + """ + initialize the global timesteps and sigmas + """ + num_train_timesteps = self.config.num_train_timesteps + shift = self.config.shift + + alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps + 1) + sigmas = 1.0 - alphas + sigmas = np.flip(shift * sigmas / (1 + (shift - 1) * sigmas))[:-1].copy() + sigmas = torch.from_numpy(sigmas) + timesteps = (sigmas * num_train_timesteps).clone() + + self._step_index = None + self._begin_index = None + self.timesteps = timesteps + self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication + + def init_sigmas_for_each_stage(self): + """ + Init the timesteps for each stage + """ + self.init_sigmas() + + stage_distance = [] + stages = self.config.stages + training_steps = self.config.num_train_timesteps + stage_range = self.config.stage_range + + # Init the start and end point of each stage + for i_s in range(stages): + # To decide the start and ends point + start_indice = int(stage_range[i_s] * training_steps) + start_indice = max(start_indice, 0) + end_indice = int(stage_range[i_s + 1] * training_steps) + end_indice = min(end_indice, training_steps) + start_sigma = self.sigmas[start_indice].item() + end_sigma = self.sigmas[end_indice].item() if end_indice < training_steps else 0.0 + self.ori_start_sigmas[i_s] = start_sigma + + if i_s != 0: + ori_sigma = 1 - start_sigma + gamma = self.config.gamma + corrected_sigma = (1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma)) * ori_sigma + # corrected_sigma = 1 / (2 - ori_sigma) * ori_sigma + start_sigma = 1 - corrected_sigma + + stage_distance.append(start_sigma - end_sigma) + self.start_sigmas[i_s] = start_sigma + self.end_sigmas[i_s] = end_sigma + + # Determine the ratio of each stage according to flow length + tot_distance = sum(stage_distance) + for i_s in range(stages): + if i_s == 0: + start_ratio = 0.0 + else: + start_ratio = sum(stage_distance[:i_s]) / tot_distance + if i_s == stages - 1: + end_ratio = 0.9999999999999999 + else: + end_ratio = sum(stage_distance[: i_s + 1]) / tot_distance + + self.timestep_ratios[i_s] = (start_ratio, end_ratio) + + # Determine the timesteps and sigmas for each stage + for i_s in range(stages): + timestep_ratio = self.timestep_ratios[i_s] + # timestep_max = self.timesteps[int(timestep_ratio[0] * training_steps)] + timestep_max = min(self.timesteps[int(timestep_ratio[0] * training_steps)], 999) + timestep_min = self.timesteps[min(int(timestep_ratio[1] * training_steps), training_steps - 1)] + timesteps = np.linspace(timestep_max, timestep_min, training_steps + 1) + self.timesteps_per_stage[i_s] = ( + timesteps[:-1] if isinstance(timesteps, torch.Tensor) else torch.from_numpy(timesteps[:-1]) + ) + stage_sigmas = np.linspace(0.999, 0, training_steps + 1) + self.sigmas_per_stage[i_s] = torch.from_numpy(stage_sigmas[:-1]) + + @property + def step_index(self): + """ + The index counter for current timestep. It will increase 1 after each scheduler step. + """ + return self._step_index + + @property + def begin_index(self): + """ + The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + """ + return self._begin_index + + def set_begin_index(self, begin_index: int = 0): + """ + Sets the begin index for the scheduler. This function should be run from pipeline before the inference. + + Args: + begin_index (`int`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + def _sigma_to_t(self, sigma): + return sigma * self.config.num_train_timesteps + + def set_timesteps( + self, + num_inference_steps: int, + stage_index: int | None = None, + device: str | torch.device = None, + sigmas: bool | None = None, + mu: bool | None = None, + is_amplify_first_chunk: bool = False, + ): + """ + Setting the timesteps and sigmas for each stage + """ + if self.config.scheduler_type == "dmd": + if is_amplify_first_chunk: + num_inference_steps = num_inference_steps * 2 + 1 + else: + num_inference_steps = num_inference_steps + 1 + + self.num_inference_steps = num_inference_steps + self.init_sigmas() + + if self.config.stages == 1: + if sigmas is None: + sigmas = np.linspace(1, 1 / self.config.num_train_timesteps, num_inference_steps + 1)[:-1].astype( + np.float32 + ) + if self.config.shift != 1.0: + assert not self.config.use_dynamic_shifting + sigmas = self.time_shift(self.config.shift, 1.0, sigmas) + timesteps = (sigmas * self.config.num_train_timesteps).copy() + sigmas = torch.from_numpy(sigmas) + else: + stage_timesteps = self.timesteps_per_stage[stage_index] + timesteps = np.linspace( + stage_timesteps[0].item(), + stage_timesteps[-1].item(), + num_inference_steps, + ) + + stage_sigmas = self.sigmas_per_stage[stage_index] + ratios = np.linspace(stage_sigmas[0].item(), stage_sigmas[-1].item(), num_inference_steps) + sigmas = torch.from_numpy(ratios) + + self.timesteps = torch.from_numpy(timesteps).to(device=device) + self.sigmas = torch.cat([sigmas, torch.zeros(1)]).to(device=device) + + self._step_index = None + self.reset_scheduler_history() + + if self.config.scheduler_type == "dmd": + self.timesteps = self.timesteps[:-1] + self.sigmas = torch.cat([self.sigmas[:-2], self.sigmas[-1:]]) + + if self.config.use_dynamic_shifting: + assert self.config.shift == 1.0 + self.sigmas = self.time_shift(mu, 1.0, self.sigmas) + if self.config.stages == 1: + self.timesteps = self.sigmas[:-1] * self.config.num_train_timesteps + else: + self.timesteps = self.timesteps_per_stage[stage_index].min() + self.sigmas[:-1] * ( + self.timesteps_per_stage[stage_index].max() - self.timesteps_per_stage[stage_index].min() + ) + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.time_shift + def time_shift(self, mu: float, sigma: float, t: torch.Tensor): + """ + Apply time shifting to the sigmas. + + Args: + mu (`float`): + The mu parameter for the time shift. + sigma (`float`): + The sigma parameter for the time shift. + t (`torch.Tensor`): + The input timesteps. + + Returns: + `torch.Tensor`: + The time-shifted timesteps. + """ + if self.config.time_shift_type == "exponential": + return self._time_shift_exponential(mu, sigma, t) + elif self.config.time_shift_type == "linear": + return self._time_shift_linear(mu, sigma, t) + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._time_shift_exponential + def _time_shift_exponential(self, mu, sigma, t): + return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._time_shift_linear + def _time_shift_linear(self, mu, sigma, t): + return mu / (mu + (1 / t - 1) ** sigma) + + # ---------------------------------- Euler ---------------------------------- + def index_for_timestep(self, timestep, schedule_timesteps=None): + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + pos = 1 if len(indices) > 1 else 0 + + return indices[pos].item() + + def _init_step_index(self, timestep): + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def step_euler( + self, + model_output: torch.FloatTensor, + timestep: float | torch.FloatTensor = None, + sample: torch.FloatTensor = None, + generator: torch.Generator | None = None, + sigma: torch.FloatTensor | None = None, + sigma_next: torch.FloatTensor | None = None, + return_dict: bool = True, + ) -> HeliosSchedulerOutput | tuple: + assert (sigma is None) == (sigma_next is None), "sigma and sigma_next must both be None or both be not None" + + if sigma is None and sigma_next is None: + if ( + isinstance(timestep, int) + or isinstance(timestep, torch.IntTensor) + or isinstance(timestep, torch.LongTensor) + ): + raise ValueError( + ( + "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass" + " one of the `scheduler.timesteps` as a timestep." + ), + ) + + if self.step_index is None: + self._step_index = 0 + + # Upcast to avoid precision issues when computing prev_sample + sample = sample.to(torch.float32) + + if sigma is None and sigma_next is None: + sigma = self.sigmas[self.step_index] + sigma_next = self.sigmas[self.step_index + 1] + + prev_sample = sample + (sigma_next - sigma) * model_output + + # Cast sample back to model compatible dtype + prev_sample = prev_sample.to(model_output.dtype) + + # upon completion increase step index by one + self._step_index += 1 + + if not return_dict: + return (prev_sample,) + + return HeliosSchedulerOutput(prev_sample=prev_sample) + + # ---------------------------------- UniPC ---------------------------------- + def _sigma_to_alpha_sigma_t(self, sigma): + if self.config.use_flow_sigmas: + alpha_t = 1 - sigma + sigma_t = torch.clamp(sigma, min=1e-8) + else: + alpha_t = 1 / ((sigma**2 + 1) ** 0.5) + sigma_t = sigma * alpha_t + + return alpha_t, sigma_t + + def convert_model_output( + self, + model_output: torch.Tensor, + *args, + sample: torch.Tensor = None, + sigma: torch.Tensor = None, + **kwargs, + ) -> torch.Tensor: + r""" + Convert the model output to the corresponding type the UniPC algorithm needs. + + Args: + model_output (`torch.Tensor`): + The direct output from the learned diffusion model. + timestep (`int`): + The current discrete timestep in the diffusion chain. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + + Returns: + `torch.Tensor`: + The converted model output. + """ + timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None) + if sample is None: + if len(args) > 1: + sample = args[1] + else: + raise ValueError("missing `sample` as a required keyword argument") + if timestep is not None: + deprecate( + "timesteps", + "1.0.0", + "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", + ) + + flag = False + if sigma is None: + flag = True + sigma = self.sigmas[self.step_index] + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) + + if self.predict_x0: + if self.config.prediction_type == "epsilon": + x0_pred = (sample - sigma_t * model_output) / alpha_t + elif self.config.prediction_type == "sample": + x0_pred = model_output + elif self.config.prediction_type == "v_prediction": + x0_pred = alpha_t * sample - sigma_t * model_output + elif self.config.prediction_type == "flow_prediction": + if flag: + sigma_t = self.sigmas[self.step_index] + else: + sigma_t = sigma + x0_pred = sample - sigma_t * model_output + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, " + "`v_prediction`, or `flow_prediction` for the UniPCMultistepScheduler." + ) + + if self.config.thresholding: + x0_pred = self._threshold_sample(x0_pred) + + return x0_pred + else: + if self.config.prediction_type == "epsilon": + return model_output + elif self.config.prediction_type == "sample": + epsilon = (sample - alpha_t * model_output) / sigma_t + return epsilon + elif self.config.prediction_type == "v_prediction": + epsilon = alpha_t * model_output + sigma_t * sample + return epsilon + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" + " `v_prediction` for the UniPCMultistepScheduler." + ) + + def multistep_uni_p_bh_update( + self, + model_output: torch.Tensor, + *args, + sample: torch.Tensor = None, + order: int = None, + sigma: torch.Tensor = None, + sigma_next: torch.Tensor = None, + **kwargs, + ) -> torch.Tensor: + """ + One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified. + + Args: + model_output (`torch.Tensor`): + The direct output from the learned diffusion model at the current timestep. + prev_timestep (`int`): + The previous discrete timestep in the diffusion chain. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + order (`int`): + The order of UniP at this timestep (corresponds to the *p* in UniPC-p). + + Returns: + `torch.Tensor`: + The sample tensor at the previous timestep. + """ + prev_timestep = args[0] if len(args) > 0 else kwargs.pop("prev_timestep", None) + if sample is None: + if len(args) > 1: + sample = args[1] + else: + raise ValueError("missing `sample` as a required keyword argument") + if order is None: + if len(args) > 2: + order = args[2] + else: + raise ValueError("missing `order` as a required keyword argument") + if prev_timestep is not None: + deprecate( + "prev_timestep", + "1.0.0", + "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", + ) + model_output_list = self.model_outputs + + s0 = self.timestep_list[-1] + m0 = model_output_list[-1] + x = sample + + if self.solver_p: + x_t = self.solver_p.step(model_output, s0, x).prev_sample + return x_t + + if sigma_next is None and sigma is None: + sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[self.step_index] + else: + sigma_t, sigma_s0 = sigma_next, sigma + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + + h = lambda_t - lambda_s0 + device = sample.device + + rks = [] + D1s = [] + for i in range(1, order): + si = self.step_index - i + mi = model_output_list[-(i + 1)] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) + + rks.append(1.0) + rks = torch.tensor(rks, device=device) + + R = [] + b = [] + + hh = -h if self.predict_x0 else h + h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1 + h_phi_k = h_phi_1 / hh - 1 + + factorial_i = 1 + + if self.config.solver_type == "bh1": + B_h = hh + elif self.config.solver_type == "bh2": + B_h = torch.expm1(hh) + else: + raise NotImplementedError() + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R = torch.stack(R) + b = torch.tensor(b, device=device) + + if len(D1s) > 0: + D1s = torch.stack(D1s, dim=1) # (B, K) + # for order 2, we use a simplified version + if order == 2: + rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device) + else: + rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype) + else: + D1s = None + + if self.predict_x0: + x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + if D1s is not None: + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) + else: + pred_res = 0 + x_t = x_t_ - alpha_t * B_h * pred_res + else: + x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 + if D1s is not None: + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) + else: + pred_res = 0 + x_t = x_t_ - sigma_t * B_h * pred_res + + x_t = x_t.to(x.dtype) + return x_t + + def multistep_uni_c_bh_update( + self, + this_model_output: torch.Tensor, + *args, + last_sample: torch.Tensor = None, + this_sample: torch.Tensor = None, + order: int = None, + sigma_before: torch.Tensor = None, + sigma: torch.Tensor = None, + **kwargs, + ) -> torch.Tensor: + """ + One step for the UniC (B(h) version). + + Args: + this_model_output (`torch.Tensor`): + The model outputs at `x_t`. + this_timestep (`int`): + The current timestep `t`. + last_sample (`torch.Tensor`): + The generated sample before the last predictor `x_{t-1}`. + this_sample (`torch.Tensor`): + The generated sample after the last predictor `x_{t}`. + order (`int`): + The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`. + + Returns: + `torch.Tensor`: + The corrected sample tensor at the current timestep. + """ + this_timestep = args[0] if len(args) > 0 else kwargs.pop("this_timestep", None) + if last_sample is None: + if len(args) > 1: + last_sample = args[1] + else: + raise ValueError("missing `last_sample` as a required keyword argument") + if this_sample is None: + if len(args) > 2: + this_sample = args[2] + else: + raise ValueError("missing `this_sample` as a required keyword argument") + if order is None: + if len(args) > 3: + order = args[3] + else: + raise ValueError("missing `order` as a required keyword argument") + if this_timestep is not None: + deprecate( + "this_timestep", + "1.0.0", + "Passing `this_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", + ) + + model_output_list = self.model_outputs + + m0 = model_output_list[-1] + x = last_sample + x_t = this_sample + model_t = this_model_output + + if sigma_before is None and sigma is None: + sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[self.step_index - 1] + else: + sigma_t, sigma_s0 = sigma, sigma_before + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + + h = lambda_t - lambda_s0 + device = this_sample.device + + rks = [] + D1s = [] + for i in range(1, order): + si = self.step_index - (i + 1) + mi = model_output_list[-(i + 1)] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) + + rks.append(1.0) + rks = torch.tensor(rks, device=device) + + R = [] + b = [] + + hh = -h if self.predict_x0 else h + h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1 + h_phi_k = h_phi_1 / hh - 1 + + factorial_i = 1 + + if self.config.solver_type == "bh1": + B_h = hh + elif self.config.solver_type == "bh2": + B_h = torch.expm1(hh) + else: + raise NotImplementedError() + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R = torch.stack(R) + b = torch.tensor(b, device=device) + + if len(D1s) > 0: + D1s = torch.stack(D1s, dim=1) + else: + D1s = None + + # for order 1, we use a simplified version + if order == 1: + rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device) + else: + rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype) + + if self.predict_x0: + x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + if D1s is not None: + corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) + else: + corr_res = 0 + D1_t = model_t - m0 + x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t) + else: + x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 + if D1s is not None: + corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) + else: + corr_res = 0 + D1_t = model_t - m0 + x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t) + x_t = x_t.to(x.dtype) + return x_t + + def step_unipc( + self, + model_output: torch.Tensor, + timestep: int | torch.Tensor = None, + sample: torch.Tensor = None, + return_dict: bool = True, + model_outputs: list = None, + timestep_list: list = None, + sigma_before: torch.Tensor = None, + sigma: torch.Tensor = None, + sigma_next: torch.Tensor = None, + cus_step_index: int = None, + cus_lower_order_num: int = None, + cus_this_order: int = None, + cus_last_sample: torch.Tensor = None, + ) -> HeliosSchedulerOutput | tuple: + if self.num_inference_steps is None: + raise ValueError( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + ) + + if cus_step_index is None: + if self.step_index is None: + self._step_index = 0 + else: + self._step_index = cus_step_index + + if cus_lower_order_num is not None: + self.lower_order_nums = cus_lower_order_num + + if cus_this_order is not None: + self.this_order = cus_this_order + + if cus_last_sample is not None: + self.last_sample = cus_last_sample + + use_corrector = ( + self.step_index > 0 and self.step_index - 1 not in self.disable_corrector and self.last_sample is not None + ) + + # Convert model output using the proper conversion method + model_output_convert = self.convert_model_output(model_output, sample=sample, sigma=sigma) + + if model_outputs is not None and timestep_list is not None: + self.model_outputs = model_outputs[:-1] + self.timestep_list = timestep_list[:-1] + + if use_corrector: + sample = self.multistep_uni_c_bh_update( + this_model_output=model_output_convert, + last_sample=self.last_sample, + this_sample=sample, + order=self.this_order, + sigma_before=sigma_before, + sigma=sigma, + ) + + if model_outputs is not None and timestep_list is not None: + model_outputs[-1] = model_output_convert + self.model_outputs = model_outputs[1:] + self.timestep_list = timestep_list[1:] + else: + for i in range(self.config.solver_order - 1): + self.model_outputs[i] = self.model_outputs[i + 1] + self.timestep_list[i] = self.timestep_list[i + 1] + self.model_outputs[-1] = model_output_convert + self.timestep_list[-1] = timestep + + if self.config.lower_order_final: + this_order = min(self.config.solver_order, len(self.timesteps) - self.step_index) + else: + this_order = self.config.solver_order + self.this_order = min(this_order, self.lower_order_nums + 1) # warmup for multistep + assert self.this_order > 0 + + self.last_sample = sample + prev_sample = self.multistep_uni_p_bh_update( + model_output=model_output, # pass the original non-converted model output, in case solver-p is used + sample=sample, + order=self.this_order, + sigma=sigma, + sigma_next=sigma_next, + ) + + if cus_lower_order_num is None: + if self.lower_order_nums < self.config.solver_order: + self.lower_order_nums += 1 + + # upon completion increase step index by one + if cus_step_index is None: + self._step_index += 1 + + if not return_dict: + return (prev_sample, model_outputs, self.last_sample, self.this_order) + + return HeliosSchedulerOutput( + prev_sample=prev_sample, + model_outputs=model_outputs, + last_sample=self.last_sample, + this_order=self.this_order, + ) + + # ---------------------------------- Merge ---------------------------------- + def step( + self, + model_output: torch.FloatTensor, + timestep: float | torch.FloatTensor = None, + sample: torch.FloatTensor = None, + generator: torch.Generator | None = None, + return_dict: bool = True, + ) -> HeliosSchedulerOutput | tuple: + if self.config.scheduler_type == "euler": + return self.step_euler( + model_output=model_output, + timestep=timestep, + sample=sample, + generator=generator, + return_dict=return_dict, + ) + elif self.config.scheduler_type == "unipc": + return self.step_unipc( + model_output=model_output, + timestep=timestep, + sample=sample, + return_dict=return_dict, + ) + else: + raise NotImplementedError + + def reset_scheduler_history(self): + self.model_outputs = [None] * self.config.solver_order + self.timestep_list = [None] * self.config.solver_order + self.lower_order_nums = 0 + self.disable_corrector = self.config.disable_corrector + self.solver_p = self.config.solver_p + self.last_sample = None + self._step_index = None + self._begin_index = None + + def __len__(self): + return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_helios_dmd.py b/src/diffusers/schedulers/scheduling_helios_dmd.py new file mode 100644 index 000000000000..1f4afa0e3128 --- /dev/null +++ b/src/diffusers/schedulers/scheduling_helios_dmd.py @@ -0,0 +1,331 @@ +# Copyright 2025 The Helios Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from dataclasses import dataclass +from typing import Literal + +import numpy as np +import torch + +from ..configuration_utils import ConfigMixin, register_to_config +from ..schedulers.scheduling_utils import SchedulerMixin +from ..utils import BaseOutput + + +@dataclass +class HeliosDMDSchedulerOutput(BaseOutput): + prev_sample: torch.FloatTensor + model_outputs: torch.FloatTensor | None = None + last_sample: torch.FloatTensor | None = None + this_order: int | None = None + + +class HeliosDMDScheduler(SchedulerMixin, ConfigMixin): + _compatibles = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 1.0, # Following Stable diffusion 3, + stages: int = 3, + stage_range: list = [0, 1 / 3, 2 / 3, 1], + gamma: float = 1 / 3, + prediction_type: str = "flow_prediction", + use_flow_sigmas: bool = True, + use_dynamic_shifting: bool = False, + time_shift_type: Literal["exponential", "linear"] = "linear", + ): + self.timestep_ratios = {} # The timestep ratio for each stage + self.timesteps_per_stage = {} # The detailed timesteps per stage (fix max and min per stage) + self.sigmas_per_stage = {} # always uniform [1000, 0] + self.start_sigmas = {} # for start point / upsample renoise + self.end_sigmas = {} # for end point + self.ori_start_sigmas = {} + + # self.init_sigmas() + self.init_sigmas_for_each_stage() + self.sigma_min = self.sigmas[-1].item() + self.sigma_max = self.sigmas[0].item() + self.gamma = gamma + + self.last_sample = None + self._step_index = None + self._begin_index = None + + def init_sigmas(self): + """ + initialize the global timesteps and sigmas + """ + num_train_timesteps = self.config.num_train_timesteps + shift = self.config.shift + + alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps + 1) + sigmas = 1.0 - alphas + sigmas = np.flip(shift * sigmas / (1 + (shift - 1) * sigmas))[:-1].copy() + sigmas = torch.from_numpy(sigmas) + timesteps = (sigmas * num_train_timesteps).clone() + + self._step_index = None + self._begin_index = None + self.timesteps = timesteps + self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication + + def init_sigmas_for_each_stage(self): + """ + Init the timesteps for each stage + """ + self.init_sigmas() + + stage_distance = [] + stages = self.config.stages + training_steps = self.config.num_train_timesteps + stage_range = self.config.stage_range + + # Init the start and end point of each stage + for i_s in range(stages): + # To decide the start and ends point + start_indice = int(stage_range[i_s] * training_steps) + start_indice = max(start_indice, 0) + end_indice = int(stage_range[i_s + 1] * training_steps) + end_indice = min(end_indice, training_steps) + start_sigma = self.sigmas[start_indice].item() + end_sigma = self.sigmas[end_indice].item() if end_indice < training_steps else 0.0 + self.ori_start_sigmas[i_s] = start_sigma + + if i_s != 0: + ori_sigma = 1 - start_sigma + gamma = self.config.gamma + corrected_sigma = (1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma)) * ori_sigma + # corrected_sigma = 1 / (2 - ori_sigma) * ori_sigma + start_sigma = 1 - corrected_sigma + + stage_distance.append(start_sigma - end_sigma) + self.start_sigmas[i_s] = start_sigma + self.end_sigmas[i_s] = end_sigma + + # Determine the ratio of each stage according to flow length + tot_distance = sum(stage_distance) + for i_s in range(stages): + if i_s == 0: + start_ratio = 0.0 + else: + start_ratio = sum(stage_distance[:i_s]) / tot_distance + if i_s == stages - 1: + end_ratio = 0.9999999999999999 + else: + end_ratio = sum(stage_distance[: i_s + 1]) / tot_distance + + self.timestep_ratios[i_s] = (start_ratio, end_ratio) + + # Determine the timesteps and sigmas for each stage + for i_s in range(stages): + timestep_ratio = self.timestep_ratios[i_s] + # timestep_max = self.timesteps[int(timestep_ratio[0] * training_steps)] + timestep_max = min(self.timesteps[int(timestep_ratio[0] * training_steps)], 999) + timestep_min = self.timesteps[min(int(timestep_ratio[1] * training_steps), training_steps - 1)] + timesteps = np.linspace(timestep_max, timestep_min, training_steps + 1) + self.timesteps_per_stage[i_s] = ( + timesteps[:-1] if isinstance(timesteps, torch.Tensor) else torch.from_numpy(timesteps[:-1]) + ) + stage_sigmas = np.linspace(0.999, 0, training_steps + 1) + self.sigmas_per_stage[i_s] = torch.from_numpy(stage_sigmas[:-1]) + + @property + def step_index(self): + """ + The index counter for current timestep. It will increase 1 after each scheduler step. + """ + return self._step_index + + @property + def begin_index(self): + """ + The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + """ + return self._begin_index + + def set_begin_index(self, begin_index: int = 0): + """ + Sets the begin index for the scheduler. This function should be run from pipeline before the inference. + + Args: + begin_index (`int`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + def _sigma_to_t(self, sigma): + return sigma * self.config.num_train_timesteps + + def set_timesteps( + self, + num_inference_steps: int, + stage_index: int | None = None, + device: str | torch.device = None, + sigmas: bool | None = None, + mu: bool | None = None, + is_amplify_first_chunk: bool = False, + ): + """ + Setting the timesteps and sigmas for each stage + """ + if is_amplify_first_chunk: + num_inference_steps = num_inference_steps * 2 + 1 + else: + num_inference_steps = num_inference_steps + 1 + + self.num_inference_steps = num_inference_steps + self.init_sigmas() + + if self.config.stages == 1: + if sigmas is None: + sigmas = np.linspace(1, 1 / self.config.num_train_timesteps, num_inference_steps + 1)[:-1].astype( + np.float32 + ) + if self.config.shift != 1.0: + assert not self.config.use_dynamic_shifting + sigmas = self.time_shift(self.config.shift, 1.0, sigmas) + timesteps = (sigmas * self.config.num_train_timesteps).copy() + sigmas = torch.from_numpy(sigmas) + else: + stage_timesteps = self.timesteps_per_stage[stage_index] + timesteps = np.linspace( + stage_timesteps[0].item(), + stage_timesteps[-1].item(), + num_inference_steps, + ) + + stage_sigmas = self.sigmas_per_stage[stage_index] + ratios = np.linspace(stage_sigmas[0].item(), stage_sigmas[-1].item(), num_inference_steps) + sigmas = torch.from_numpy(ratios) + + self.timesteps = torch.from_numpy(timesteps).to(device=device) + self.sigmas = torch.cat([sigmas, torch.zeros(1)]).to(device=device) + + self._step_index = None + self.reset_scheduler_history() + + self.timesteps = self.timesteps[:-1] + self.sigmas = torch.cat([self.sigmas[:-2], self.sigmas[-1:]]) + + if self.config.use_dynamic_shifting: + assert self.config.shift == 1.0 + self.sigmas = self.time_shift(mu, 1.0, self.sigmas) + if self.config.stages == 1: + self.timesteps = self.sigmas[:-1] * self.config.num_train_timesteps + else: + self.timesteps = self.timesteps_per_stage[stage_index].min() + self.sigmas[:-1] * ( + self.timesteps_per_stage[stage_index].max() - self.timesteps_per_stage[stage_index].min() + ) + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.time_shift + def time_shift(self, mu: float, sigma: float, t: torch.Tensor): + """ + Apply time shifting to the sigmas. + + Args: + mu (`float`): + The mu parameter for the time shift. + sigma (`float`): + The sigma parameter for the time shift. + t (`torch.Tensor`): + The input timesteps. + + Returns: + `torch.Tensor`: + The time-shifted timesteps. + """ + if self.config.time_shift_type == "exponential": + return self._time_shift_exponential(mu, sigma, t) + elif self.config.time_shift_type == "linear": + return self._time_shift_linear(mu, sigma, t) + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._time_shift_exponential + def _time_shift_exponential(self, mu, sigma, t): + return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._time_shift_linear + def _time_shift_linear(self, mu, sigma, t): + return mu / (mu + (1 / t - 1) ** sigma) + + # ---------------------------------- For DMD ---------------------------------- + def add_noise(self, original_samples, noise, timestep, sigmas, timesteps): + sigmas = sigmas.to(noise.device) + timesteps = timesteps.to(noise.device) + timestep_id = torch.argmin((timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma = sigmas[timestep_id].reshape(-1, 1, 1, 1, 1) + sample = (1 - sigma) * original_samples + sigma * noise + return sample.type_as(noise) + + def convert_flow_pred_to_x0(self, flow_pred, xt, timestep, sigmas, timesteps): + # use higher precision for calculations + original_dtype = flow_pred.dtype + device = flow_pred.device + flow_pred, xt, sigmas, timesteps = (x.double().to(device) for x in (flow_pred, xt, sigmas, timesteps)) + + timestep_id = torch.argmin((timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1, 1) + x0_pred = xt - sigma_t * flow_pred + return x0_pred.to(original_dtype) + + def step( + self, + model_output: torch.FloatTensor, + timestep: float | torch.FloatTensor = None, + sample: torch.FloatTensor = None, + generator: torch.Generator | None = None, + return_dict: bool = True, + cur_sampling_step: int = 0, + dmd_noisy_tensor: torch.FloatTensor | None = None, + dmd_sigmas: torch.FloatTensor | None = None, + dmd_timesteps: torch.FloatTensor | None = None, + all_timesteps: torch.FloatTensor | None = None, + ) -> HeliosDMDSchedulerOutput | tuple: + pred_image_or_video = self.convert_flow_pred_to_x0( + flow_pred=model_output, + xt=sample, + timestep=torch.full((model_output.shape[0],), timestep, dtype=torch.long, device=model_output.device), + sigmas=dmd_sigmas, + timesteps=dmd_timesteps, + ) + if cur_sampling_step < len(all_timesteps) - 1: + prev_sample = self.add_noise( + pred_image_or_video, + dmd_noisy_tensor, + torch.full( + (model_output.shape[0],), + all_timesteps[cur_sampling_step + 1], + dtype=torch.long, + device=model_output.device, + ), + sigmas=dmd_sigmas, + timesteps=dmd_timesteps, + ) + else: + prev_sample = pred_image_or_video + + if not return_dict: + return (prev_sample,) + + return HeliosDMDSchedulerOutput(prev_sample=prev_sample) + + def reset_scheduler_history(self): + self._step_index = None + self._begin_index = None + + def __len__(self): + return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_heun_discrete.py b/src/diffusers/schedulers/scheduling_heun_discrete.py index b113f9b49832..df2593493aeb 100644 --- a/src/diffusers/schedulers/scheduling_heun_discrete.py +++ b/src/diffusers/schedulers/scheduling_heun_discrete.py @@ -14,7 +14,7 @@ import math from dataclasses import dataclass -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -44,14 +44,14 @@ class HeunDiscreteSchedulerOutput(BaseOutput): """ prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None + pred_original_sample: torch.Tensor | None = None # Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -65,8 +65,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -77,6 +77,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -144,13 +151,13 @@ def __init__( num_train_timesteps: int = 1000, beta_start: float = 0.00085, # sensible defaults beta_end: float = 0.012, - beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2", "exp"] = "linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, - prediction_type: Literal["epsilon", "sample", "v_prediction"] = "epsilon", - use_karras_sigmas: Optional[bool] = False, - use_exponential_sigmas: Optional[bool] = False, - use_beta_sigmas: Optional[bool] = False, - clip_sample: Optional[bool] = False, + beta_schedule: str = "linear", + trained_betas: np.ndarray | list[float] | None = None, + prediction_type: str = "epsilon", + use_karras_sigmas: bool = False, + use_exponential_sigmas: bool = False, + use_beta_sigmas: bool = False, + clip_sample: bool = False, clip_sample_range: float = 1.0, timestep_spacing: Literal["linspace", "leading", "trailing"] = "linspace", steps_offset: int = 0, @@ -189,7 +196,7 @@ def __init__( # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep def index_for_timestep( - self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None + self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None ) -> int: """ Find the index of a given timestep in the timestep schedule. @@ -254,7 +261,7 @@ def set_begin_index(self, begin_index: int = 0) -> None: def scale_model_input( self, sample: torch.Tensor, - timestep: Union[float, torch.Tensor], + timestep: float | torch.Tensor, ) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the @@ -279,11 +286,11 @@ def scale_model_input( def set_timesteps( self, - num_inference_steps: Optional[int] = None, - device: Union[str, torch.device] = None, - num_train_timesteps: Optional[int] = None, - timesteps: Optional[List[int]] = None, - ) -> None: + num_inference_steps: int | None = None, + device: str | torch.device = None, + num_train_timesteps: int | None = None, + timesteps: list[int] | None = None, + ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -295,7 +302,7 @@ def set_timesteps( num_train_timesteps (`int`, *optional*, defaults to `None`): The number of diffusion steps used when training the model. If `None`, the default `num_train_timesteps` attribute is used. - timesteps (`List[int]`, *optional*, defaults to `None`): + timesteps (`list[int]`, *optional*): Custom timesteps used to support arbitrary spacing between timesteps. If `None`, timesteps will be generated based on the `timestep_spacing` attribute. If `timesteps` is passed, `num_inference_steps` must be `None`, and `timestep_spacing` attribute will be ignored. @@ -533,7 +540,7 @@ def state_in_first_order(self): return self.dt is None # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: + def _init_step_index(self, timestep: float | torch.Tensor) -> None: """ Initialize the step index for the scheduler based on the given timestep. @@ -550,11 +557,11 @@ def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: def step( self, - model_output: Union[torch.Tensor, np.ndarray], - timestep: Union[float, torch.Tensor], - sample: Union[torch.Tensor, np.ndarray], + model_output: torch.Tensor | np.ndarray, + timestep: float | torch.Tensor, + sample: torch.Tensor | np.ndarray, return_dict: bool = True, - ) -> Union[HeunDiscreteSchedulerOutput, Tuple]: + ) -> HeunDiscreteSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). diff --git a/src/diffusers/schedulers/scheduling_ipndm.py b/src/diffusers/schedulers/scheduling_ipndm.py index da188fe8297c..f2e744e0f2af 100644 --- a/src/diffusers/schedulers/scheduling_ipndm.py +++ b/src/diffusers/schedulers/scheduling_ipndm.py @@ -13,7 +13,6 @@ # limitations under the License. import math -from typing import List, Optional, Tuple, Union import numpy as np import torch @@ -32,7 +31,7 @@ class IPNDMScheduler(SchedulerMixin, ConfigMixin): Args: num_train_timesteps (`int`, defaults to 1000): The number of diffusion steps to train the model. - trained_betas (`np.ndarray`, *optional*): + trained_betas (`np.ndarray` or `List[float]`, *optional*): Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. """ @@ -40,7 +39,9 @@ class IPNDMScheduler(SchedulerMixin, ConfigMixin): @register_to_config def __init__( - self, num_train_timesteps: int = 1000, trained_betas: Optional[Union[np.ndarray, List[float]]] = None + self, + num_train_timesteps: int = 1000, + trained_betas: np.ndarray | list[float] | None = None, ): # set `betas`, `alphas`, `timesteps` self.set_timesteps(num_train_timesteps) @@ -59,21 +60,29 @@ def __init__( self._begin_index = None @property - def step_index(self): + def step_index(self) -> int | None: """ The index counter for current timestep. It will increase 1 after each scheduler step. + + Returns: + `int` or `None`: + The index counter for current timestep. """ return self._step_index @property - def begin_index(self): + def begin_index(self) -> int | None: """ The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + + Returns: + `int` or `None`: + The index for the first timestep. """ return self._begin_index # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index - def set_begin_index(self, begin_index: int = 0): + def set_begin_index(self, begin_index: int = 0) -> None: """ Sets the begin index for the scheduler. This function should be run from pipeline before the inference. @@ -83,7 +92,7 @@ def set_begin_index(self, begin_index: int = 0): """ self._begin_index = begin_index - def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None): + def set_timesteps(self, num_inference_steps: int, device: str | torch.device | None = None): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -113,7 +122,7 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep def index_for_timestep( - self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None + self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None ) -> int: """ Find the index of a given timestep in the timestep schedule. @@ -143,7 +152,7 @@ def index_for_timestep( return indices[pos].item() # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: + def _init_step_index(self, timestep: float | torch.Tensor) -> None: """ Initialize the step index for the scheduler based on the given timestep. @@ -161,10 +170,10 @@ def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: def step( self, model_output: torch.Tensor, - timestep: Union[int, torch.Tensor], + timestep: int | torch.Tensor, sample: torch.Tensor, return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: + ) -> SchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the linear multistep method. It performs one forward pass multiple times to approximate the solution. @@ -172,7 +181,7 @@ def step( Args: model_output (`torch.Tensor`): The direct output from learned diffusion model. - timestep (`int`): + timestep (`int` or `torch.Tensor`): The current discrete timestep in the diffusion chain. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. @@ -231,7 +240,30 @@ def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tens """ return sample - def _get_prev_sample(self, sample, timestep_index, prev_timestep_index, ets): + def _get_prev_sample( + self, + sample: torch.Tensor, + timestep_index: int, + prev_timestep_index: int, + ets: torch.Tensor, + ) -> torch.Tensor: + """ + Predicts the previous sample based on the current sample, timestep indices, and running model outputs. + + Args: + sample (`torch.Tensor`): + The current sample. + timestep_index (`int`): + Index of the current timestep in the schedule. + prev_timestep_index (`int`): + Index of the previous timestep in the schedule. + ets (`torch.Tensor`): + The running sequence of model outputs. + + Returns: + `torch.Tensor`: + The predicted previous sample. + """ alpha = self.alphas[timestep_index] sigma = self.betas[timestep_index] @@ -243,5 +275,5 @@ def _get_prev_sample(self, sample, timestep_index, prev_timestep_index, ets): return prev_sample - def __len__(self): + def __len__(self) -> int: return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py index da40bed635e1..706190e291a4 100644 --- a/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py +++ b/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py @@ -14,7 +14,7 @@ import math from dataclasses import dataclass -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -45,14 +45,14 @@ class KDPM2AncestralDiscreteSchedulerOutput(BaseOutput): """ prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None + pred_original_sample: torch.Tensor | None = None # Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -66,8 +66,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -78,6 +78,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -143,10 +150,10 @@ def __init__( beta_start: float = 0.00085, # sensible defaults beta_end: float = 0.012, beta_schedule: str = "linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, - use_karras_sigmas: Optional[bool] = False, - use_exponential_sigmas: Optional[bool] = False, - use_beta_sigmas: Optional[bool] = False, + trained_betas: np.ndarray | list[float] | None = None, + use_karras_sigmas: bool = False, + use_exponential_sigmas: bool = False, + use_beta_sigmas: bool = False, prediction_type: str = "epsilon", timestep_spacing: str = "linspace", steps_offset: int = 0, @@ -215,7 +222,7 @@ def set_begin_index(self, begin_index: int = 0): def scale_model_input( self, sample: torch.Tensor, - timestep: Union[float, torch.Tensor], + timestep: float | torch.Tensor, ) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the @@ -245,8 +252,8 @@ def scale_model_input( def set_timesteps( self, num_inference_steps: int, - device: Union[str, torch.device] = None, - num_train_timesteps: Optional[int] = None, + device: str | torch.device | None = None, + num_train_timesteps: int | None = None, ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -506,7 +513,7 @@ def state_in_first_order(self): # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep def index_for_timestep( - self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None + self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None ) -> int: """ Find the index of a given timestep in the timestep schedule. @@ -536,7 +543,7 @@ def index_for_timestep( return indices[pos].item() # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: + def _init_step_index(self, timestep: float | torch.Tensor) -> None: """ Initialize the step index for the scheduler based on the given timestep. @@ -553,12 +560,12 @@ def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: def step( self, - model_output: Union[torch.Tensor, np.ndarray], - timestep: Union[float, torch.Tensor], - sample: Union[torch.Tensor, np.ndarray], - generator: Optional[torch.Generator] = None, + model_output: torch.Tensor | np.ndarray, + timestep: float | torch.Tensor, + sample: torch.Tensor | np.ndarray, + generator: torch.Generator | None = None, return_dict: bool = True, - ) -> Union[KDPM2AncestralDiscreteSchedulerOutput, Tuple]: + ) -> KDPM2AncestralDiscreteSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). diff --git a/src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py b/src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py index 6dc08d4d0a86..3092c610b46d 100644 --- a/src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py +++ b/src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py @@ -14,7 +14,7 @@ import math from dataclasses import dataclass -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -44,14 +44,14 @@ class KDPM2DiscreteSchedulerOutput(BaseOutput): """ prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None + pred_original_sample: torch.Tensor | None = None # Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -65,8 +65,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -77,6 +77,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -142,10 +149,10 @@ def __init__( beta_start: float = 0.00085, # sensible defaults beta_end: float = 0.012, beta_schedule: str = "linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, - use_karras_sigmas: Optional[bool] = False, - use_exponential_sigmas: Optional[bool] = False, - use_beta_sigmas: Optional[bool] = False, + trained_betas: np.ndarray | list[float] | None = None, + use_karras_sigmas: bool = False, + use_exponential_sigmas: bool = False, + use_beta_sigmas: bool = False, prediction_type: str = "epsilon", timestep_spacing: str = "linspace", steps_offset: int = 0, @@ -215,7 +222,7 @@ def set_begin_index(self, begin_index: int = 0): def scale_model_input( self, sample: torch.Tensor, - timestep: Union[float, torch.Tensor], + timestep: float | torch.Tensor, ) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the @@ -245,8 +252,8 @@ def scale_model_input( def set_timesteps( self, num_inference_steps: int, - device: Union[str, torch.device] = None, - num_train_timesteps: Optional[int] = None, + device: str | torch.device = None, + num_train_timesteps: int = None, ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -332,7 +339,7 @@ def state_in_first_order(self): # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep def index_for_timestep( - self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None + self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None ) -> int: """ Find the index of a given timestep in the timestep schedule. @@ -362,7 +369,7 @@ def index_for_timestep( return indices[pos].item() # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: + def _init_step_index(self, timestep: float | torch.Tensor) -> None: """ Initialize the step index for the scheduler based on the given timestep. @@ -538,11 +545,11 @@ def _convert_to_beta( def step( self, - model_output: Union[torch.Tensor, np.ndarray], - timestep: Union[float, torch.Tensor], - sample: Union[torch.Tensor, np.ndarray], + model_output: torch.Tensor | np.ndarray, + timestep: float | torch.Tensor, + sample: torch.Tensor | np.ndarray, return_dict: bool = True, - ) -> Union[KDPM2DiscreteSchedulerOutput, Tuple]: + ) -> KDPM2DiscreteSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). diff --git a/src/diffusers/schedulers/scheduling_karras_ve_flax.py b/src/diffusers/schedulers/scheduling_karras_ve_flax.py index bacfbd61006d..04c46220fcca 100644 --- a/src/diffusers/schedulers/scheduling_karras_ve_flax.py +++ b/src/diffusers/schedulers/scheduling_karras_ve_flax.py @@ -14,7 +14,6 @@ from dataclasses import dataclass -from typing import Optional, Tuple, Union import flax import jax @@ -22,16 +21,19 @@ from jax import random from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput +from ..utils import BaseOutput, logging from .scheduling_utils_flax import FlaxSchedulerMixin +logger = logging.get_logger(__name__) + + @flax.struct.dataclass class KarrasVeSchedulerState: # setable values - num_inference_steps: Optional[int] = None - timesteps: Optional[jnp.ndarray] = None - schedule: Optional[jnp.ndarray] = None # sigma(t_i) + num_inference_steps: int = None + timesteps: jnp.ndarray | None = None + schedule: jnp.ndarray | None = None # sigma(t_i) @classmethod def create(cls): @@ -102,13 +104,16 @@ def __init__( s_min: float = 0.05, s_max: float = 50, ): - pass + logger.warning( + "Flax classes are deprecated and will be removed in Diffusers v1.0.0. We " + "recommend migrating to PyTorch classes or pinning your version of Diffusers." + ) def create_state(self): return KarrasVeSchedulerState.create() def set_timesteps( - self, state: KarrasVeSchedulerState, num_inference_steps: int, shape: Tuple = () + self, state: KarrasVeSchedulerState, num_inference_steps: int, shape: tuple = () ) -> KarrasVeSchedulerState: """ Sets the continuous timesteps used for the diffusion chain. Supporting function to be run before inference. @@ -141,7 +146,7 @@ def add_noise_to_input( sample: jnp.ndarray, sigma: float, key: jax.Array, - ) -> Tuple[jnp.ndarray, float]: + ) -> tuple[jnp.ndarray, float]: """ Explicit Langevin-like "churn" step of adding noise to the sample according to a factor gamma_i ≥ 0 to reach a higher noise level sigma_hat = sigma_i + gamma_i*sigma_i. @@ -169,7 +174,7 @@ def step( sigma_prev: float, sample_hat: jnp.ndarray, return_dict: bool = True, - ) -> Union[FlaxKarrasVeOutput, Tuple]: + ) -> FlaxKarrasVeOutput | tuple: """ Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion process from the learned model outputs (most often the predicted noise). @@ -207,7 +212,7 @@ def step_correct( sample_prev: jnp.ndarray, derivative: jnp.ndarray, return_dict: bool = True, - ) -> Union[FlaxKarrasVeOutput, Tuple]: + ) -> FlaxKarrasVeOutput | tuple: """ Correct the predicted sample based on the output model_output of the network. TODO complete description diff --git a/src/diffusers/schedulers/scheduling_lcm.py b/src/diffusers/schedulers/scheduling_lcm.py index 0527f3533851..e70bc4745a24 100644 --- a/src/diffusers/schedulers/scheduling_lcm.py +++ b/src/diffusers/schedulers/scheduling_lcm.py @@ -17,7 +17,7 @@ import math from dataclasses import dataclass -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -46,14 +46,14 @@ class LCMSchedulerOutput(BaseOutput): """ prev_sample: torch.Tensor - denoised: Optional[torch.Tensor] = None + denoised: torch.Tensor | None = None # Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -67,8 +67,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -79,6 +79,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -200,7 +207,7 @@ def __init__( beta_start: float = 0.00085, beta_end: float = 0.012, beta_schedule: str = "scaled_linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, + trained_betas: np.ndarray | list[float] | None = None, original_inference_steps: int = 50, clip_sample: bool = False, clip_sample_range: float = 1.0, @@ -253,7 +260,7 @@ def __init__( # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep def index_for_timestep( - self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None + self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None ) -> int: """ Find the index of a given timestep in the timestep schedule. @@ -283,7 +290,7 @@ def index_for_timestep( return indices[pos].item() # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: + def _init_step_index(self, timestep: float | torch.Tensor) -> None: """ Initialize the step index for the scheduler based on the given timestep. @@ -320,7 +327,7 @@ def set_begin_index(self, begin_index: int = 0): """ self._begin_index = begin_index - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. @@ -382,10 +389,10 @@ def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: def set_timesteps( self, - num_inference_steps: Optional[int] = None, - device: Union[str, torch.device] = None, - original_inference_steps: Optional[int] = None, - timesteps: Optional[List[int]] = None, + num_inference_steps: int | None = None, + device: str | torch.device = None, + original_inference_steps: int | None = None, + timesteps: list[int] | None = None, strength: int = 1.0, ): """ @@ -402,7 +409,7 @@ def set_timesteps( schedule (which is different from the standard `diffusers` implementation). We will then take `num_inference_steps` timesteps from this schedule, evenly spaced in terms of indices, and use that as our final timestep schedule. If not set, this will default to the `original_inference_steps` attribute. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default timestep spacing strategy of equal spacing between timesteps on the training/distillation timestep schedule is used. If `timesteps` is passed, `num_inference_steps` must be `None`. @@ -533,9 +540,9 @@ def step( model_output: torch.Tensor, timestep: int, sample: torch.Tensor, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, return_dict: bool = True, - ) -> Union[LCMSchedulerOutput, Tuple]: + ) -> LCMSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). @@ -715,7 +722,7 @@ def previous_timestep(self, timestep): The current timestep. Returns: - `int`: + `int` or `torch.Tensor`: The previous timestep. """ if self.custom_timesteps or self.num_inference_steps: diff --git a/src/diffusers/schedulers/scheduling_lms_discrete.py b/src/diffusers/schedulers/scheduling_lms_discrete.py index 276af6eeacb7..e917adda9516 100644 --- a/src/diffusers/schedulers/scheduling_lms_discrete.py +++ b/src/diffusers/schedulers/scheduling_lms_discrete.py @@ -14,18 +14,21 @@ import math import warnings from dataclasses import dataclass -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np -import scipy.stats import torch -from scipy import integrate from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput +from ..utils import BaseOutput, is_scipy_available from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin +if is_scipy_available(): + import scipy.stats + from scipy import integrate + + @dataclass # Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->LMSDiscrete class LMSDiscreteSchedulerOutput(BaseOutput): @@ -42,14 +45,14 @@ class LMSDiscreteSchedulerOutput(BaseOutput): """ prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None + pred_original_sample: torch.Tensor | None = None # Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -63,8 +66,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -75,6 +78,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -137,11 +147,11 @@ def __init__( num_train_timesteps: int = 1000, beta_start: float = 0.0001, beta_end: float = 0.02, - beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, - use_karras_sigmas: Optional[bool] = False, - use_exponential_sigmas: Optional[bool] = False, - use_beta_sigmas: Optional[bool] = False, + beta_schedule: str = "linear", + trained_betas: np.ndarray | list[float] | None = None, + use_karras_sigmas: bool = False, + use_exponential_sigmas: bool = False, + use_beta_sigmas: bool = False, prediction_type: Literal["epsilon", "sample", "v_prediction"] = "epsilon", timestep_spacing: Literal["linspace", "leading", "trailing"] = "linspace", steps_offset: int = 0, @@ -182,7 +192,7 @@ def __init__( self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication @property - def init_noise_sigma(self) -> Union[float, torch.Tensor]: + def init_noise_sigma(self) -> float | torch.Tensor: """ The standard deviation of the initial noise distribution. @@ -198,7 +208,7 @@ def init_noise_sigma(self) -> Union[float, torch.Tensor]: return (self.sigmas.max() ** 2 + 1) ** 0.5 @property - def step_index(self) -> Optional[int]: + def step_index(self) -> int: """ The index counter for current timestep. It will increase by 1 after each scheduler step. @@ -209,7 +219,7 @@ def step_index(self) -> Optional[int]: return self._step_index @property - def begin_index(self) -> Optional[int]: + def begin_index(self) -> int: """ The index for the first timestep. It should be set from pipeline with `set_begin_index` method. @@ -230,7 +240,7 @@ def set_begin_index(self, begin_index: int = 0) -> None: """ self._begin_index = begin_index - def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: float | torch.Tensor) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. @@ -283,7 +293,7 @@ def lms_derivative(tau): return integrated_coeff - def set_timesteps(self, num_inference_steps: int, device: Optional[Union[str, torch.device]] = None) -> None: + def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -343,7 +353,7 @@ def set_timesteps(self, num_inference_steps: int, device: Optional[Union[str, to # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep def index_for_timestep( - self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None + self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None ) -> int: """ Find the index of a given timestep in the timestep schedule. @@ -373,7 +383,7 @@ def index_for_timestep( return indices[pos].item() # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: + def _init_step_index(self, timestep: float | torch.Tensor) -> None: """ Initialize the step index for the scheduler based on the given timestep. @@ -535,11 +545,11 @@ def _convert_to_beta( def step( self, model_output: torch.Tensor, - timestep: Union[float, torch.Tensor], + timestep: float | torch.Tensor, sample: torch.Tensor, order: int = 4, return_dict: bool = True, - ) -> Union[LMSDiscreteSchedulerOutput, Tuple]: + ) -> LMSDiscreteSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). diff --git a/src/diffusers/schedulers/scheduling_lms_discrete_flax.py b/src/diffusers/schedulers/scheduling_lms_discrete_flax.py index 3fd4dc8a5d61..65902678e1d9 100644 --- a/src/diffusers/schedulers/scheduling_lms_discrete_flax.py +++ b/src/diffusers/schedulers/scheduling_lms_discrete_flax.py @@ -13,13 +13,13 @@ # limitations under the License. from dataclasses import dataclass -from typing import Optional, Tuple, Union import flax import jax.numpy as jnp from scipy import integrate from ..configuration_utils import ConfigMixin, register_to_config +from ..utils import logging from .scheduling_utils_flax import ( CommonSchedulerState, FlaxKarrasDiffusionSchedulers, @@ -29,6 +29,9 @@ ) +logger = logging.get_logger(__name__) + + @flax.struct.dataclass class LMSDiscreteSchedulerState: common: CommonSchedulerState @@ -37,16 +40,25 @@ class LMSDiscreteSchedulerState: init_noise_sigma: jnp.ndarray timesteps: jnp.ndarray sigmas: jnp.ndarray - num_inference_steps: Optional[int] = None + num_inference_steps: int = None # running values - derivatives: Optional[jnp.ndarray] = None + derivatives: jnp.ndarray | None = None @classmethod def create( - cls, common: CommonSchedulerState, init_noise_sigma: jnp.ndarray, timesteps: jnp.ndarray, sigmas: jnp.ndarray + cls, + common: CommonSchedulerState, + init_noise_sigma: jnp.ndarray, + timesteps: jnp.ndarray, + sigmas: jnp.ndarray, ): - return cls(common=common, init_noise_sigma=init_noise_sigma, timesteps=timesteps, sigmas=sigmas) + return cls( + common=common, + init_noise_sigma=init_noise_sigma, + timesteps=timesteps, + sigmas=sigmas, + ) @dataclass @@ -97,13 +109,17 @@ def __init__( beta_start: float = 0.0001, beta_end: float = 0.02, beta_schedule: str = "linear", - trained_betas: Optional[jnp.ndarray] = None, + trained_betas: jnp.ndarray | None = None, prediction_type: str = "epsilon", dtype: jnp.dtype = jnp.float32, ): + logger.warning( + "Flax classes are deprecated and will be removed in Diffusers v1.0.0. We " + "recommend migrating to PyTorch classes or pinning your version of Diffusers." + ) self.dtype = dtype - def create_state(self, common: Optional[CommonSchedulerState] = None) -> LMSDiscreteSchedulerState: + def create_state(self, common: CommonSchedulerState | None = None) -> LMSDiscreteSchedulerState: if common is None: common = CommonSchedulerState.create(self) @@ -165,7 +181,10 @@ def lms_derivative(tau): return integrated_coeff def set_timesteps( - self, state: LMSDiscreteSchedulerState, num_inference_steps: int, shape: Tuple = () + self, + state: LMSDiscreteSchedulerState, + num_inference_steps: int, + shape: tuple = (), ) -> LMSDiscreteSchedulerState: """ Sets the timesteps used for the diffusion chain. Supporting function to be run before inference. @@ -177,7 +196,12 @@ def set_timesteps( the number of diffusion steps used when generating samples with a pre-trained model. """ - timesteps = jnp.linspace(self.config.num_train_timesteps - 1, 0, num_inference_steps, dtype=self.dtype) + timesteps = jnp.linspace( + self.config.num_train_timesteps - 1, + 0, + num_inference_steps, + dtype=self.dtype, + ) low_idx = jnp.floor(timesteps).astype(jnp.int32) high_idx = jnp.ceil(timesteps).astype(jnp.int32) @@ -208,7 +232,7 @@ def step( sample: jnp.ndarray, order: int = 4, return_dict: bool = True, - ) -> Union[FlaxLMSSchedulerOutput, Tuple]: + ) -> FlaxLMSSchedulerOutput | tuple: """ Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion process from the learned model outputs (most often the predicted noise). diff --git a/src/diffusers/schedulers/scheduling_ltx_euler_ancestral_rf.py b/src/diffusers/schedulers/scheduling_ltx_euler_ancestral_rf.py new file mode 100644 index 000000000000..453c8515c301 --- /dev/null +++ b/src/diffusers/schedulers/scheduling_ltx_euler_ancestral_rf.py @@ -0,0 +1,385 @@ +# Copyright 2025 Lightricks and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +LTXEulerAncestralRFScheduler + +This scheduler implements a K-diffusion style Euler-Ancestral sampler specialized for flow / CONST parameterization, +closely mirroring ComfyUI's `sample_euler_ancestral_RF` implementation used for LTX-Video. + +Reference implementation (ComfyUI): + comfy.k_diffusion.sampling.sample_euler_ancestral_RF +""" + +from dataclasses import dataclass + +import torch + +from ..configuration_utils import ConfigMixin, register_to_config +from ..utils import BaseOutput, logging +from ..utils.torch_utils import randn_tensor +from .scheduling_utils import SchedulerMixin + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class LTXEulerAncestralRFSchedulerOutput(BaseOutput): + """ + Output class for the scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor`): + Updated sample for the next step in the denoising process. + """ + + prev_sample: torch.FloatTensor + + +class LTXEulerAncestralRFScheduler(SchedulerMixin, ConfigMixin): + """ + Euler-Ancestral scheduler for LTX-Video (RF / CONST parametrization). + + This scheduler is intended for models where the network is trained with a CONST-like parameterization (as in LTXV / + FLUX). It approximates ComfyUI's `sample_euler_ancestral_RF` sampler and is useful when reproducing ComfyUI + workflows inside diffusers. + + The scheduler can either: + - reuse the [`FlowMatchEulerDiscreteScheduler`] sigma / timestep logic when only `num_inference_steps` is provided + (default diffusers-style usage), or + - follow an explicit ComfyUI-style sigma schedule when `sigmas` (or `timesteps`) are passed to [`set_timesteps`]. + + Args: + num_train_timesteps (`int`, defaults to 1000): + Included for config compatibility; not used to build the schedule. + eta (`float`, defaults to 1.0): + Stochasticity parameter. `eta=0.0` yields deterministic DDIM-like sampling; `eta=1.0` matches ComfyUI's + default RF behavior. + s_noise (`float`, defaults to 1.0): + Global scaling factor for the stochastic noise term. + """ + + # Allow config migration from the flow-match scheduler and back. + _compatibles = ["FlowMatchEulerDiscreteScheduler"] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + eta: float = 1.0, + s_noise: float = 1.0, + ): + # Note: num_train_timesteps is kept only for config compatibility. + self.num_inference_steps: int = None + self.sigmas: torch.Tensor | None = None + self.timesteps: torch.Tensor | None = None + self._step_index: int = None + self._begin_index: int = None + + @property + def step_index(self) -> int: + return self._step_index + + @property + def begin_index(self) -> int: + """ + The index for the first timestep. It can be set from a pipeline with `set_begin_index` to support + image-to-image like workflows that start denoising part-way through the schedule. + """ + return self._begin_index + + def set_begin_index(self, begin_index: int = 0): + """ + Included for API compatibility; not strictly needed here but kept to allow pipelines that call + `set_begin_index`. + """ + self._begin_index = begin_index + + def index_for_timestep( + self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None + ) -> int: + """ + Map a (continuous) `timestep` value to an index into `self.timesteps`. + + This follows the convention used in other discrete schedulers: if the same timestep value appears multiple + times in the schedule (which can happen when starting in the middle of the schedule), the *second* occurrence + is used for the first `step` call so that no sigma is accidentally skipped. + """ + if schedule_timesteps is None: + if self.timesteps is None: + raise ValueError("Timesteps have not been set. Call `set_timesteps` first.") + schedule_timesteps = self.timesteps + + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(schedule_timesteps.device) + + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + pos = 1 if len(indices) > 1 else 0 + + if len(indices) == 0: + raise ValueError( + "Passed `timestep` is not in `self.timesteps`. Make sure to use values from `scheduler.timesteps`." + ) + + return indices[pos].item() + + def _init_step_index(self, timestep: float | torch.Tensor): + """ + Initialize the internal step index based on a given timestep. + """ + if self.timesteps is None: + raise ValueError("Timesteps have not been set. Call `set_timesteps` first.") + + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def set_timesteps( + self, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + sigmas: list[float] | torch.Tensor | None = None, + timesteps: list[float] | torch.Tensor | None = None, + mu: float | None = None, + **kwargs, + ): + """ + Set the sigma / timestep schedule for sampling. + + When `sigmas` or `timesteps` are provided explicitly, they are used as the RF sigma schedule (ComfyUI-style) + and are expected to include the terminal 0.0. When both are `None`, the scheduler reuses the + [`FlowMatchEulerDiscreteScheduler`] logic to generate sigmas from `num_inference_steps` and the stored config + (including any resolution-dependent shifting, Karras/beta schedules, etc.). + + Args: + num_inference_steps (`int`, *optional*): + Number of denoising steps. If provided together with explicit `sigmas`/`timesteps`, they are expected + to be consistent and are otherwise ignored with a warning. + device (`str` or `torch.device`, *optional*): + Device to move the internal tensors to. + sigmas (`list[float]` or `torch.Tensor`, *optional*): + Explicit sigma schedule, e.g. `[1.0, 0.99, ..., 0.0]`. + timesteps (`list[float]` or `torch.Tensor`, *optional*): + Optional alias for `sigmas`. If `sigmas` is None and `timesteps` is provided, timesteps are treated as + sigmas. + mu (`float`, *optional*): + Optional shift parameter used when delegating to [`FlowMatchEulerDiscreteScheduler.set_timesteps`] and + `config.use_dynamic_shifting` is `True`. + """ + # 1. Auto-generate schedule (FlowMatch-style) when no explicit sigmas/timesteps are given + if sigmas is None and timesteps is None: + if num_inference_steps is None: + raise ValueError( + "LTXEulerAncestralRFScheduler.set_timesteps requires either explicit `sigmas`/`timesteps` " + "or a `num_inference_steps` value." + ) + + # We reuse FlowMatchEulerDiscreteScheduler to construct a sigma schedule that is + # consistent with the original LTX training setup (including optional time shifting, + # Karras / exponential / beta schedules, etc.). + from .scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler + + base_scheduler = FlowMatchEulerDiscreteScheduler.from_config(self.config) + base_scheduler.set_timesteps( + num_inference_steps=num_inference_steps, + device=device, + sigmas=None, + mu=mu, + timesteps=None, + ) + + self.num_inference_steps = base_scheduler.num_inference_steps + # Keep sigmas / timesteps on the requested device so step() can operate on-device without + # extra transfers. + self.sigmas = base_scheduler.sigmas.to(device=device) + self.timesteps = base_scheduler.timesteps.to(device=device) + self._step_index = None + self._begin_index = None + return + + # 2. Explicit sigma schedule (ComfyUI-style path) + if sigmas is None: + # `timesteps` is treated as sigmas in RF / flow-matching setups. + sigmas = timesteps + + if isinstance(sigmas, list): + sigmas_tensor = torch.tensor(sigmas, dtype=torch.float32) + elif isinstance(sigmas, torch.Tensor): + sigmas_tensor = sigmas.to(dtype=torch.float32) + else: + raise TypeError(f"`sigmas` must be a list or torch.Tensor, got {type(sigmas)}.") + + if sigmas_tensor.ndim != 1: + raise ValueError(f"`sigmas` must be a 1D tensor, got shape {tuple(sigmas_tensor.shape)}.") + + if sigmas_tensor[-1].abs().item() > 1e-6: + logger.warning( + "The last sigma in the schedule is not zero (%.6f). " + "For best compatibility with ComfyUI's RF sampler, the terminal sigma " + "should be 0.0.", + sigmas_tensor[-1].item(), + ) + + # Move to device once, then derive timesteps. + if device is not None: + sigmas_tensor = sigmas_tensor.to(device) + + # Internal sigma schedule stays in [0, 1] (as provided). + self.sigmas = sigmas_tensor + # Timesteps are scaled to match the training setup of LTX (FlowMatch-style), + # where the network expects timesteps on [0, num_train_timesteps]. + # This keeps the transformer conditioning in the expected range while the RF + # scheduler still operates on the raw sigma values. + num_train = float(getattr(self.config, "num_train_timesteps", 1000)) + self.timesteps = sigmas_tensor * num_train + + if num_inference_steps is not None and num_inference_steps != len(sigmas) - 1: + logger.warning( + "Provided `num_inference_steps=%d` does not match `len(sigmas)-1=%d`. " + "Overriding `num_inference_steps` with `len(sigmas)-1`.", + num_inference_steps, + len(sigmas) - 1, + ) + + self.num_inference_steps = len(sigmas) - 1 + self._step_index = None + self._begin_index = None + + def _sigma_broadcast(self, sigma: torch.Tensor, sample: torch.Tensor) -> torch.Tensor: + """ + Helper to broadcast a scalar sigma to the shape of `sample`. + """ + while sigma.ndim < sample.ndim: + sigma = sigma.view(*sigma.shape, 1) + return sigma + + def step( + self, + model_output: torch.FloatTensor, + timestep: float | torch.Tensor, + sample: torch.FloatTensor, + generator: torch.Generator | None = None, + return_dict: bool = True, + ) -> LTXEulerAncestralRFSchedulerOutput | tuple[torch.FloatTensor]: + """ + Perform a single Euler-Ancestral RF update step. + + Args: + model_output (`torch.FloatTensor`): + Raw model output at the current step. Interpreted under the CONST parametrization as `v_t`, with + denoised state reconstructed as `x0 = x_t - sigma_t * v_t`. + timestep (`float` or `torch.Tensor`): + The current sigma value (must match one entry in `self.timesteps`). + sample (`torch.FloatTensor`): + Current latent sample `x_t`. + generator (`torch.Generator`, *optional*): + Optional generator for reproducible noise. + return_dict (`bool`): + If `True`, return a `LTXEulerAncestralRFSchedulerOutput`; otherwise return a tuple where the first + element is the updated sample. + """ + + if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)): + raise ValueError( + ( + "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `LTXEulerAncestralRFScheduler.step()` is not supported. Make sure to pass" + " one of the `scheduler.timesteps` values as `timestep`." + ), + ) + + if self.sigmas is None or self.timesteps is None: + raise ValueError("Scheduler has not been initialized. Call `set_timesteps` before `step`.") + + if self._step_index is None: + self._init_step_index(timestep) + + i = self._step_index + if i >= len(self.sigmas) - 1: + # Already at the end; simply return the current sample. + prev_sample = sample + else: + # Work in float32 for numerical stability + sample_f = sample.to(torch.float32) + model_output_f = model_output.to(torch.float32) + + sigma = self.sigmas[i] + sigma_next = self.sigmas[i + 1] + + sigma_b = self._sigma_broadcast(sigma.view(1), sample_f) + sigma_next_b = self._sigma_broadcast(sigma_next.view(1), sample_f) + + # Approximate denoised x0 under CONST parametrization: + # x0 = x_t - sigma_t * v_t + denoised = sample_f - sigma_b * model_output_f + + if sigma_next.abs().item() < 1e-8: + # Final denoising step + x = denoised + else: + eta = float(self.config.eta) + s_noise = float(self.config.s_noise) + + # Downstep computation (ComfyUI RF variant) + downstep_ratio = 1.0 + (sigma_next / sigma - 1.0) * eta + sigma_down = sigma_next * downstep_ratio + + alpha_ip1 = 1.0 - sigma_next + alpha_down = 1.0 - sigma_down + + # Deterministic part (Euler step in (x, x0)-space) + sigma_down_b = self._sigma_broadcast(sigma_down.view(1), sample_f) + alpha_ip1_b = self._sigma_broadcast(alpha_ip1.view(1), sample_f) + alpha_down_b = self._sigma_broadcast(alpha_down.view(1), sample_f) + + sigma_ratio = sigma_down_b / sigma_b + x = sigma_ratio * sample_f + (1.0 - sigma_ratio) * denoised + + # Stochastic ancestral noise + if eta > 0.0 and s_noise > 0.0: + renoise_coeff = ( + (sigma_next_b**2 - sigma_down_b**2 * alpha_ip1_b**2 / (alpha_down_b**2 + 1e-12)) + .clamp(min=0.0) + .sqrt() + ) + + noise = randn_tensor( + sample_f.shape, generator=generator, device=sample_f.device, dtype=sample_f.dtype + ) + x = (alpha_ip1_b / (alpha_down_b + 1e-12)) * x + noise * renoise_coeff * s_noise + + prev_sample = x.to(sample.dtype) + + # Advance internal step index + self._step_index = min(self._step_index + 1, len(self.sigmas) - 1) + + if not return_dict: + return (prev_sample,) + + return LTXEulerAncestralRFSchedulerOutput(prev_sample=prev_sample) + + def __len__(self) -> int: + # For compatibility with other schedulers; used e.g. in some training + # utilities to infer the maximum number of training timesteps. + return int(getattr(self.config, "num_train_timesteps", 1000)) diff --git a/src/diffusers/schedulers/scheduling_pndm.py b/src/diffusers/schedulers/scheduling_pndm.py index 651532b06ddb..f87dff14375d 100644 --- a/src/diffusers/schedulers/scheduling_pndm.py +++ b/src/diffusers/schedulers/scheduling_pndm.py @@ -15,7 +15,7 @@ # DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim import math -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -28,7 +28,7 @@ def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -42,8 +42,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -54,6 +54,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -115,8 +122,8 @@ def __init__( num_train_timesteps: int = 1000, beta_start: float = 0.0001, beta_end: float = 0.02, - beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, + beta_schedule: str = "linear", + trained_betas: np.ndarray | list[float] | None = None, skip_prk_steps: bool = False, set_alpha_to_one: bool = False, prediction_type: Literal["epsilon", "v_prediction"] = "epsilon", @@ -162,7 +169,7 @@ def __init__( self.plms_timesteps = None self.timesteps = None - def set_timesteps(self, num_inference_steps: int, device: Optional[Union[str, torch.device]] = None) -> None: + def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -228,7 +235,7 @@ def step( timestep: int, sample: torch.Tensor, return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: + ) -> SchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise), and calls [`~PNDMScheduler.step_prk`] @@ -261,7 +268,7 @@ def step_prk( timestep: int, sample: torch.Tensor, return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: + ) -> SchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the Runge-Kutta method. It performs four forward passes to approximate the solution to the differential @@ -320,7 +327,7 @@ def step_plms( timestep: int, sample: torch.Tensor, return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: + ) -> SchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the linear multistep method. It performs one forward pass multiple times to approximate the solution. diff --git a/src/diffusers/schedulers/scheduling_pndm_flax.py b/src/diffusers/schedulers/scheduling_pndm_flax.py index 44bafccd5520..377e2d35ba0a 100644 --- a/src/diffusers/schedulers/scheduling_pndm_flax.py +++ b/src/diffusers/schedulers/scheduling_pndm_flax.py @@ -15,13 +15,13 @@ # DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim from dataclasses import dataclass -from typing import Optional, Tuple, Union import flax import jax import jax.numpy as jnp from ..configuration_utils import ConfigMixin, register_to_config +from ..utils import logging from .scheduling_utils_flax import ( CommonSchedulerState, FlaxKarrasDiffusionSchedulers, @@ -31,6 +31,9 @@ ) +logger = logging.get_logger(__name__) + + @flax.struct.dataclass class PNDMSchedulerState: common: CommonSchedulerState @@ -39,15 +42,15 @@ class PNDMSchedulerState: # setable values init_noise_sigma: jnp.ndarray timesteps: jnp.ndarray - num_inference_steps: Optional[int] = None - prk_timesteps: Optional[jnp.ndarray] = None - plms_timesteps: Optional[jnp.ndarray] = None + num_inference_steps: int = None + prk_timesteps: jnp.ndarray | None = None + plms_timesteps: jnp.ndarray | None = None # running values - cur_model_output: Optional[jnp.ndarray] = None - counter: Optional[jnp.int32] = None - cur_sample: Optional[jnp.ndarray] = None - ets: Optional[jnp.ndarray] = None + cur_model_output: jnp.ndarray | None = None + counter: jnp.int32 | None = None + cur_sample: jnp.ndarray | None = None + ets: jnp.ndarray | None = None @classmethod def create( @@ -124,13 +127,17 @@ def __init__( beta_start: float = 0.0001, beta_end: float = 0.02, beta_schedule: str = "linear", - trained_betas: Optional[jnp.ndarray] = None, + trained_betas: jnp.ndarray | None = None, skip_prk_steps: bool = False, set_alpha_to_one: bool = False, steps_offset: int = 0, prediction_type: str = "epsilon", dtype: jnp.dtype = jnp.float32, ): + logger.warning( + "Flax classes are deprecated and will be removed in Diffusers v1.0.0. We " + "recommend migrating to PyTorch classes or pinning your version of Diffusers." + ) self.dtype = dtype # For now we only support F-PNDM, i.e. the runge-kutta method @@ -138,7 +145,7 @@ def __init__( # mainly at formula (9), (12), (13) and the Algorithm 2. self.pndm_order = 4 - def create_state(self, common: Optional[CommonSchedulerState] = None) -> PNDMSchedulerState: + def create_state(self, common: CommonSchedulerState | None = None) -> PNDMSchedulerState: if common is None: common = CommonSchedulerState.create(self) @@ -162,7 +169,7 @@ def create_state(self, common: Optional[CommonSchedulerState] = None) -> PNDMSch timesteps=timesteps, ) - def set_timesteps(self, state: PNDMSchedulerState, num_inference_steps: int, shape: Tuple) -> PNDMSchedulerState: + def set_timesteps(self, state: PNDMSchedulerState, num_inference_steps: int, shape: tuple) -> PNDMSchedulerState: """ Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference. @@ -171,7 +178,7 @@ def set_timesteps(self, state: PNDMSchedulerState, num_inference_steps: int, sha the `FlaxPNDMScheduler` state data class instance. num_inference_steps (`int`): the number of diffusion steps used when generating samples with a pre-trained model. - shape (`Tuple`): + shape (`tuple`): the shape of the samples to be generated. """ @@ -190,7 +197,10 @@ def set_timesteps(self, state: PNDMSchedulerState, num_inference_steps: int, sha else: prk_timesteps = _timesteps[-self.pndm_order :].repeat(2) + jnp.tile( - jnp.array([0, self.config.num_train_timesteps // num_inference_steps // 2], dtype=jnp.int32), + jnp.array( + [0, self.config.num_train_timesteps // num_inference_steps // 2], + dtype=jnp.int32, + ), self.pndm_order, ) @@ -218,7 +228,10 @@ def set_timesteps(self, state: PNDMSchedulerState, num_inference_steps: int, sha ) def scale_model_input( - self, state: PNDMSchedulerState, sample: jnp.ndarray, timestep: Optional[int] = None + self, + state: PNDMSchedulerState, + sample: jnp.ndarray, + timestep: int | None = None, ) -> jnp.ndarray: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the @@ -241,7 +254,7 @@ def step( timestep: int, sample: jnp.ndarray, return_dict: bool = True, - ) -> Union[FlaxPNDMSchedulerOutput, Tuple]: + ) -> FlaxPNDMSchedulerOutput | tuple: """ Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion process from the learned model outputs (most often the predicted noise). @@ -295,7 +308,7 @@ def step_prk( model_output: jnp.ndarray, timestep: int, sample: jnp.ndarray, - ) -> Union[FlaxPNDMSchedulerOutput, Tuple]: + ) -> FlaxPNDMSchedulerOutput | tuple: """ Step function propagating the sample with the Runge-Kutta method. RK takes 4 forward passes to approximate the solution to the differential equation. @@ -320,7 +333,9 @@ def step_prk( ) diff_to_prev = jnp.where( - state.counter % 2, 0, self.config.num_train_timesteps // state.num_inference_steps // 2 + state.counter % 2, + 0, + self.config.num_train_timesteps // state.num_inference_steps // 2, ) prev_timestep = timestep - diff_to_prev timestep = state.prk_timesteps[state.counter // 4 * 4] @@ -363,7 +378,7 @@ def step_plms( model_output: jnp.ndarray, timestep: int, sample: jnp.ndarray, - ) -> Union[FlaxPNDMSchedulerOutput, Tuple]: + ) -> FlaxPNDMSchedulerOutput | tuple: """ Step function propagating the sample with the linear multi-step method. This has one forward pass with multiple times to approximate the solution. @@ -401,7 +416,9 @@ def step_plms( prev_timestep = jnp.where(state.counter == 1, timestep, prev_timestep) timestep = jnp.where( - state.counter == 1, timestep + self.config.num_train_timesteps // state.num_inference_steps, timestep + state.counter == 1, + timestep + self.config.num_train_timesteps // state.num_inference_steps, + timestep, ) # Reference: @@ -466,7 +483,9 @@ def _get_prev_sample(self, state: PNDMSchedulerState, sample, timestep, prev_tim # prev_sample -> x_(t−δ) alpha_prod_t = state.common.alphas_cumprod[timestep] alpha_prod_t_prev = jnp.where( - prev_timestep >= 0, state.common.alphas_cumprod[prev_timestep], state.final_alpha_cumprod + prev_timestep >= 0, + state.common.alphas_cumprod[prev_timestep], + state.final_alpha_cumprod, ) beta_prod_t = 1 - alpha_prod_t beta_prod_t_prev = 1 - alpha_prod_t_prev diff --git a/src/diffusers/schedulers/scheduling_repaint.py b/src/diffusers/schedulers/scheduling_repaint.py index a2eaf8eb3abd..43665f6c68e8 100644 --- a/src/diffusers/schedulers/scheduling_repaint.py +++ b/src/diffusers/schedulers/scheduling_repaint.py @@ -14,7 +14,7 @@ import math from dataclasses import dataclass -from typing import Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -47,7 +47,7 @@ class RePaintSchedulerOutput(BaseOutput): def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -61,8 +61,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -73,6 +73,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -126,7 +133,7 @@ def __init__( beta_end: float = 0.02, beta_schedule: str = "linear", eta: float = 0.0, - trained_betas: Optional[np.ndarray] = None, + trained_betas: np.ndarray | None = None, clip_sample: bool = True, ): if trained_betas is not None: @@ -161,7 +168,7 @@ def __init__( self.eta = eta - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. @@ -183,7 +190,7 @@ def set_timesteps( num_inference_steps: int, jump_length: int = 10, jump_n_sample: int = 10, - device: Union[str, torch.device] = None, + device: str | torch.device = None, ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -251,9 +258,9 @@ def step( sample: torch.Tensor, original_image: torch.Tensor, mask: torch.Tensor, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, return_dict: bool = True, - ) -> Union[RePaintSchedulerOutput, Tuple]: + ) -> RePaintSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). diff --git a/src/diffusers/schedulers/scheduling_sasolver.py b/src/diffusers/schedulers/scheduling_sasolver.py index 5783e20de69d..28369e9e3384 100644 --- a/src/diffusers/schedulers/scheduling_sasolver.py +++ b/src/diffusers/schedulers/scheduling_sasolver.py @@ -16,7 +16,7 @@ # The codebase is modified based on https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py import math -from typing import Callable, List, Literal, Optional, Tuple, Union +from typing import Callable, Literal import numpy as np import torch @@ -35,7 +35,7 @@ def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -49,8 +49,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -61,6 +61,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -155,23 +162,23 @@ def __init__( beta_start: float = 0.0001, beta_end: float = 0.02, beta_schedule: str = "linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, + trained_betas: np.ndarray | list[float] | None = None, predictor_order: int = 2, corrector_order: int = 2, prediction_type: str = "epsilon", - tau_func: Optional[Callable] = None, + tau_func: Callable | None = None, thresholding: bool = False, dynamic_thresholding_ratio: float = 0.995, sample_max_value: float = 1.0, algorithm_type: str = "data_prediction", lower_order_final: bool = True, - use_karras_sigmas: Optional[bool] = False, - use_exponential_sigmas: Optional[bool] = False, - use_beta_sigmas: Optional[bool] = False, - use_flow_sigmas: Optional[bool] = False, - flow_shift: Optional[float] = 1.0, + use_karras_sigmas: bool = False, + use_exponential_sigmas: bool = False, + use_beta_sigmas: bool = False, + use_flow_sigmas: bool = False, + flow_shift: float = 1.0, lambda_min_clipped: float = -float("inf"), - variance_type: Optional[str] = None, + variance_type: str | None = None, timestep_spacing: str = "linspace", steps_offset: int = 0, ): @@ -259,7 +266,7 @@ def set_begin_index(self, begin_index: int = 0): """ self._begin_index = begin_index - def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torch.device] = None): + def set_timesteps(self, num_inference_steps: int = None, device: str | torch.device = None): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -431,7 +438,7 @@ def _sigma_to_alpha_sigma_t(self, sigma): The sigma value(s) to convert. Returns: - `Tuple[torch.Tensor, torch.Tensor]`: + `tuple[torch.Tensor, torch.Tensor]`: A tuple containing (alpha_t, sigma_t) values. """ if self.config.use_flow_sigmas: @@ -1115,7 +1122,9 @@ def stochastic_adams_moulton_update( # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep def index_for_timestep( - self, timestep: Union[int, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None + self, + timestep: int | torch.Tensor, + schedule_timesteps: torch.Tensor | None = None, ) -> int: """ Find the index for a given timestep in the schedule. @@ -1172,7 +1181,7 @@ def step( sample: torch.Tensor, generator=None, return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: + ) -> SchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the SA-Solver. diff --git a/src/diffusers/schedulers/scheduling_scm.py b/src/diffusers/schedulers/scheduling_scm.py index 7b01d886299c..fb3bbc40d726 100644 --- a/src/diffusers/schedulers/scheduling_scm.py +++ b/src/diffusers/schedulers/scheduling_scm.py @@ -16,7 +16,6 @@ # and https://github.com/hojonathanho/diffusion from dataclasses import dataclass -from typing import Optional, Tuple, Union import numpy as np import torch @@ -46,7 +45,7 @@ class SCMSchedulerOutput(BaseOutput): """ prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None + pred_original_sample: torch.Tensor | None = None class SCMScheduler(SchedulerMixin, ConfigMixin): @@ -118,7 +117,7 @@ def set_timesteps( self, num_inference_steps: int, timesteps: torch.Tensor = None, - device: Union[str, torch.device] = None, + device: str | torch.device = None, max_timesteps: float = 1.57080, intermediate_timesteps: float = 1.3, ): @@ -173,7 +172,7 @@ def set_timesteps( self._begin_index = None # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: + def _init_step_index(self, timestep: float | torch.Tensor) -> None: """ Initialize the step index for the scheduler based on the given timestep. @@ -190,7 +189,7 @@ def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep def index_for_timestep( - self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None + self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None ) -> int: """ Find the index of a given timestep in the timestep schedule. @@ -226,7 +225,7 @@ def step( sample: torch.FloatTensor, generator: torch.Generator = None, return_dict: bool = True, - ) -> Union[SCMSchedulerOutput, Tuple]: + ) -> SCMSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). diff --git a/src/diffusers/schedulers/scheduling_sde_ve.py b/src/diffusers/schedulers/scheduling_sde_ve.py index 1bfc08cce5e9..e4e37666c60d 100644 --- a/src/diffusers/schedulers/scheduling_sde_ve.py +++ b/src/diffusers/schedulers/scheduling_sde_ve.py @@ -16,7 +16,6 @@ import math from dataclasses import dataclass -from typing import Optional, Tuple, Union import torch @@ -86,7 +85,7 @@ def __init__( self.set_sigmas(num_train_timesteps, sigma_min, sigma_max, sampling_eps) - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: int = None) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. @@ -103,9 +102,7 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None """ return sample - def set_timesteps( - self, num_inference_steps: int, sampling_eps: float = None, device: Union[str, torch.device] = None - ): + def set_timesteps(self, num_inference_steps: int, sampling_eps: float = None, device: str | torch.device = None): """ Sets the continuous timesteps used for the diffusion chain (to be run before inference). @@ -162,9 +159,9 @@ def step_pred( model_output: torch.Tensor, timestep: int, sample: torch.Tensor, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, return_dict: bool = True, - ) -> Union[SdeVeOutput, Tuple]: + ) -> SdeVeOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). @@ -229,9 +226,9 @@ def step_correct( self, model_output: torch.Tensor, sample: torch.Tensor, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: + ) -> SchedulerOutput | tuple: """ Correct the predicted sample based on the `model_output` of the network. This is often run repeatedly after making the prediction for the previous timestep. diff --git a/src/diffusers/schedulers/scheduling_sde_ve_flax.py b/src/diffusers/schedulers/scheduling_sde_ve_flax.py index 09cd081462b3..62a54c7dc948 100644 --- a/src/diffusers/schedulers/scheduling_sde_ve_flax.py +++ b/src/diffusers/schedulers/scheduling_sde_ve_flax.py @@ -15,7 +15,6 @@ # DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch from dataclasses import dataclass -from typing import Optional, Tuple, Union import flax import jax @@ -23,15 +22,23 @@ from jax import random from ..configuration_utils import ConfigMixin, register_to_config -from .scheduling_utils_flax import FlaxSchedulerMixin, FlaxSchedulerOutput, broadcast_to_shape_from_left +from ..utils import logging +from .scheduling_utils_flax import ( + FlaxSchedulerMixin, + FlaxSchedulerOutput, + broadcast_to_shape_from_left, +) + + +logger = logging.get_logger(__name__) @flax.struct.dataclass class ScoreSdeVeSchedulerState: # setable values - timesteps: Optional[jnp.ndarray] = None - discrete_sigmas: Optional[jnp.ndarray] = None - sigmas: Optional[jnp.ndarray] = None + timesteps: jnp.ndarray | None = None + discrete_sigmas: jnp.ndarray | None = None + sigmas: jnp.ndarray | None = None @classmethod def create(cls): @@ -54,7 +61,7 @@ class FlaxSdeVeOutput(FlaxSchedulerOutput): state: ScoreSdeVeSchedulerState prev_sample: jnp.ndarray - prev_sample_mean: Optional[jnp.ndarray] = None + prev_sample_mean: jnp.ndarray | None = None class FlaxScoreSdeVeScheduler(FlaxSchedulerMixin, ConfigMixin): @@ -95,7 +102,10 @@ def __init__( sampling_eps: float = 1e-5, correct_steps: int = 1, ): - pass + logger.warning( + "Flax classes are deprecated and will be removed in Diffusers v1.0.0. We " + "recommend migrating to PyTorch classes or pinning your version of Diffusers." + ) def create_state(self): state = ScoreSdeVeSchedulerState.create() @@ -108,7 +118,11 @@ def create_state(self): ) def set_timesteps( - self, state: ScoreSdeVeSchedulerState, num_inference_steps: int, shape: Tuple = (), sampling_eps: float = None + self, + state: ScoreSdeVeSchedulerState, + num_inference_steps: int, + shape: tuple = (), + sampling_eps: float = None, ) -> ScoreSdeVeSchedulerState: """ Sets the continuous timesteps used for the diffusion chain. Supporting function to be run before inference. @@ -172,7 +186,7 @@ def step_pred( sample: jnp.ndarray, key: jax.Array, return_dict: bool = True, - ) -> Union[FlaxSdeVeOutput, Tuple]: + ) -> FlaxSdeVeOutput | tuple: """ Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion process from the learned model outputs (most often the predicted noise). @@ -231,7 +245,7 @@ def step_correct( sample: jnp.ndarray, key: jax.Array, return_dict: bool = True, - ) -> Union[FlaxSdeVeOutput, Tuple]: + ) -> FlaxSdeVeOutput | tuple: """ Correct the predicted sample based on the output model_output of the network. This is often run repeatedly after making the prediction for the previous timestep. diff --git a/src/diffusers/schedulers/scheduling_tcd.py b/src/diffusers/schedulers/scheduling_tcd.py index 7b4840ffdb19..838d040c04a6 100644 --- a/src/diffusers/schedulers/scheduling_tcd.py +++ b/src/diffusers/schedulers/scheduling_tcd.py @@ -17,7 +17,7 @@ import math from dataclasses import dataclass -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -45,14 +45,14 @@ class TCDSchedulerOutput(BaseOutput): """ prev_sample: torch.Tensor - pred_noised_sample: Optional[torch.Tensor] = None + pred_noised_sample: torch.Tensor | None = None # Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -66,8 +66,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -78,6 +78,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -201,7 +208,7 @@ def __init__( beta_start: float = 0.00085, beta_end: float = 0.012, beta_schedule: str = "scaled_linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, + trained_betas: np.ndarray | list[float] | None = None, original_inference_steps: int = 50, clip_sample: bool = False, clip_sample_range: float = 1.0, @@ -254,7 +261,7 @@ def __init__( # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep def index_for_timestep( - self, timestep: Union[float, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None + self, timestep: float | torch.Tensor, schedule_timesteps: torch.Tensor | None = None ) -> int: """ Find the index of a given timestep in the timestep schedule. @@ -284,7 +291,7 @@ def index_for_timestep( return indices[pos].item() # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index - def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: + def _init_step_index(self, timestep: float | torch.Tensor) -> None: """ Initialize the step index for the scheduler based on the given timestep. @@ -321,7 +328,7 @@ def set_begin_index(self, begin_index: int = 0): """ self._begin_index = begin_index - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. @@ -413,10 +420,10 @@ def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: def set_timesteps( self, - num_inference_steps: Optional[int] = None, - device: Union[str, torch.device] = None, - original_inference_steps: Optional[int] = None, - timesteps: Optional[List[int]] = None, + num_inference_steps: int | None = None, + device: str | torch.device = None, + original_inference_steps: int | None = None, + timesteps: list[int] | None = None, strength: float = 1.0, ): """ @@ -433,7 +440,7 @@ def set_timesteps( schedule (which is different from the standard `diffusers` implementation). We will then take `num_inference_steps` timesteps from this schedule, evenly spaced in terms of indices, and use that as our final timestep schedule. If not set, this will default to the `original_inference_steps` attribute. - timesteps (`List[int]`, *optional*): + timesteps (`list[int]`, *optional*): Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default timestep spacing strategy of equal spacing between timesteps on the training/distillation timestep schedule is used. If `timesteps` is passed, `num_inference_steps` must be `None`. @@ -579,9 +586,9 @@ def step( timestep: int, sample: torch.Tensor, eta: float = 0.3, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, return_dict: bool = True, - ) -> Union[TCDSchedulerOutput, Tuple]: + ) -> TCDSchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). @@ -770,7 +777,7 @@ def previous_timestep(self, timestep): The current timestep. Returns: - `int`: + `int` or `torch.Tensor`: The previous timestep. """ if self.custom_timesteps or self.num_inference_steps: diff --git a/src/diffusers/schedulers/scheduling_unclip.py b/src/diffusers/schedulers/scheduling_unclip.py index 5a978dec649b..26226b4eb35b 100644 --- a/src/diffusers/schedulers/scheduling_unclip.py +++ b/src/diffusers/schedulers/scheduling_unclip.py @@ -14,7 +14,7 @@ import math from dataclasses import dataclass -from typing import Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -41,14 +41,14 @@ class UnCLIPSchedulerOutput(BaseOutput): """ prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None + pred_original_sample: torch.Tensor | None = None # Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -62,8 +62,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -74,6 +74,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -125,7 +132,7 @@ def __init__( num_train_timesteps: int = 1000, variance_type: str = "fixed_small_log", clip_sample: bool = True, - clip_sample_range: Optional[float] = 1.0, + clip_sample_range: float = 1.0, prediction_type: str = "epsilon", beta_schedule: str = "squaredcos_cap_v2", ): @@ -147,7 +154,7 @@ def __init__( self.variance_type = variance_type - def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: + def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. @@ -161,7 +168,7 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None """ return sample - def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None): + def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None): """ Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference. @@ -219,10 +226,10 @@ def step( model_output: torch.Tensor, timestep: int, sample: torch.Tensor, - prev_timestep: Optional[int] = None, + prev_timestep: int | None = None, generator=None, return_dict: bool = True, - ) -> Union[UnCLIPSchedulerOutput, Tuple]: + ) -> UnCLIPSchedulerOutput | tuple: """ Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion process from the learned model outputs (most often the predicted noise). diff --git a/src/diffusers/schedulers/scheduling_unipc_multistep.py b/src/diffusers/schedulers/scheduling_unipc_multistep.py index 5ea56b300be2..71a5444491ed 100644 --- a/src/diffusers/schedulers/scheduling_unipc_multistep.py +++ b/src/diffusers/schedulers/scheduling_unipc_multistep.py @@ -16,7 +16,7 @@ # The codebase is modified based on https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py import math -from typing import List, Literal, Optional, Tuple, Union +from typing import Literal import numpy as np import torch @@ -34,7 +34,7 @@ def betas_for_alpha_bar( num_diffusion_timesteps: int, max_beta: float = 0.999, - alpha_transform_type: Literal["cosine", "exp"] = "cosine", + alpha_transform_type: Literal["cosine", "exp", "laplace"] = "cosine", ) -> torch.Tensor: """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of @@ -48,8 +48,8 @@ def betas_for_alpha_bar( The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. - alpha_transform_type (`"cosine"` or `"exp"`, defaults to `"cosine"`): - The type of noise schedule for `alpha_bar`. Choose from `cosine` or `exp`. + alpha_transform_type (`str`, defaults to `"cosine"`): + The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: @@ -60,6 +60,13 @@ def betas_for_alpha_bar( def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + elif alpha_transform_type == "laplace": + + def alpha_bar_fn(t): + lmb = -0.5 * math.copysign(1, 0.5 - t) * math.log(1 - 2 * math.fabs(0.5 - t) + 1e-6) + snr = math.exp(lmb) + return math.sqrt(snr / (1 + snr)) + elif alpha_transform_type == "exp": def alpha_bar_fn(t): @@ -194,8 +201,8 @@ def __init__( num_train_timesteps: int = 1000, beta_start: float = 0.0001, beta_end: float = 0.02, - beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear", - trained_betas: Optional[Union[np.ndarray, List[float]]] = None, + beta_schedule: str = "linear", + trained_betas: np.ndarray | list[float] | None = None, solver_order: int = 2, prediction_type: Literal["epsilon", "sample", "v_prediction", "flow_prediction"] = "epsilon", thresholding: bool = False, @@ -204,21 +211,22 @@ def __init__( predict_x0: bool = True, solver_type: Literal["bh1", "bh2"] = "bh2", lower_order_final: bool = True, - disable_corrector: List[int] = [], - solver_p: Optional[SchedulerMixin] = None, - use_karras_sigmas: Optional[bool] = False, - use_exponential_sigmas: Optional[bool] = False, - use_beta_sigmas: Optional[bool] = False, - use_flow_sigmas: Optional[bool] = False, - flow_shift: Optional[float] = 1.0, + disable_corrector: list[int] = [], + solver_p: SchedulerMixin = None, + use_karras_sigmas: bool = False, + use_exponential_sigmas: bool = False, + use_beta_sigmas: bool = False, + use_flow_sigmas: bool = False, + flow_shift: float = 1.0, timestep_spacing: Literal["linspace", "leading", "trailing"] = "linspace", steps_offset: int = 0, - final_sigmas_type: Optional[Literal["zero", "sigma_min"]] = "zero", + final_sigmas_type: Literal["zero", "sigma_min"] = "zero", rescale_betas_zero_snr: bool = False, use_dynamic_shifting: bool = False, time_shift_type: Literal["exponential"] = "exponential", - sigma_min: Optional[float] = None, - sigma_max: Optional[float] = None, + sigma_min: float | None = None, + sigma_max: bool | None = None, + shift_terminal: bool | None = None, ) -> None: if self.config.use_beta_sigmas and not is_scipy_available(): raise ImportError("Make sure to install scipy if you want to use beta sigmas.") @@ -238,6 +246,8 @@ def __init__( self.betas = betas_for_alpha_bar(num_train_timesteps) else: raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}") + if shift_terminal is not None and not use_flow_sigmas: + raise ValueError("`shift_terminal` is only supported when `use_flow_sigmas=True`.") if rescale_betas_zero_snr: self.betas = rescale_zero_terminal_snr(self.betas) @@ -281,14 +291,14 @@ def __init__( self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication @property - def step_index(self) -> Optional[int]: + def step_index(self) -> int: """ The index counter for current timestep. It will increase 1 after each scheduler step. """ return self._step_index @property - def begin_index(self) -> Optional[int]: + def begin_index(self) -> int: """ The index for the first timestep. It should be set from pipeline with `set_begin_index` method. """ @@ -306,8 +316,12 @@ def set_begin_index(self, begin_index: int = 0) -> None: self._begin_index = begin_index def set_timesteps( - self, num_inference_steps: int, device: Optional[Union[str, torch.device]] = None, mu: Optional[float] = None - ) -> None: + self, + num_inference_steps: int | None = None, + device: str | torch.device = None, + sigmas: list[float] | None = None, + mu: bool | None = None, + ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -316,13 +330,24 @@ def set_timesteps( The number of diffusion steps used when generating samples with a pre-trained model. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + sigmas (`List[float]`, *optional*): + Custom values for sigmas to be used for each diffusion step. If `None`, the sigmas are computed + automatically. mu (`float`, *optional*): Optional mu parameter for dynamic shifting when using exponential time shift type. """ + if self.config.use_dynamic_shifting and mu is None: + raise ValueError("`mu` must be passed when `use_dynamic_shifting` is set to be `True`") + + if sigmas is not None: + if not self.config.use_flow_sigmas: + raise ValueError( + "Passing `sigmas` is only supported when `use_flow_sigmas=True`. " + "Please set `use_flow_sigmas=True` during scheduler initialization." + ) + num_inference_steps = len(sigmas) + # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 - if mu is not None: - assert self.config.use_dynamic_shifting and self.config.time_shift_type == "exponential" - self.config.flow_shift = np.exp(mu) if self.config.timestep_spacing == "linspace": timesteps = ( np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps + 1) @@ -347,8 +372,9 @@ def set_timesteps( f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'." ) - sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) if self.config.use_karras_sigmas: + if sigmas is None: + sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) log_sigmas = np.log(sigmas) sigmas = np.flip(sigmas).copy() sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) @@ -368,6 +394,8 @@ def set_timesteps( ) sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) elif self.config.use_exponential_sigmas: + if sigmas is None: + sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) log_sigmas = np.log(sigmas) sigmas = np.flip(sigmas).copy() sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) @@ -382,6 +410,8 @@ def set_timesteps( ) sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) elif self.config.use_beta_sigmas: + if sigmas is None: + sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) log_sigmas = np.log(sigmas) sigmas = np.flip(sigmas).copy() sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) @@ -396,9 +426,18 @@ def set_timesteps( ) sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) elif self.config.use_flow_sigmas: - alphas = np.linspace(1, 1 / self.config.num_train_timesteps, num_inference_steps + 1) - sigmas = 1.0 - alphas - sigmas = np.flip(self.config.flow_shift * sigmas / (1 + (self.config.flow_shift - 1) * sigmas))[:-1].copy() + if sigmas is None: + sigmas = np.linspace(1, 1 / self.config.num_train_timesteps, num_inference_steps + 1)[:-1] + if self.config.use_dynamic_shifting: + sigmas = self.time_shift(mu, 1.0, sigmas) + else: + sigmas = self.config.flow_shift * sigmas / (1 + (self.config.flow_shift - 1) * sigmas) + if self.config.shift_terminal: + sigmas = self.stretch_shift_to_terminal(sigmas) + eps = 1e-6 + if np.fabs(sigmas[0] - 1) < eps: + # to avoid inf torch.log(alpha_si) in multistep_uni_p_bh_update during first/second update + sigmas[0] -= eps timesteps = (sigmas * self.config.num_train_timesteps).copy() if self.config.final_sigmas_type == "sigma_min": sigma_last = sigmas[-1] @@ -410,6 +449,8 @@ def set_timesteps( ) sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) else: + if sigmas is None: + sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) if self.config.final_sigmas_type == "sigma_min": sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5 @@ -439,6 +480,58 @@ def set_timesteps( self._begin_index = None self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.time_shift + def time_shift(self, mu: float, sigma: float, t: torch.Tensor): + """ + Apply time shifting to the sigmas. + + Args: + mu (`float`): + The mu parameter for the time shift. + sigma (`float`): + The sigma parameter for the time shift. + t (`torch.Tensor`): + The input timesteps. + + Returns: + `torch.Tensor`: + The time-shifted timesteps. + """ + if self.config.time_shift_type == "exponential": + return self._time_shift_exponential(mu, sigma, t) + elif self.config.time_shift_type == "linear": + return self._time_shift_linear(mu, sigma, t) + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.stretch_shift_to_terminal + def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor: + r""" + Stretches and shifts the timestep schedule to ensure it terminates at the configured `shift_terminal` config + value. + + Reference: + https://github.com/Lightricks/LTX-Video/blob/a01a171f8fe3d99dce2728d60a73fecf4d4238ae/ltx_video/schedulers/rf.py#L51 + + Args: + t (`torch.Tensor`): + A tensor of timesteps to be stretched and shifted. + + Returns: + `torch.Tensor`: + A tensor of adjusted timesteps such that the final value equals `self.config.shift_terminal`. + """ + one_minus_z = 1 - t + scale_factor = one_minus_z[-1] / (1 - self.config.shift_terminal) + stretched_t = 1 - (one_minus_z / scale_factor) + return stretched_t + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._time_shift_exponential + def _time_shift_exponential(self, mu, sigma, t): + return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._time_shift_linear + def _time_shift_linear(self, mu, sigma, t): + return mu / (mu + (1 / t - 1) ** sigma) + # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: """ @@ -521,7 +614,7 @@ def _sigma_to_t(self, sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray: return t # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._sigma_to_alpha_sigma_t - def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """ Convert sigma values to alpha_t and sigma_t values. @@ -530,7 +623,7 @@ def _sigma_to_alpha_sigma_t(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, to The sigma value(s) to convert. Returns: - `Tuple[torch.Tensor, torch.Tensor]`: + `tuple[torch.Tensor, torch.Tensor]`: A tuple containing (alpha_t, sigma_t) values. """ if self.config.use_flow_sigmas: @@ -1005,7 +1098,9 @@ def multistep_uni_c_bh_update( # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep def index_for_timestep( - self, timestep: Union[int, torch.Tensor], schedule_timesteps: Optional[torch.Tensor] = None + self, + timestep: int | torch.Tensor, + schedule_timesteps: torch.Tensor | None = None, ) -> int: """ Find the index for a given timestep in the schedule. @@ -1039,7 +1134,7 @@ def index_for_timestep( return step_index # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index - def _init_step_index(self, timestep: Union[int, torch.Tensor]) -> None: + def _init_step_index(self, timestep: int | torch.Tensor) -> None: """ Initialize the step_index counter for the scheduler. @@ -1058,10 +1153,10 @@ def _init_step_index(self, timestep: Union[int, torch.Tensor]) -> None: def step( self, model_output: torch.Tensor, - timestep: Union[int, torch.Tensor], + timestep: int | torch.Tensor, sample: torch.Tensor, return_dict: bool = True, - ) -> Union[SchedulerOutput, Tuple]: + ) -> SchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the multistep UniPC. diff --git a/src/diffusers/schedulers/scheduling_utils.py b/src/diffusers/schedulers/scheduling_utils.py index a355c7bb1a51..0d48d791ac42 100644 --- a/src/diffusers/schedulers/scheduling_utils.py +++ b/src/diffusers/schedulers/scheduling_utils.py @@ -15,7 +15,6 @@ import os from dataclasses import dataclass from enum import Enum -from typing import Optional, Union import torch from huggingface_hub.utils import validate_hf_hub_args @@ -83,7 +82,7 @@ class SchedulerMixin(PushToHubMixin): the scheduler's `__init__` function, and the attributes can be accessed by `scheduler.config.num_train_timesteps`. Class attributes: - - **_compatibles** (`List[str]`) -- A list of scheduler classes that are compatible with the parent scheduler + - **_compatibles** (`list[str]`) -- A list of scheduler classes that are compatible with the parent scheduler class. Use [`~ConfigMixin.from_config`] to load a different compatible scheduler class (should be overridden by parent class). """ @@ -96,8 +95,8 @@ class SchedulerMixin(PushToHubMixin): @validate_hf_hub_args def from_pretrained( cls, - pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None, - subfolder: Optional[str] = None, + pretrained_model_name_or_path: str | os.PathLike | None = None, + subfolder: str | None = None, return_unused_kwargs=False, **kwargs, ) -> Self: @@ -116,14 +115,14 @@ def from_pretrained( The subfolder location of a model file within a larger model repository on the Hub or locally. return_unused_kwargs (`bool`, *optional*, defaults to `False`): Whether kwargs that are not consumed by the Python class should be returned or not. - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. output_loading_info(`bool`, *optional*, defaults to `False`): @@ -153,7 +152,7 @@ def from_pretrained( ) return cls.from_config(config, return_unused_kwargs=return_unused_kwargs, **kwargs) - def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs): + def save_pretrained(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs): """ Save a scheduler configuration object to a directory so that it can be reloaded using the [`~SchedulerMixin.from_pretrained`] class method. @@ -165,7 +164,7 @@ def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the repository you want to push to with `repo_id` (will default to the name of `save_directory` in your namespace). - kwargs (`Dict[str, Any]`, *optional*): + kwargs (`dict[str, Any]`, *optional*): Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. """ self.save_config(save_directory=save_directory, push_to_hub=push_to_hub, **kwargs) @@ -176,7 +175,7 @@ def compatibles(self): Returns all schedulers that are compatible with this scheduler Returns: - `List[SchedulerMixin]`: List of compatible schedulers + `list[SchedulerMixin]`: list of compatible schedulers """ return self._get_compatibles() diff --git a/src/diffusers/schedulers/scheduling_utils_flax.py b/src/diffusers/schedulers/scheduling_utils_flax.py index 0534e47d8a30..44de56a3980d 100644 --- a/src/diffusers/schedulers/scheduling_utils_flax.py +++ b/src/diffusers/schedulers/scheduling_utils_flax.py @@ -16,7 +16,6 @@ import os from dataclasses import dataclass from enum import Enum -from typing import Optional, Tuple, Union import flax import jax.numpy as jnp @@ -62,7 +61,7 @@ class FlaxSchedulerMixin(PushToHubMixin): Mixin containing common functions for the schedulers. Class attributes: - - **_compatibles** (`List[str]`) -- A list of classes that are compatible with the parent class, so that + - **_compatibles** (`list[str]`) -- A list of classes that are compatible with the parent class, so that `from_config` can be used from a class different than the one used to save the config (should be overridden by parent class). """ @@ -76,8 +75,8 @@ class FlaxSchedulerMixin(PushToHubMixin): @validate_hf_hub_args def from_pretrained( cls, - pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None, - subfolder: Optional[str] = None, + pretrained_model_name_or_path: str | os.PathLike | None = None, + subfolder: str | None = None, return_unused_kwargs=False, **kwargs, ): @@ -98,14 +97,14 @@ def from_pretrained( return_unused_kwargs (`bool`, *optional*, defaults to `False`): Whether kwargs that are not consumed by the Python class should be returned or not. - cache_dir (`Union[str, os.PathLike]`, *optional*): + cache_dir (`str | os.PathLike`, *optional*): Path to a directory in which a downloaded pretrained model configuration should be cached if the standard cache should not be used. force_download (`bool`, *optional*, defaults to `False`): Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. output_loading_info(`bool`, *optional*, defaults to `False`): @@ -148,7 +147,7 @@ def from_pretrained( return scheduler, state - def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs): + def save_pretrained(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs): """ Save a scheduler configuration object to the directory `save_directory`, so that it can be re-loaded using the [`~FlaxSchedulerMixin.from_pretrained`] class method. @@ -160,7 +159,7 @@ def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the repository you want to push to with `repo_id` (will default to the name of `save_directory` in your namespace). - kwargs (`Dict[str, Any]`, *optional*): + kwargs (`dict[str, Any]`, *optional*): Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. """ self.save_config(save_directory=save_directory, push_to_hub=push_to_hub, **kwargs) @@ -171,7 +170,7 @@ def compatibles(self): Returns all schedulers that are compatible with this scheduler Returns: - `List[SchedulerMixin]`: List of compatible schedulers + `list[SchedulerMixin]`: list of compatible schedulers """ return self._get_compatibles() @@ -185,7 +184,7 @@ def _get_compatibles(cls): return compatible_classes -def broadcast_to_shape_from_left(x: jnp.ndarray, shape: Tuple[int]) -> jnp.ndarray: +def broadcast_to_shape_from_left(x: jnp.ndarray, shape: tuple[int]) -> jnp.ndarray: assert len(shape) >= x.ndim return jnp.broadcast_to(x.reshape(x.shape + (1,) * (len(shape) - x.ndim)), shape) diff --git a/src/diffusers/schedulers/scheduling_vq_diffusion.py b/src/diffusers/schedulers/scheduling_vq_diffusion.py index 57306301d023..7ba3a8e6f5c4 100644 --- a/src/diffusers/schedulers/scheduling_vq_diffusion.py +++ b/src/diffusers/schedulers/scheduling_vq_diffusion.py @@ -13,7 +13,6 @@ # limitations under the License. from dataclasses import dataclass -from typing import Optional, Tuple, Union import numpy as np import torch @@ -59,7 +58,7 @@ def index_to_log_onehot(x: torch.LongTensor, num_classes: int) -> torch.Tensor: return log_x -def gumbel_noised(logits: torch.Tensor, generator: Optional[torch.Generator]) -> torch.Tensor: +def gumbel_noised(logits: torch.Tensor, generator: torch.Generator | None) -> torch.Tensor: """ Apply gumbel noise to `logits` """ @@ -175,7 +174,7 @@ def __init__( self.num_inference_steps = None self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy()) - def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None): + def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -202,9 +201,9 @@ def step( model_output: torch.Tensor, timestep: torch.long, sample: torch.LongTensor, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, return_dict: bool = True, - ) -> Union[VQDiffusionSchedulerOutput, Tuple]: + ) -> VQDiffusionSchedulerOutput | tuple: """ Predict the sample from the previous timestep by the reverse transition distribution. See [`~VQDiffusionScheduler.q_posterior`] for more details about how the distribution is computer. diff --git a/src/diffusers/training_utils.py b/src/diffusers/training_utils.py index 7a98fa3da14a..6c07a30c2ccc 100644 --- a/src/diffusers/training_utils.py +++ b/src/diffusers/training_utils.py @@ -6,11 +6,18 @@ import re import warnings from contextlib import contextmanager -from typing import Any, Dict, Iterable, List, Optional, Tuple, Union +from functools import partial +from typing import Any, Iterable import numpy as np import torch + +if getattr(torch, "distributed", None) is not None: + from torch.distributed.fsdp import CPUOffload, ShardingStrategy + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy + from .models import UNet2DConditionModel from .pipelines import DiffusionPipeline from .schedulers import SchedulerMixin @@ -18,6 +25,7 @@ convert_state_dict_to_diffusers, convert_state_dict_to_peft, deprecate, + is_accelerate_available, is_peft_available, is_torch_npu_available, is_torchvision_available, @@ -31,6 +39,9 @@ if transformers.integrations.deepspeed.is_deepspeed_zero3_enabled(): import deepspeed +if is_accelerate_available(): + from accelerate.logging import get_logger + if is_peft_available(): from peft import set_peft_model_state_dict @@ -151,7 +162,7 @@ def compute_dream_and_update_latents( target: torch.Tensor, encoder_hidden_states: torch.Tensor, dream_detail_preservation: float = 1.0, -) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: +) -> tuple[torch.Tensor, torch.Tensor]: """ Implements "DREAM (Diffusion Rectification and Estimation-Adaptive Models)" from https://huggingface.co/papers/2312.00210. DREAM helps align training with sampling to help training be more @@ -196,7 +207,7 @@ def compute_dream_and_update_latents( return _noisy_latents, _target -def unet_lora_state_dict(unet: UNet2DConditionModel) -> Dict[str, torch.Tensor]: +def unet_lora_state_dict(unet: UNet2DConditionModel) -> dict[str, torch.Tensor]: r""" Returns: A state dict containing just the LoRA parameters. @@ -215,7 +226,7 @@ def unet_lora_state_dict(unet: UNet2DConditionModel) -> Dict[str, torch.Tensor]: return lora_state_dict -def cast_training_params(model: Union[torch.nn.Module, List[torch.nn.Module]], dtype=torch.float32): +def cast_training_params(model: torch.nn.Module | list[torch.nn.Module], dtype=torch.float32): """ Casts the training parameters of the model to the specified data type. @@ -233,7 +244,7 @@ def cast_training_params(model: Union[torch.nn.Module, List[torch.nn.Module]], d def _set_state_dict_into_text_encoder( - lora_state_dict: Dict[str, torch.Tensor], prefix: str, text_encoder: torch.nn.Module + lora_state_dict: dict[str, torch.Tensor], prefix: str, text_encoder: torch.nn.Module ): """ Sets the `lora_state_dict` into `text_encoder` coming from `transformers`. @@ -251,7 +262,7 @@ def _set_state_dict_into_text_encoder( set_peft_model_state_dict(text_encoder, text_encoder_state_dict, adapter_name="default") -def _collate_lora_metadata(modules_to_save: Dict[str, torch.nn.Module]) -> Dict[str, Any]: +def _collate_lora_metadata(modules_to_save: dict[str, torch.nn.Module]) -> dict[str, Any]: metadatas = {} for module_name, module in modules_to_save.items(): if module is not None: @@ -265,8 +276,8 @@ def compute_density_for_timestep_sampling( logit_mean: float = None, logit_std: float = None, mode_scale: float = None, - device: Union[torch.device, str] = "cpu", - generator: Optional[torch.Generator] = None, + device: torch.device | str = "cpu", + generator: torch.Generator | None = None, ): """ Compute the density for sampling the timesteps when doing SD3 training. @@ -321,9 +332,7 @@ def free_memory(): @contextmanager -def offload_models( - *modules: Union[torch.nn.Module, DiffusionPipeline], device: Union[str, torch.device], offload: bool = True -): +def offload_models(*modules: torch.nn.Module | DiffusionPipeline, device: str | torch.device, offload: bool = True): """ Context manager that, if offload=True, moves each module to `device` on enter, then moves it back to its original device on exit. @@ -394,6 +403,83 @@ def find_nearest_bucket(h, w, bucket_options): return best_bucket_idx +def _to_cpu_contiguous(state_dicts) -> dict: + return {k: v.detach().cpu().contiguous() if isinstance(v, torch.Tensor) else v for k, v in state_dicts.items()} + + +def get_fsdp_kwargs_from_accelerator(accelerator) -> dict: + """ + Extract and convert FSDP config from Accelerator into PyTorch FSDP kwargs. + """ + + kwargs = {} + fsdp_state = getattr(accelerator.state, "fsdp_plugin", None) + + if fsdp_state is None: + raise ValueError("Accelerate isn't configured to handle FSDP. Please update your installation.") + + fsdp_plugin = accelerator.state.fsdp_plugin + + if fsdp_plugin is None: + # FSDP not enabled in Accelerator + kwargs["sharding_strategy"] = ShardingStrategy.FULL_SHARD + else: + # FSDP is enabled → use plugin's strategy, or default if None + kwargs["sharding_strategy"] = fsdp_plugin.sharding_strategy or ShardingStrategy.FULL_SHARD + + return kwargs + + +def wrap_with_fsdp( + model: torch.nn.Module, + device: str | torch.device, + offload: bool = True, + use_orig_params: bool = True, + limit_all_gathers: bool = True, + fsdp_kwargs: dict[str, Any] | None = None, + transformer_layer_cls: set[type[torch.nn.Module]] | None = None, +) -> FSDP: + """ + Wrap a model with FSDP using common defaults and optional transformer auto-wrapping. + + Args: + model: Model to wrap + device: Target device (e.g., accelerator.device) + offload: Whether to enable CPU parameter offloading + use_orig_params: Whether to use original parameters + limit_all_gathers: Whether to limit all gathers + fsdp_kwargs: FSDP arguments (sharding_strategy, etc.) — usually from Accelerate config + transformer_layer_cls: Classes for auto-wrapping (if not using policy from fsdp_kwargs) + + Returns: + FSDP-wrapped model + """ + + logger = get_logger(__name__) + + if transformer_layer_cls is None: + # Set the default layers if transformer_layer_cls is not provided + transformer_layer_cls = type(model.model.language_model.layers[0]) + logger.info(f"transformer_layer_cls is not provided, auto-inferred as {transformer_layer_cls.__name__}") + + # Add auto-wrap policy if transformer layers specified + auto_wrap_policy = partial(transformer_auto_wrap_policy, transformer_layer_cls={transformer_layer_cls}) + + config = { + "device_id": device, + "cpu_offload": CPUOffload(offload_params=offload) if offload else None, + "use_orig_params": use_orig_params, + "limit_all_gathers": limit_all_gathers, + "auto_wrap_policy": auto_wrap_policy, + } + + if fsdp_kwargs: + config.update(fsdp_kwargs) + + fsdp_model = FSDP(model, **config) + return fsdp_model + + # Adapted from torch-ema https://github.com/fadel/pytorch_ema/blob/master/torch_ema/ema.py#L14 class EMAModel: """ @@ -407,11 +493,11 @@ def __init__( min_decay: float = 0.0, update_after_step: int = 0, use_ema_warmup: bool = False, - inv_gamma: Union[float, int] = 1.0, - power: Union[float, int] = 2 / 3, + inv_gamma: float | int = 1.0, + power: float | int = 2 / 3, foreach: bool = False, - model_cls: Optional[Any] = None, - model_config: Dict[str, Any] = None, + model_cls: Any | None = None, + model_config: dict[str, Any] | None = None, **kwargs, ): """ @@ -425,7 +511,7 @@ def __init__( Inverse multiplicative factor of EMA warmup. Default: 1. Only used if `use_ema_warmup` is True. power (float): Exponential factor of EMA warmup. Default: 2/3. Only used if `use_ema_warmup` is True. foreach (bool): Use torch._foreach functions for updating shadow parameters. Should be faster. - device (Optional[Union[str, torch.device]]): The device to store the EMA weights on. If None, the EMA + device (str | torch.device | None): The device to store the EMA weights on. If None, the EMA weights will be stored on CPU. @crowsonkb's notes on EMA Warmup: diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index 440e4539e720..23d7ac7c6c2d 100644 --- a/src/diffusers/utils/__init__.py +++ b/src/diffusers/utils/__init__.py @@ -23,6 +23,7 @@ DEFAULT_HF_PARALLEL_LOADING_WORKERS, DEPRECATED_REVISION_ARGS, DIFFUSERS_DYNAMIC_MODULE_NAME, + DIFFUSERS_LOAD_ID_FIELDS, FLAX_WEIGHTS_NAME, GGUF_FILE_EXTENSION, HF_ENABLE_PARALLEL_LOADING, @@ -66,6 +67,7 @@ is_accelerate_version, is_aiter_available, is_aiter_version, + is_av_available, is_better_profanity_available, is_bitsandbytes_available, is_bitsandbytes_version, @@ -83,9 +85,8 @@ is_hpu_available, is_inflect_available, is_invisible_watermark_available, - is_k_diffusion_available, - is_k_diffusion_version, is_kernels_available, + is_kernels_version, is_kornia_available, is_librosa_available, is_matplotlib_available, @@ -129,6 +130,7 @@ from .logging import get_logger from .outputs import BaseOutput from .peft_utils import ( + apply_lora_scale, check_peft_version, delete_adapter_layers, get_adapter_name, diff --git a/src/diffusers/utils/constants.py b/src/diffusers/utils/constants.py index c46fa4363483..4f94df656a65 100644 --- a/src/diffusers/utils/constants.py +++ b/src/diffusers/utils/constants.py @@ -73,3 +73,11 @@ ENCODE_ENDPOINT_SD_V1 = "https://qc6479g0aac6qwy9.us-east-1.aws.endpoints.huggingface.cloud/" ENCODE_ENDPOINT_SD_XL = "https://xjqqhmyn62rog84g.us-east-1.aws.endpoints.huggingface.cloud/" ENCODE_ENDPOINT_FLUX = "https://ptccx55jz97f9zgo.us-east-1.aws.endpoints.huggingface.cloud/" + + +DIFFUSERS_LOAD_ID_FIELDS = [ + "pretrained_model_name_or_path", + "subfolder", + "variant", + "revision", +] diff --git a/src/diffusers/utils/deprecation_utils.py b/src/diffusers/utils/deprecation_utils.py index d76623541b9f..7ec612499e9a 100644 --- a/src/diffusers/utils/deprecation_utils.py +++ b/src/diffusers/utils/deprecation_utils.py @@ -1,6 +1,6 @@ import inspect import warnings -from typing import Any, Dict, Optional, Union +from typing import Any from packaging import version @@ -26,7 +26,7 @@ } -def _maybe_remap_transformers_class(class_name: str) -> Optional[str]: +def _maybe_remap_transformers_class(class_name: str) -> str | None: """ Check if a Transformers class should be remapped to a newer version. @@ -53,7 +53,7 @@ def _maybe_remap_transformers_class(class_name: str) -> Optional[str]: return None -def deprecate(*args, take_from: Optional[Union[Dict, Any]] = None, standard_warn=True, stacklevel=2): +def deprecate(*args, take_from: dict | Any | None = None, standard_warn=True, stacklevel=2): from .. import __version__ deprecated_kwargs = take_from diff --git a/src/diffusers/utils/distributed_utils.py b/src/diffusers/utils/distributed_utils.py new file mode 100644 index 000000000000..239b7b26200d --- /dev/null +++ b/src/diffusers/utils/distributed_utils.py @@ -0,0 +1,36 @@ +# Copyright 2025 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +try: + import torch +except ImportError: + torch = None + + +def is_torch_dist_rank_zero() -> bool: + if torch is None: + return True + + dist_module = getattr(torch, "distributed", None) + if dist_module is None or not dist_module.is_available(): + return True + + if not dist_module.is_initialized(): + return True + + try: + return dist_module.get_rank() == 0 + except (RuntimeError, ValueError): + return True diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index eb956a96a30f..a607d41dcba1 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -227,6 +227,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class MagCacheConfig(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class PyramidAttentionBroadcastConfig(metaclass=DummyObject): _backends = ["torch"] @@ -284,6 +299,10 @@ def apply_layer_skip(*args, **kwargs): requires_backends(apply_layer_skip, ["torch"]) +def apply_mag_cache(*args, **kwargs): + requires_backends(apply_mag_cache, ["torch"]) + + def apply_pyramid_attention_broadcast(*args, **kwargs): requires_backends(apply_pyramid_attention_broadcast, ["torch"]) @@ -502,6 +521,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class AutoencoderKLLTX2Audio(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + +class AutoencoderKLLTX2Video(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class AutoencoderKLLTXVideo(metaclass=DummyObject): _backends = ["torch"] @@ -607,6 +656,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class AutoencoderRAE(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class AutoencoderTiny(metaclass=DummyObject): _backends = ["torch"] @@ -697,6 +761,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class ChromaRadianceTransformer2DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class ChronoEditTransformer3DModel(metaclass=DummyObject): _backends = ["torch"] @@ -847,6 +926,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class CosmosControlNetModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class CosmosTransformer3DModel(metaclass=DummyObject): _backends = ["torch"] @@ -952,6 +1046,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class GlmImageTransformer2DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + +class HeliosTransformer3DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class HiDreamImageTransformer2DModel(metaclass=DummyObject): _backends = ["torch"] @@ -1147,6 +1271,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class LTX2VideoTransformer3DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class LTXVideoTransformer3DModel(metaclass=DummyObject): _backends = ["torch"] @@ -1811,6 +1950,21 @@ def attention_backend(*args, **kwargs): requires_backends(attention_backend, ["torch"]) +class AutoPipelineBlocks(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class ComponentsManager(metaclass=DummyObject): _backends = ["torch"] @@ -1841,6 +1995,66 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class ConditionalPipelineBlocks(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + +class ConfigSpec(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + +class InputParam(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + +class LoopSequentialPipelineBlocks(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class ModularPipeline(metaclass=DummyObject): _backends = ["torch"] @@ -1871,6 +2085,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class OutputParam(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + +class SequentialPipelineBlocks(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + def get_constant_schedule(*args, **kwargs): requires_backends(get_constant_schedule, ["torch"]) @@ -2544,6 +2788,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class HeliosDMDScheduler(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + +class HeliosScheduler(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class HeunDiscreteScheduler(metaclass=DummyObject): _backends = ["torch"] @@ -2634,6 +2908,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class LTXEulerAncestralRFScheduler(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class PNDMScheduler(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_and_k_diffusion_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_and_k_diffusion_objects.py deleted file mode 100644 index 2ab00c54ce2d..000000000000 --- a/src/diffusers/utils/dummy_torch_and_transformers_and_k_diffusion_objects.py +++ /dev/null @@ -1,32 +0,0 @@ -# This file is autogenerated by the command `make fix-copies`, do not edit. -from ..utils import DummyObject, requires_backends - - -class StableDiffusionKDiffusionPipeline(metaclass=DummyObject): - _backends = ["torch", "transformers", "k_diffusion"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "transformers", "k_diffusion"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "k_diffusion"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "k_diffusion"]) - - -class StableDiffusionXLKDiffusionPipeline(metaclass=DummyObject): - _backends = ["torch", "transformers", "k_diffusion"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "transformers", "k_diffusion"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "k_diffusion"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers", "k_diffusion"]) diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 6c28e87581b9..475e076a072d 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -17,6 +17,66 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class Flux2KleinAutoBlocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class Flux2KleinBaseAutoBlocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class Flux2KleinBaseModularPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class Flux2KleinModularPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class Flux2ModularPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] @@ -167,6 +227,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class QwenImageLayeredAutoBlocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class QwenImageLayeredModularPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class QwenImageModularPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] @@ -212,7 +302,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) -class Wan22AutoBlocks(metaclass=DummyObject): +class Wan22Blocks(metaclass=DummyObject): _backends = ["torch", "transformers"] def __init__(self, *args, **kwargs): @@ -227,7 +317,82 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) -class WanAutoBlocks(metaclass=DummyObject): +class Wan22Image2VideoBlocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class Wan22Image2VideoModularPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class Wan22ModularPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class WanBlocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class WanImage2VideoAutoBlocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class WanImage2VideoModularPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] def __init__(self, *args, **kwargs): @@ -557,6 +722,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class BriaFiboEditPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class BriaFiboPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] @@ -602,6 +782,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class ChromaInpaintPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class ChromaPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] @@ -617,6 +812,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class ChromaRadiancePipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class ChronoEditPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] @@ -782,6 +992,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class Cosmos2_5_TransferPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class Cosmos2TextToImagePipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] @@ -902,6 +1127,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class Flux2KleinPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class Flux2Pipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] @@ -1112,6 +1352,51 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class GlmImagePipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class HeliosPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class HeliosPyramidPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class HiDreamImagePipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] @@ -1877,6 +2162,66 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class LTX2ConditionPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class LTX2ImageToVideoPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class LTX2LatentUpsamplePipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class LTX2Pipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class LTXConditionPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] @@ -1892,6 +2237,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class LTXI2VLongMultiPromptPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class LTXImageToVideoPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] @@ -3917,6 +4277,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class ZImageInpaintPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class ZImageOmniPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/src/diffusers/utils/dynamic_modules_utils.py b/src/diffusers/utils/dynamic_modules_utils.py index 1c0734cf35bb..856966dd29b5 100644 --- a/src/diffusers/utils/dynamic_modules_utils.py +++ b/src/diffusers/utils/dynamic_modules_utils.py @@ -24,7 +24,6 @@ import threading from pathlib import Path from types import ModuleType -from typing import Dict, Optional, Union from urllib import request from huggingface_hub import hf_hub_download, model_info @@ -65,7 +64,7 @@ def init_hf_modules(): init_path.touch() -def create_dynamic_module(name: Union[str, os.PathLike]): +def create_dynamic_module(name: str | os.PathLike): """ Creates a dynamic module in the cache directory for modules. """ @@ -197,7 +196,7 @@ def get_class_in_module(class_name, module_path, force_reload=False): if force_reload: sys.modules.pop(name, None) importlib.invalidate_caches() - cached_module: Optional[ModuleType] = sys.modules.get(name) + cached_module: ModuleType | None = sys.modules.get(name) module_spec = importlib.util.spec_from_file_location(name, location=module_file) module: ModuleType @@ -245,16 +244,16 @@ def find_pipeline_class(loaded_module): @validate_hf_hub_args def get_cached_module_file( - pretrained_model_name_or_path: Union[str, os.PathLike], + pretrained_model_name_or_path: str | os.PathLike, module_file: str, - subfolder: Optional[str] = None, - cache_dir: Optional[Union[str, os.PathLike]] = None, + subfolder: str | None = None, + cache_dir: str | os.PathLike | None = None, force_download: bool = False, - proxies: Optional[Dict[str, str]] = None, - token: Optional[Union[bool, str]] = None, - revision: Optional[str] = None, + proxies: dict[str, str] | None = None, + token: bool | str | None = None, + revision: str | None = None, local_files_only: bool = False, - local_dir: Optional[str] = None, + local_dir: str | None = None, ): """ Prepares Downloads a module from a local folder or a distant repo and returns its path inside the cached @@ -278,7 +277,7 @@ def get_cached_module_file( force_download (`bool`, *optional*, defaults to `False`): Whether or not to force to (re-)download the configuration files and override the cached versions if they exist. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request. token (`str` or *bool*, *optional*): @@ -300,7 +299,10 @@ def get_cached_module_file( # Download and cache module_file from the repo `pretrained_model_name_or_path` of grab it if it's a local file. pretrained_model_name_or_path = str(pretrained_model_name_or_path) - module_file_or_url = os.path.join(pretrained_model_name_or_path, module_file) + if subfolder is not None: + module_file_or_url = os.path.join(pretrained_model_name_or_path, subfolder, module_file) + else: + module_file_or_url = os.path.join(pretrained_model_name_or_path, module_file) if os.path.isfile(module_file_or_url): resolved_module_file = module_file_or_url @@ -385,7 +387,11 @@ def get_cached_module_file( if not os.path.exists(submodule_path / module_folder): os.makedirs(submodule_path / module_folder) module_needed = f"{module_needed}.py" - shutil.copyfile(os.path.join(pretrained_model_name_or_path, module_needed), submodule_path / module_needed) + if subfolder is not None: + source_path = os.path.join(pretrained_model_name_or_path, subfolder, module_needed) + else: + source_path = os.path.join(pretrained_model_name_or_path, module_needed) + shutil.copyfile(source_path, submodule_path / module_needed) else: # Get the commit hash # TODO: we will get this info in the etag soon, so retrieve it from there and not here. @@ -426,17 +432,17 @@ def get_cached_module_file( @validate_hf_hub_args def get_class_from_dynamic_module( - pretrained_model_name_or_path: Union[str, os.PathLike], + pretrained_model_name_or_path: str | os.PathLike, module_file: str, - subfolder: Optional[str] = None, - class_name: Optional[str] = None, - cache_dir: Optional[Union[str, os.PathLike]] = None, + subfolder: str | None = None, + class_name: str | None = None, + cache_dir: str | os.PathLike | None = None, force_download: bool = False, - proxies: Optional[Dict[str, str]] = None, - token: Optional[Union[bool, str]] = None, - revision: Optional[str] = None, + proxies: dict[str, str] | None = None, + token: bool | str | None = None, + revision: str | None = None, local_files_only: bool = False, - local_dir: Optional[str] = None, + local_dir: str | None = None, ): """ Extracts a class from a module file, present in the local folder or repository of a model. @@ -464,7 +470,7 @@ def get_class_from_dynamic_module( force_download (`bool`, *optional*, defaults to `False`): Whether or not to force to (re-)download the configuration files and override the cached versions if they exist. - proxies (`Dict[str, str]`, *optional*): + proxies (`dict[str, str]`, *optional*): A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request. token (`str` or `bool`, *optional*): diff --git a/src/diffusers/utils/export_utils.py b/src/diffusers/utils/export_utils.py index 07cf46928a44..d9b4e173a862 100644 --- a/src/diffusers/utils/export_utils.py +++ b/src/diffusers/utils/export_utils.py @@ -3,7 +3,6 @@ import struct import tempfile from contextlib import contextmanager -from typing import List, Optional, Union import numpy as np import PIL.Image @@ -25,7 +24,7 @@ def buffered_writer(raw_f): f.flush() -def export_to_gif(image: List[PIL.Image.Image], output_gif_path: str = None, fps: int = 10) -> str: +def export_to_gif(image: list[PIL.Image.Image], output_gif_path: str = None, fps: int = 10) -> str: if output_gif_path is None: output_gif_path = tempfile.NamedTemporaryFile(suffix=".gif").name @@ -113,7 +112,7 @@ def export_to_obj(mesh, output_obj_path: str = None): def _legacy_export_to_video( - video_frames: Union[List[np.ndarray], List[PIL.Image.Image]], output_video_path: str = None, fps: int = 10 + video_frames: list[np.ndarray] | list[PIL.Image.Image], output_video_path: str = None, fps: int = 10 ): if is_opencv_available(): import cv2 @@ -139,12 +138,12 @@ def _legacy_export_to_video( def export_to_video( - video_frames: Union[List[np.ndarray], List[PIL.Image.Image]], + video_frames: list[np.ndarray] | list[PIL.Image.Image], output_video_path: str = None, fps: int = 10, quality: float = 5.0, - bitrate: Optional[int] = None, - macro_block_size: Optional[int] = 16, + bitrate: int | None = None, + macro_block_size: int | None = 16, ) -> str: """ quality: diff --git a/src/diffusers/utils/hub_utils.py b/src/diffusers/utils/hub_utils.py index d0b05c7d9541..b5eb9ab2e17f 100644 --- a/src/diffusers/utils/hub_utils.py +++ b/src/diffusers/utils/hub_utils.py @@ -21,7 +21,6 @@ import tempfile import warnings from pathlib import Path -from typing import Dict, List, Optional, Union from uuid import uuid4 from huggingface_hub import ( @@ -72,7 +71,7 @@ SESSION_ID = uuid4().hex -def http_user_agent(user_agent: Union[Dict, str, None] = None) -> str: +def http_user_agent(user_agent: dict | str | None = None) -> str: """ Formats a user-agent string with basic info about a request. """ @@ -98,15 +97,17 @@ def http_user_agent(user_agent: Union[Dict, str, None] = None) -> str: def load_or_create_model_card( repo_id_or_path: str = None, - token: Optional[str] = None, + token: str | None = None, is_pipeline: bool = False, from_training: bool = False, - model_description: Optional[str] = None, + model_description: str | None = None, base_model: str = None, - prompt: Optional[str] = None, - license: Optional[str] = None, - widget: Optional[List[dict]] = None, - inference: Optional[bool] = None, + prompt: str | None = None, + license: str | None = None, + widget: list[dict] | None = None, + inference: bool | None = None, + is_modular: bool = False, + update_model_card: bool = False, ) -> ModelCard: """ Loads or creates a model card. @@ -128,9 +129,14 @@ def load_or_create_model_card( prompt (`str`, *optional*): Prompt used for training. Useful for DreamBooth-like training. license: (`str`, *optional*): License of the output artifact. Helpful when using `load_or_create_model_card` from a training script. - widget (`List[dict]`, *optional*): Widget to accompany a gallery template. + widget (`list[dict]`, *optional*): Widget to accompany a gallery template. inference: (`bool`, optional): Whether to turn on inference widget. Helpful when using `load_or_create_model_card` from a training script. + is_modular: (`bool`, optional): Boolean flag to denote if the model card is for a modular pipeline. + When True, uses model_description as-is without additional template formatting. + update_model_card: (`bool`, optional): When True, regenerates the model card content even if one + already exists on the remote repo. Existing card metadata (tags, license, etc.) is preserved. Only + supported for modular pipelines (i.e., `is_modular=True`). """ if not is_jinja_available(): raise ValueError( @@ -139,9 +145,17 @@ def load_or_create_model_card( " To install it, please run `pip install Jinja2`." ) + if update_model_card and not is_modular: + raise ValueError("`update_model_card=True` is only supported for modular pipelines (`is_modular=True`).") + try: # Check if the model card is present on the remote repo model_card = ModelCard.load(repo_id_or_path, token=token) + # For modular pipelines, regenerate card content when requested (preserve existing metadata) + if update_model_card and is_modular and model_description is not None: + existing_data = model_card.data + model_card = ModelCard(model_description) + model_card.data = existing_data except (EntryNotFoundError, RepositoryNotFoundError): # Otherwise create a model card from template if from_training: @@ -159,15 +173,19 @@ def load_or_create_model_card( ) else: card_data = ModelCardData() - component = "pipeline" if is_pipeline else "model" - if model_description is None: - model_description = f"This is the model card of a 🧨 diffusers {component} that has been pushed on the Hub. This model card has been automatically generated." - model_card = ModelCard.from_template(card_data, model_description=model_description) + if is_modular and model_description is not None: + model_card = ModelCard(model_description) + model_card.data = card_data + else: + component = "pipeline" if is_pipeline else "model" + if model_description is None: + model_description = f"This is the model card of a 🧨 diffusers {component} that has been pushed on the Hub. This model card has been automatically generated." + model_card = ModelCard.from_template(card_data, model_description=model_description) return model_card -def populate_model_card(model_card: ModelCard, tags: Union[str, List[str]] = None) -> ModelCard: +def populate_model_card(model_card: ModelCard, tags: str | list[str] | None = None) -> ModelCard: """Populates the `model_card` with library name and optional tags.""" if model_card.data.library_name is None: model_card.data.library_name = "diffusers" @@ -183,7 +201,7 @@ def populate_model_card(model_card: ModelCard, tags: Union[str, List[str]] = Non return model_card -def extract_commit_hash(resolved_file: Optional[str], commit_hash: Optional[str] = None): +def extract_commit_hash(resolved_file: str | None, commit_hash: str | None = None): """ Extracts the commit hash from a resolved filename toward a cache file. """ @@ -197,7 +215,7 @@ def extract_commit_hash(resolved_file: Optional[str], commit_hash: Optional[str] return commit_hash if REGEX_COMMIT_HASH.match(commit_hash) else None -def _add_variant(weights_name: str, variant: Optional[str] = None) -> str: +def _add_variant(weights_name: str, variant: str | None = None) -> str: if variant is not None: splits = weights_name.split(".") splits = splits[:-1] + [variant] + splits[-1:] @@ -208,19 +226,19 @@ def _add_variant(weights_name: str, variant: Optional[str] = None) -> str: @validate_hf_hub_args def _get_model_file( - pretrained_model_name_or_path: Union[str, Path], + pretrained_model_name_or_path: str | Path, *, weights_name: str, - subfolder: Optional[str] = None, - cache_dir: Optional[str] = None, + subfolder: str | None = None, + cache_dir: str | None = None, force_download: bool = False, - proxies: Optional[Dict] = None, + proxies: dict | None = None, local_files_only: bool = False, - token: Optional[str] = None, - user_agent: Optional[Union[Dict, str]] = None, - revision: Optional[str] = None, - commit_hash: Optional[str] = None, - dduf_entries: Optional[Dict[str, DDUFEntry]] = None, + token: str | None = None, + user_agent: dict | str | None = None, + revision: str | None = None, + commit_hash: str | None = None, + dduf_entries: dict[str, DDUFEntry] | None = None, ): pretrained_model_name_or_path = str(pretrained_model_name_or_path) @@ -348,7 +366,7 @@ def _get_checkpoint_shard_files( user_agent=None, revision=None, subfolder="", - dduf_entries: Optional[Dict[str, DDUFEntry]] = None, + dduf_entries: dict[str, DDUFEntry] | None = None, ): """ For a given model: @@ -449,7 +467,7 @@ def _get_checkpoint_shard_files( return cached_filenames, sharded_metadata -def _check_legacy_sharding_variant_format(folder: str = None, filenames: List[str] = None, variant: str = None): +def _check_legacy_sharding_variant_format(folder: str = None, filenames: list[str] = None, variant: str = None): if filenames and folder: raise ValueError("Both `filenames` and `folder` cannot be provided.") if not filenames: @@ -469,12 +487,12 @@ class PushToHubMixin: def _upload_folder( self, - working_dir: Union[str, os.PathLike], + working_dir: str | os.PathLike, repo_id: str, - token: Optional[str] = None, - commit_message: Optional[str] = None, + token: str | None = None, + commit_message: str | None = None, create_pr: bool = False, - subfolder: Optional[str] = None, + subfolder: str | None = None, ): """ Uploads all files in `working_dir` to `repo_id`. @@ -500,13 +518,13 @@ def _upload_folder( def push_to_hub( self, repo_id: str, - commit_message: Optional[str] = None, - private: Optional[bool] = None, - token: Optional[str] = None, + commit_message: str | None = None, + private: bool | None = None, + token: str | None = None, create_pr: bool = False, safe_serialization: bool = True, - variant: Optional[str] = None, - subfolder: Optional[str] = None, + variant: str | None = None, + subfolder: str | None = None, ) -> str: """ Upload model, scheduler, or pipeline files to the 🤗 Hugging Face Hub. diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py index 57b0a337922a..551fa358a28d 100644 --- a/src/diffusers/utils/import_utils.py +++ b/src/diffusers/utils/import_utils.py @@ -24,7 +24,7 @@ from functools import lru_cache as cache from itertools import chain from types import ModuleType -from typing import Any, Tuple, Union +from typing import Any from huggingface_hub.utils import is_jinja_available # noqa: F401 from packaging.version import Version, parse @@ -59,7 +59,7 @@ _is_google_colab = "google.colab" in sys.modules or any(k.startswith("COLAB_") for k in os.environ) -def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> Tuple[bool, str]: +def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> tuple[bool, str]: global _package_map pkg_exists = importlib.util.find_spec(pkg_name) is not None pkg_version = "N/A" @@ -198,7 +198,7 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> Tuple[b _kernels_available, _kernels_version = _is_package_available("kernels") _inflect_available, _inflect_version = _is_package_available("inflect") _unidecode_available, _unidecode_version = _is_package_available("unidecode") -_k_diffusion_available, _k_diffusion_version = _is_package_available("k_diffusion") + _note_seq_available, _note_seq_version = _is_package_available("note_seq") _wandb_available, _wandb_version = _is_package_available("wandb") _tensorboard_available, _tensorboard_version = _is_package_available("tensorboard") @@ -227,9 +227,10 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> Tuple[b _sageattention_available, _sageattention_version = _is_package_available("sageattention") _flash_attn_available, _flash_attn_version = _is_package_available("flash_attn") _flash_attn_3_available, _flash_attn_3_version = _is_package_available("flash_attn_3") -_aiter_available, _aiter_version = _is_package_available("aiter") +_aiter_available, _aiter_version = _is_package_available("aiter", get_dist_name=True) _kornia_available, _kornia_version = _is_package_available("kornia") _nvidia_modelopt_available, _nvidia_modelopt_version = _is_package_available("modelopt", get_dist_name=True) +_av_available, _av_version = _is_package_available("av") def is_torch_available(): @@ -292,10 +293,6 @@ def is_kernels_available(): return _kernels_available -def is_k_diffusion_available(): - return _k_diffusion_available - - def is_note_seq_available(): return _note_seq_available @@ -420,6 +417,10 @@ def is_kornia_available(): return _kornia_available +def is_av_available(): + return _av_available + + # docstyle-ignore FLAX_IMPORT_ERROR = """ {0} requires the FLAX library but it was not found in your environment. Checkout the instructions on the @@ -474,12 +475,6 @@ def is_kornia_available(): Unidecode` """ -# docstyle-ignore -K_DIFFUSION_IMPORT_ERROR = """ -{0} requires the k-diffusion library but it was not found in your environment. You can install it with pip: `pip -install k-diffusion` -""" - # docstyle-ignore NOTE_SEQ_IMPORT_ERROR = """ {0} requires the note-seq library but it was not found in your environment. You can install it with pip: `pip @@ -596,7 +591,6 @@ def is_kornia_available(): ("transformers", (is_transformers_available, TRANSFORMERS_IMPORT_ERROR)), ("unidecode", (is_unidecode_available, UNIDECODE_IMPORT_ERROR)), ("librosa", (is_librosa_available, LIBROSA_IMPORT_ERROR)), - ("k_diffusion", (is_k_diffusion_available, K_DIFFUSION_IMPORT_ERROR)), ("note_seq", (is_note_seq_available, NOTE_SEQ_IMPORT_ERROR)), ("wandb", (is_wandb_available, WANDB_IMPORT_ERROR)), ("tensorboard", (is_tensorboard_available, TENSORBOARD_IMPORT_ERROR)), @@ -663,7 +657,7 @@ def __getattr__(cls, key): # This function was copied from: https://github.com/huggingface/accelerate/blob/874c4967d94badd24f893064cc3bef45f57cadf7/src/accelerate/utils/versions.py#L319 -def compare_versions(library_or_version: Union[str, Version], operation: str, requirement_version: str): +def compare_versions(library_or_version: str | Version, operation: str, requirement_version: str): """ Compares a library version to some requirement using a given operation. @@ -730,6 +724,22 @@ def is_transformers_version(operation: str, version: str): return compare_versions(parse(_transformers_version), operation, version) +@cache +def is_kernels_version(operation: str, version: str): + """ + Compares the current Kernels version to a given reference with an operation. + + Args: + operation (`str`): + A string representation of an operator, such as `">"` or `"<="` + version (`str`): + A version string + """ + if not _kernels_available: + return False + return compare_versions(parse(_kernels_version), operation, version) + + @cache def is_hf_hub_version(operation: str, version: str): """ @@ -825,22 +835,6 @@ def is_torchao_version(operation: str, version: str): return compare_versions(parse(_torchao_version), operation, version) -@cache -def is_k_diffusion_version(operation: str, version: str): - """ - Compares the current k-diffusion version to a given reference with an operation. - - Args: - operation (`str`): - A string representation of an operator, such as `">"` or `"<="` - version (`str`): - A version string - """ - if not _k_diffusion_available: - return False - return compare_versions(parse(_k_diffusion_version), operation, version) - - @cache def is_optimum_quanto_version(operation: str, version: str): """ diff --git a/src/diffusers/utils/loading_utils.py b/src/diffusers/utils/loading_utils.py index dd23ae73c861..c4fee0cfdd83 100644 --- a/src/diffusers/utils/loading_utils.py +++ b/src/diffusers/utils/loading_utils.py @@ -1,6 +1,6 @@ import os import tempfile -from typing import Any, Callable, List, Optional, Tuple, Union +from typing import Any, Callable from urllib.parse import unquote, urlparse import PIL.Image @@ -12,7 +12,7 @@ def load_image( - image: Union[str, PIL.Image.Image], convert_method: Optional[Callable[[PIL.Image.Image], PIL.Image.Image]] = None + image: str | PIL.Image.Image, convert_method: Callable[[PIL.Image.Image], PIL.Image.Image] | None = None ) -> PIL.Image.Image: """ Loads `image` to a PIL Image. @@ -56,20 +56,20 @@ def load_image( def load_video( video: str, - convert_method: Optional[Callable[[List[PIL.Image.Image]], List[PIL.Image.Image]]] = None, -) -> List[PIL.Image.Image]: + convert_method: Callable[[list[PIL.Image.Image]], list[PIL.Image.Image]] | None = None, +) -> list[PIL.Image.Image]: """ Loads `video` to a list of PIL Image. Args: video (`str`): A URL or Path to a video to convert to a list of PIL Image format. - convert_method (Callable[[List[PIL.Image.Image]], List[PIL.Image.Image]], *optional*): + convert_method (Callable[[list[PIL.Image.Image]], list[PIL.Image.Image]], *optional*): A conversion method to apply to the video after loading it. When set to `None` the images will be converted to "RGB". Returns: - `List[PIL.Image.Image]`: + `list[PIL.Image.Image]`: The video as a list of PIL images. """ is_url = video.startswith("http://") or video.startswith("https://") @@ -139,7 +139,7 @@ def load_video( # Taken from `transformers`. -def get_module_from_name(module, tensor_name: str) -> Tuple[Any, str]: +def get_module_from_name(module, tensor_name: str) -> tuple[Any, str]: if "." in tensor_name: splits = tensor_name.split(".") for split in splits[:-1]: diff --git a/src/diffusers/utils/logging.py b/src/diffusers/utils/logging.py index 2ad6d3a47607..21c4bba8a830 100644 --- a/src/diffusers/utils/logging.py +++ b/src/diffusers/utils/logging.py @@ -28,13 +28,14 @@ WARN, # NOQA WARNING, # NOQA ) -from typing import Dict, Optional from tqdm import auto as tqdm_lib +from .distributed_utils import is_torch_dist_rank_zero + _lock = threading.Lock() -_default_handler: Optional[logging.Handler] = None +_default_handler: logging.Handler | None = None log_levels = { "debug": logging.DEBUG, @@ -47,6 +48,23 @@ _default_log_level = logging.WARNING _tqdm_active = True +_rank_zero_filter = None + + +class _RankZeroFilter(logging.Filter): + def filter(self, record): + # Always allow rank-zero logs, but keep debug-level messages from all ranks for troubleshooting. + return is_torch_dist_rank_zero() or record.levelno <= logging.DEBUG + + +def _ensure_rank_zero_filter(logger: logging.Logger) -> None: + global _rank_zero_filter + + if _rank_zero_filter is None: + _rank_zero_filter = _RankZeroFilter() + + if not any(isinstance(f, _RankZeroFilter) for f in logger.filters): + logger.addFilter(_rank_zero_filter) def _get_default_logging_level() -> int: @@ -90,6 +108,7 @@ def _configure_library_root_logger() -> None: library_root_logger.addHandler(_default_handler) library_root_logger.setLevel(_get_default_logging_level()) library_root_logger.propagate = False + _ensure_rank_zero_filter(library_root_logger) def _reset_library_root_logger() -> None: @@ -105,11 +124,11 @@ def _reset_library_root_logger() -> None: _default_handler = None -def get_log_levels_dict() -> Dict[str, int]: +def get_log_levels_dict() -> dict[str, int]: return log_levels -def get_logger(name: Optional[str] = None) -> logging.Logger: +def get_logger(name: str | None = None) -> logging.Logger: """ Return a logger with the specified name. @@ -120,7 +139,9 @@ def get_logger(name: Optional[str] = None) -> logging.Logger: name = _get_library_name() _configure_library_root_logger() - return logging.getLogger(name) + logger = logging.getLogger(name) + _ensure_rank_zero_filter(logger) + return logger def get_verbosity() -> int: diff --git a/src/diffusers/utils/outputs.py b/src/diffusers/utils/outputs.py index 2b20f6120ce3..b7d9a29ccfce 100644 --- a/src/diffusers/utils/outputs.py +++ b/src/diffusers/utils/outputs.py @@ -17,7 +17,7 @@ from collections import OrderedDict from dataclasses import fields, is_dataclass -from typing import Any, Tuple +from typing import Any import numpy as np @@ -127,7 +127,7 @@ def __reduce__(self): args = tuple(getattr(self, field.name) for field in fields(self)) return callable, args, *remaining - def to_tuple(self) -> Tuple[Any, ...]: + def to_tuple(self) -> tuple[Any, ...]: """ Convert self to a tuple containing all the attributes/keys that are not `None`. """ diff --git a/src/diffusers/utils/peft_utils.py b/src/diffusers/utils/peft_utils.py index 12066ee3f89b..65bcfe631e97 100644 --- a/src/diffusers/utils/peft_utils.py +++ b/src/diffusers/utils/peft_utils.py @@ -16,8 +16,8 @@ """ import collections +import functools import importlib -from typing import Optional from packaging import version @@ -123,7 +123,7 @@ def scale_lora_layers(model, weight): module.scale_layer(weight) -def unscale_lora_layers(model, weight: Optional[float] = None): +def unscale_lora_layers(model, weight: float | None = None): """ Removes the previously passed weight given to the LoRA layers of the model. @@ -275,6 +275,55 @@ def get_module_weight(weight_for_adapter, module_name): module.set_scale(adapter_name, get_module_weight(weight, module_name)) +def apply_lora_scale(kwargs_name: str = "joint_attention_kwargs"): + """ + Decorator to automatically handle LoRA layer scaling/unscaling in forward methods. + + This decorator extracts the `lora_scale` from the specified kwargs parameter, applies scaling before the forward + pass, and ensures unscaling happens after, even if an exception occurs. + + Args: + kwargs_name (`str`, defaults to `"joint_attention_kwargs"`): + The name of the keyword argument that contains the LoRA scale. Common values include + "joint_attention_kwargs", "attention_kwargs", "cross_attention_kwargs", etc. + """ + + def decorator(forward_fn): + @functools.wraps(forward_fn) + def wrapper(self, *args, **kwargs): + from . import USE_PEFT_BACKEND + + lora_scale = 1.0 + attention_kwargs = kwargs.get(kwargs_name) + + if attention_kwargs is not None: + attention_kwargs = attention_kwargs.copy() + kwargs[kwargs_name] = attention_kwargs + lora_scale = attention_kwargs.pop("scale", 1.0) + + if not USE_PEFT_BACKEND and lora_scale != 1.0: + logger.warning( + f"Passing `scale` via `{kwargs_name}` when not using the PEFT backend is ineffective." + ) + + # Apply LoRA scaling if using PEFT backend + if USE_PEFT_BACKEND: + scale_lora_layers(self, lora_scale) + + try: + # Execute the forward pass + result = forward_fn(self, *args, **kwargs) + return result + finally: + # Always unscale, even if forward pass raises an exception + if USE_PEFT_BACKEND: + unscale_lora_layers(self, lora_scale) + + return wrapper + + return decorator + + def check_peft_version(min_version: str) -> None: r""" Checks if the version of PEFT is compatible. diff --git a/src/diffusers/utils/pil_utils.py b/src/diffusers/utils/pil_utils.py index 76678070b697..72d4704fa945 100644 --- a/src/diffusers/utils/pil_utils.py +++ b/src/diffusers/utils/pil_utils.py @@ -1,5 +1,3 @@ -from typing import List - import PIL.Image import PIL.ImageOps from packaging import version @@ -50,7 +48,7 @@ def numpy_to_pil(images): return pil_images -def make_image_grid(images: List[PIL.Image.Image], rows: int, cols: int, resize: int = None) -> PIL.Image.Image: +def make_image_grid(images: list[PIL.Image.Image], rows: int, cols: int, resize: int = None) -> PIL.Image.Image: """ Prepares a single grid of images. Useful for visualization purposes. """ diff --git a/src/diffusers/utils/remote_utils.py b/src/diffusers/utils/remote_utils.py index 6494dc14171a..2580bbf5dda9 100644 --- a/src/diffusers/utils/remote_utils.py +++ b/src/diffusers/utils/remote_utils.py @@ -13,9 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import io import json -from typing import List, Literal, Optional, Union, cast +from typing import Literal, cast import requests @@ -58,18 +60,18 @@ def detect_image_type(data: bytes) -> str: def check_inputs_decode( endpoint: str, tensor: "torch.Tensor", - processor: Optional[Union["VaeImageProcessor", "VideoProcessor"]] = None, + processor: "VaeImageProcessor" | "VideoProcessor" | None = None, do_scaling: bool = True, - scaling_factor: Optional[float] = None, - shift_factor: Optional[float] = None, + scaling_factor: float | None = None, + shift_factor: float | None = None, output_type: Literal["mp4", "pil", "pt"] = "pil", return_type: Literal["mp4", "pil", "pt"] = "pil", image_format: Literal["png", "jpg"] = "jpg", partial_postprocess: bool = False, input_tensor_type: Literal["binary"] = "binary", output_tensor_type: Literal["binary"] = "binary", - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, ): if tensor.ndim == 3 and height is None and width is None: raise ValueError("`height` and `width` required for packed latents.") @@ -91,7 +93,7 @@ def check_inputs_decode( def postprocess_decode( response: requests.Response, - processor: Optional[Union["VaeImageProcessor", "VideoProcessor"]] = None, + processor: "VaeImageProcessor" | "VideoProcessor" | None = None, output_type: Literal["mp4", "pil", "pt"] = "pil", return_type: Literal["mp4", "pil", "pt"] = "pil", partial_postprocess: bool = False, @@ -117,7 +119,7 @@ def postprocess_decode( else: if isinstance(processor, VideoProcessor): output = cast( - List[Image.Image], + list[Image.Image], processor.postprocess_video(output_tensor, output_type="pil")[0], ) else: @@ -144,15 +146,15 @@ def postprocess_decode( def prepare_decode( tensor: "torch.Tensor", - processor: Optional[Union["VaeImageProcessor", "VideoProcessor"]] = None, + processor: "VaeImageProcessor" | "VideoProcessor" | None = None, do_scaling: bool = True, - scaling_factor: Optional[float] = None, - shift_factor: Optional[float] = None, + scaling_factor: float | None = None, + shift_factor: float | None = None, output_type: Literal["mp4", "pil", "pt"] = "pil", image_format: Literal["png", "jpg"] = "jpg", partial_postprocess: bool = False, - height: Optional[int] = None, - width: Optional[int] = None, + height: int | None = None, + width: int | None = None, ): headers = {} parameters = { @@ -188,19 +190,19 @@ def prepare_decode( def remote_decode( endpoint: str, tensor: "torch.Tensor", - processor: Optional[Union["VaeImageProcessor", "VideoProcessor"]] = None, + processor: "VaeImageProcessor" | "VideoProcessor" | None = None, do_scaling: bool = True, - scaling_factor: Optional[float] = None, - shift_factor: Optional[float] = None, + scaling_factor: float | None = None, + shift_factor: float | None = None, output_type: Literal["mp4", "pil", "pt"] = "pil", return_type: Literal["mp4", "pil", "pt"] = "pil", image_format: Literal["png", "jpg"] = "jpg", partial_postprocess: bool = False, input_tensor_type: Literal["binary"] = "binary", output_tensor_type: Literal["binary"] = "binary", - height: Optional[int] = None, - width: Optional[int] = None, -) -> Union[Image.Image, List[Image.Image], bytes, "torch.Tensor"]: + height: int | None = None, + width: int | None = None, +) -> Image.Image | list[Image.Image] | bytes | "torch.Tensor": """ Hugging Face Hybrid Inference that allow running VAE decode remotely. @@ -275,7 +277,7 @@ def remote_decode( Required for `"packed"` latents. Returns: - output (`Image.Image` or `List[Image.Image]` or `bytes` or `torch.Tensor`). + output (`Image.Image` or `list[Image.Image]` or `bytes` or `torch.Tensor`). """ if input_tensor_type == "base64": deprecate( @@ -336,9 +338,9 @@ def remote_decode( def check_inputs_encode( endpoint: str, - image: Union["torch.Tensor", Image.Image], - scaling_factor: Optional[float] = None, - shift_factor: Optional[float] = None, + image: "torch.Tensor" | Image.Image, + scaling_factor: float | None = None, + shift_factor: float | None = None, ): pass @@ -356,9 +358,9 @@ def postprocess_encode( def prepare_encode( - image: Union["torch.Tensor", Image.Image], - scaling_factor: Optional[float] = None, - shift_factor: Optional[float] = None, + image: "torch.Tensor" | Image.Image, + scaling_factor: float | None = None, + shift_factor: float | None = None, ): headers = {} parameters = {} @@ -379,9 +381,9 @@ def prepare_encode( def remote_encode( endpoint: str, - image: Union["torch.Tensor", Image.Image], - scaling_factor: Optional[float] = None, - shift_factor: Optional[float] = None, + image: "torch.Tensor" | Image.Image, + scaling_factor: float | None = None, + shift_factor: float | None = None, ) -> "torch.Tensor": """ Hugging Face Hybrid Inference that allow running VAE encode remotely. diff --git a/src/diffusers/utils/testing_utils.py b/src/diffusers/utils/testing_utils.py index 3297bb5fdcd6..cb7bf942c648 100644 --- a/src/diffusers/utils/testing_utils.py +++ b/src/diffusers/utils/testing_utils.py @@ -19,7 +19,7 @@ from contextlib import contextmanager from io import BytesIO, StringIO from pathlib import Path -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple, Union +from typing import TYPE_CHECKING, Any, Callable, Set import numpy as np import PIL.Image @@ -676,7 +676,7 @@ def get_python_version(): return major, minor -def load_numpy(arry: Union[str, np.ndarray], local_path: Optional[str] = None) -> np.ndarray: +def load_numpy(arry: str | np.ndarray, local_path: str | None = None) -> np.ndarray: if isinstance(arry, str): if local_path is not None: # local_path can be passed to correct images of tests @@ -702,14 +702,14 @@ def load_numpy(arry: Union[str, np.ndarray], local_path: Optional[str] = None) - return arry -def load_pt(url: str, map_location: Optional[str] = None, weights_only: Optional[bool] = True): +def load_pt(url: str, map_location: str | None = None, weights_only: bool | None = True): response = requests.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT) response.raise_for_status() arry = torch.load(BytesIO(response.content), map_location=map_location, weights_only=weights_only) return arry -def load_image(image: Union[str, PIL.Image.Image]) -> PIL.Image.Image: +def load_image(image: str | PIL.Image.Image) -> PIL.Image.Image: """ Loads `image` to a PIL Image. @@ -750,7 +750,7 @@ def preprocess_image(image: PIL.Image, batch_size: int): return 2.0 * image - 1.0 -def export_to_gif(image: List[PIL.Image.Image], output_gif_path: str = None) -> str: +def export_to_gif(image: list[PIL.Image.Image], output_gif_path: str = None) -> str: if output_gif_path is None: output_gif_path = tempfile.NamedTemporaryFile(suffix=".gif").name @@ -843,7 +843,7 @@ def export_to_obj(mesh, output_obj_path: str = None): f.writelines("\n".join(combined_data)) -def export_to_video(video_frames: List[np.ndarray], output_video_path: str = None) -> str: +def export_to_video(video_frames: list[np.ndarray], output_video_path: str = None) -> str: if is_opencv_available(): import cv2 else: @@ -1024,7 +1024,7 @@ def summary_failures_short(tr): # Adapted from https://github.com/huggingface/transformers/blob/000e52aec8850d3fe2f360adc6fd256e5b47fe4c/src/transformers/testing_utils.py#L1905 -def is_flaky(max_attempts: int = 5, wait_before_retry: Optional[float] = None, description: Optional[str] = None): +def is_flaky(max_attempts: int = 5, wait_before_retry: float | None = None, description: str | None = None): """ To decorate flaky tests (methods or entire classes). They will be retried on failures. @@ -1281,7 +1281,7 @@ def _is_torch_fp64_available(device): # This dispatches a defined function according to the accelerator from the function definitions. -def _device_agnostic_dispatch(device: str, dispatch_table: Dict[str, Callable], *args, **kwargs): +def _device_agnostic_dispatch(device: str, dispatch_table: dict[str, Callable], *args, **kwargs): if device not in dispatch_table: return dispatch_table["default"](*args, **kwargs) @@ -1381,7 +1381,7 @@ def backend_supports_training(device: str): # Guard for when Torch is not available if is_torch_available(): # Update device function dict mapping - def update_mapping_from_spec(device_fn_dict: Dict[str, Callable], attribute_name: str): + def update_mapping_from_spec(device_fn_dict: dict[str, Callable], attribute_name: str): try: # Try to import the function directly spec_fn = getattr(device_spec_module, attribute_name) @@ -1430,7 +1430,7 @@ def update_mapping_from_spec(device_fn_dict: Dict[str, Callable], attribute_name # Modified from https://github.com/huggingface/transformers/blob/cdfb018d0300fef3b07d9220f3efe9c2a9974662/src/transformers/testing_utils.py#L3090 # Type definition of key used in `Expectations` class. -DeviceProperties = Tuple[Union[str, None], Union[int, None]] +DeviceProperties = tuple[str | None, int | None] @functools.lru_cache @@ -1477,7 +1477,7 @@ def _get_expected_safetensors_files( module: torch.nn.Module, offload_to_disk_path: str, offload_type: str, - num_blocks_per_group: Optional[int] = None, + num_blocks_per_group: int | None = None, ) -> Set[str]: expected_files = set() @@ -1545,7 +1545,7 @@ def _check_safetensors_serialization( module: torch.nn.Module, offload_to_disk_path: str, offload_type: str, - num_blocks_per_group: Optional[int] = None, + num_blocks_per_group: int | None = None, ) -> bool: if not os.path.isdir(offload_to_disk_path): return False, None, None diff --git a/src/diffusers/utils/torch_utils.py b/src/diffusers/utils/torch_utils.py index 3b66fdadbef8..7f4cb3e12766 100644 --- a/src/diffusers/utils/torch_utils.py +++ b/src/diffusers/utils/torch_utils.py @@ -15,9 +15,10 @@ PyTorch utilities: Utilities related to PyTorch """ +from __future__ import annotations + import functools import os -from typing import Callable, Dict, List, Optional, Tuple, Union from . import logging from .import_utils import is_torch_available, is_torch_mlu_available, is_torch_npu_available, is_torch_version @@ -88,7 +89,7 @@ def maybe_allow_in_graph(cls): # This dispatches a defined function according to the accelerator from the function definitions. -def _device_agnostic_dispatch(device: str, dispatch_table: Dict[str, Callable], *args, **kwargs): +def _device_agnostic_dispatch(device: str, dispatch_table: dict[str, callable], *args, **kwargs): if device not in dispatch_table: return dispatch_table["default"](*args, **kwargs) @@ -144,11 +145,11 @@ def backend_supports_training(device: str): def randn_tensor( - shape: Union[Tuple, List], - generator: Optional[Union[List["torch.Generator"], "torch.Generator"]] = None, - device: Optional[Union[str, "torch.device"]] = None, - dtype: Optional["torch.dtype"] = None, - layout: Optional["torch.layout"] = None, + shape: tuple | list, + generator: list["torch.Generator"] | "torch.Generator" | None = None, + device: str | "torch.device" | None = None, + dtype: "torch.dtype" | None = None, + layout: "torch.layout" | None = None, ): """A helper function to create random tensors on the desired `device` with the desired `dtype`. When passing a list of generators, you can seed each batch size individually. If CPU generators are passed, the tensor @@ -241,9 +242,9 @@ def fourier_filter(x_in: "torch.Tensor", threshold: int, scale: int) -> "torch.T def apply_freeu( resolution_idx: int, hidden_states: "torch.Tensor", res_hidden_states: "torch.Tensor", **freeu_kwargs -) -> Tuple["torch.Tensor", "torch.Tensor"]: - """Applies the FreeU mechanism as introduced in https://huggingface.co/papers/2309.11497. Adapted from the official - code repository: https://github.com/ChenyangSi/FreeU. +) -> tuple["torch.Tensor", "torch.Tensor"]: + """Applies the FreeU mechanism as introduced in https: + //arxiv.org/abs/2309.11497. Adapted from the official code repository: https://github.com/ChenyangSi/FreeU. Args: resolution_idx (`int`): Integer denoting the UNet block where FreeU is being applied. @@ -292,7 +293,7 @@ def get_device(): return "cpu" -def empty_device_cache(device_type: Optional[str] = None): +def empty_device_cache(device_type: str | None = None): if device_type is None: device_type = get_device() if device_type in ["cpu"]: @@ -301,7 +302,7 @@ def empty_device_cache(device_type: Optional[str] = None): device_mod.empty_cache() -def device_synchronize(device_type: Optional[str] = None): +def device_synchronize(device_type: str | None = None): if device_type is None: device_type = get_device() device_mod = getattr(torch, device_type, torch.cuda) diff --git a/src/diffusers/utils/typing_utils.py b/src/diffusers/utils/typing_utils.py index 2b5b1a4f5ab5..e910932f33d9 100644 --- a/src/diffusers/utils/typing_utils.py +++ b/src/diffusers/utils/typing_utils.py @@ -18,7 +18,7 @@ from typing import Any, Dict, List, Set, Tuple, Type, Union, get_args, get_origin -def _is_valid_type(obj: Any, class_or_tuple: Union[Type, Tuple[Type, ...]]) -> bool: +def _is_valid_type(obj: Any, class_or_tuple: Type | tuple[Type, ...]) -> bool: """ Checks if an object is an instance of any of the provided types. For collections, it checks if every element is of the correct type as well. diff --git a/src/diffusers/utils/versions.py b/src/diffusers/utils/versions.py index 945a3977ce62..ccede7b5ee57 100644 --- a/src/diffusers/utils/versions.py +++ b/src/diffusers/utils/versions.py @@ -19,7 +19,6 @@ import operator import re import sys -from typing import Optional from packaging import version @@ -46,7 +45,7 @@ def _compare_versions(op, got_ver, want_ver, requirement, pkg, hint): ) -def require_version(requirement: str, hint: Optional[str] = None) -> None: +def require_version(requirement: str, hint: str | None = None) -> None: """ Perform a runtime check of the dependency versions, using the exact same syntax used by pip. diff --git a/src/diffusers/video_processor.py b/src/diffusers/video_processor.py index abeb30bca102..0c51b4b38f23 100644 --- a/src/diffusers/video_processor.py +++ b/src/diffusers/video_processor.py @@ -13,7 +13,6 @@ # limitations under the License. import warnings -from typing import List, Optional, Tuple, Union import numpy as np import PIL @@ -26,20 +25,20 @@ class VideoProcessor(VaeImageProcessor): r"""Simple video processor.""" - def preprocess_video(self, video, height: Optional[int] = None, width: Optional[int] = None) -> torch.Tensor: + def preprocess_video(self, video, height: int | None = None, width: int | None = None, **kwargs) -> torch.Tensor: r""" - Preprocesses input video(s). + Preprocesses input video(s). Keyword arguments will be forwarded to `VaeImageProcessor.preprocess`. Args: - video (`List[PIL.Image]`, `List[List[PIL.Image]]`, `torch.Tensor`, `np.array`, `List[torch.Tensor]`, `List[np.array]`): + video (`list[PIL.Image]`, `list[list[PIL.Image]]`, `torch.Tensor`, `np.array`, `list[torch.Tensor]`, `list[np.array]`): The input video. It can be one of the following: - * List of the PIL images. - * List of list of PIL images. + * list of the PIL images. + * list of list of PIL images. * 4D Torch tensors (expected shape for each tensor `(num_frames, num_channels, height, width)`). * 4D NumPy arrays (expected shape for each array `(num_frames, height, width, num_channels)`). - * List of 4D Torch tensors (expected shape for each tensor `(num_frames, num_channels, height, + * list of 4D Torch tensors (expected shape for each tensor `(num_frames, num_channels, height, width)`). - * List of 4D NumPy arrays (expected shape for each array `(num_frames, height, width, num_channels)`). + * list of 4D NumPy arrays (expected shape for each array `(num_frames, height, width, num_channels)`). * 5D NumPy arrays: expected shape for each array `(batch_size, num_frames, height, width, num_channels)`. * 5D Torch tensors: expected shape for each array `(batch_size, num_frames, num_channels, height, @@ -50,6 +49,10 @@ def preprocess_video(self, video, height: Optional[int] = None, width: Optional[ width (`int`, *optional*`, defaults to `None`): The width in preprocessed frames of the video. If `None`, will use get_default_height_width()` to get the default width. + + Returns: + `torch.Tensor` of shape `(batch_size, num_channels, num_frames, height, width)`: + A 5D tensor holding the batched channels-first video(s). """ if isinstance(video, list) and isinstance(video[0], np.ndarray) and video[0].ndim == 5: warnings.warn( @@ -80,7 +83,7 @@ def preprocess_video(self, video, height: Optional[int] = None, width: Optional[ "Input is in incorrect format. Currently, we only support numpy.ndarray, torch.Tensor, PIL.Image.Image" ) - video = torch.stack([self.preprocess(img, height=height, width=width) for img in video], dim=0) + video = torch.stack([self.preprocess(img, height=height, width=width, **kwargs) for img in video], dim=0) # move the number of channels before the number of frames. video = video.permute(0, 2, 1, 3, 4) @@ -88,10 +91,11 @@ def preprocess_video(self, video, height: Optional[int] = None, width: Optional[ return video def postprocess_video( - self, video: torch.Tensor, output_type: str = "np" - ) -> Union[np.ndarray, torch.Tensor, List[PIL.Image.Image]]: + self, video: torch.Tensor, output_type: str = "np", **kwargs + ) -> np.ndarray | torch.Tensor | list[PIL.Image.Image]: r""" - Converts a video tensor to a list of frames for export. + Converts a video tensor to a list of frames for export. Keyword arguments will be forwarded to + `VaeImageProcessor.postprocess`. Args: video (`torch.Tensor`): The video as a tensor. @@ -101,7 +105,7 @@ def postprocess_video( outputs = [] for batch_idx in range(batch_size): batch_vid = video[batch_idx].permute(1, 0, 2, 3) - batch_output = self.postprocess(batch_vid, output_type) + batch_output = self.postprocess(batch_vid, output_type, **kwargs) outputs.append(batch_output) if output_type == "np": @@ -114,7 +118,7 @@ def postprocess_video( return outputs @staticmethod - def classify_height_width_bin(height: int, width: int, ratios: dict) -> Tuple[int, int]: + def classify_height_width_bin(height: int, width: int, ratios: dict) -> tuple[int, int]: r""" Returns the binned height and width based on the aspect ratio. @@ -124,7 +128,7 @@ def classify_height_width_bin(height: int, width: int, ratios: dict) -> Tuple[in ratios (`dict`): A dictionary where keys are aspect ratios and values are tuples of (height, width). Returns: - `Tuple[int, int]`: The closest binned height and width. + `tuple[int, int]`: The closest binned height and width. """ ar = float(height / width) closest_ratio = min(ratios.keys(), key=lambda ratio: abs(float(ratio) - ar)) diff --git a/tests/conftest.py b/tests/conftest.py index fd76d1c84ee7..1d7e83467f03 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -32,6 +32,24 @@ def pytest_configure(config): config.addinivalue_line("markers", "big_accelerator: marks tests as requiring big accelerator resources") + config.addinivalue_line("markers", "lora: marks tests for LoRA/PEFT functionality") + config.addinivalue_line("markers", "ip_adapter: marks tests for IP Adapter functionality") + config.addinivalue_line("markers", "training: marks tests for training functionality") + config.addinivalue_line("markers", "attention: marks tests for attention processor functionality") + config.addinivalue_line("markers", "memory: marks tests for memory optimization functionality") + config.addinivalue_line("markers", "cpu_offload: marks tests for CPU offloading functionality") + config.addinivalue_line("markers", "group_offload: marks tests for group offloading functionality") + config.addinivalue_line("markers", "compile: marks tests for torch.compile functionality") + config.addinivalue_line("markers", "single_file: marks tests for single file checkpoint loading") + config.addinivalue_line("markers", "quantization: marks tests for quantization functionality") + config.addinivalue_line("markers", "bitsandbytes: marks tests for BitsAndBytes quantization functionality") + config.addinivalue_line("markers", "quanto: marks tests for Quanto quantization functionality") + config.addinivalue_line("markers", "torchao: marks tests for TorchAO quantization functionality") + config.addinivalue_line("markers", "gguf: marks tests for GGUF quantization functionality") + config.addinivalue_line("markers", "modelopt: marks tests for NVIDIA ModelOpt quantization functionality") + config.addinivalue_line("markers", "context_parallel: marks tests for context parallel inference functionality") + config.addinivalue_line("markers", "slow: mark test as slow") + config.addinivalue_line("markers", "nightly: mark test as nightly") def pytest_addoption(parser): diff --git a/tests/fixtures/custom_pipeline/pipeline.py b/tests/fixtures/custom_pipeline/pipeline.py index 25673e566549..2d569d0f8abe 100644 --- a/tests/fixtures/custom_pipeline/pipeline.py +++ b/tests/fixtures/custom_pipeline/pipeline.py @@ -15,8 +15,6 @@ # limitations under the License. -from typing import Optional, Tuple, Union - import torch from diffusers import DiffusionPipeline, ImagePipelineOutput, SchedulerMixin, UNet2DModel @@ -42,12 +40,12 @@ def __init__(self, unet: UNet2DModel, scheduler: SchedulerMixin): def __call__( self, batch_size: int = 1, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, num_inference_steps: int = 50, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, **kwargs, - ) -> Union[ImagePipelineOutput, Tuple]: + ) -> ImagePipelineOutput | tuple: r""" Args: batch_size (`int`, *optional*, defaults to 1): diff --git a/tests/fixtures/custom_pipeline/what_ever.py b/tests/fixtures/custom_pipeline/what_ever.py index 7504940780e8..afb140f64d32 100644 --- a/tests/fixtures/custom_pipeline/what_ever.py +++ b/tests/fixtures/custom_pipeline/what_ever.py @@ -15,8 +15,6 @@ # limitations under the License. -from typing import Optional, Tuple, Union - import torch from diffusers import SchedulerMixin, UNet2DModel @@ -43,12 +41,12 @@ def __init__(self, unet: UNet2DModel, scheduler: SchedulerMixin): def __call__( self, batch_size: int = 1, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, num_inference_steps: int = 50, - output_type: Optional[str] = "pil", + output_type: str | None = "pil", return_dict: bool = True, **kwargs, - ) -> Union[ImagePipelineOutput, Tuple]: + ) -> ImagePipelineOutput | tuple: r""" Args: batch_size (`int`, *optional*, defaults to 1): diff --git a/tests/hooks/test_group_offloading.py b/tests/hooks/test_group_offloading.py index 236094109d07..108a7247bcc6 100644 --- a/tests/hooks/test_group_offloading.py +++ b/tests/hooks/test_group_offloading.py @@ -566,3 +566,127 @@ def get_autoencoder_kl_config(self, block_out_channels=None, norm_num_groups=Non "layers_per_block": 1, } return init_dict + + +# Model with conditionally-executed modules, simulating Helios patch_short/patch_mid/patch_long behavior. +# These modules are only called when optional inputs are provided, which means the lazy prefetch +# execution order tracer may not see them on the first forward pass. This can cause a device mismatch +# on subsequent calls when the modules ARE invoked but their weights were never onloaded. +# See: https://github.com/huggingface/diffusers/pull/13211 +class DummyModelWithConditionalModules(ModelMixin): + def __init__(self, in_features: int, hidden_features: int, out_features: int, num_layers: int) -> None: + super().__init__() + + self.linear_1 = torch.nn.Linear(in_features, hidden_features) + self.activation = torch.nn.ReLU() + self.blocks = torch.nn.ModuleList( + [DummyBlock(hidden_features, hidden_features, hidden_features) for _ in range(num_layers)] + ) + self.linear_2 = torch.nn.Linear(hidden_features, out_features) + + # These modules are only invoked when optional_input is not None. + # Output dimension matches hidden_features so they can be added after linear_1. + self.optional_proj_1 = torch.nn.Linear(in_features, hidden_features) + self.optional_proj_2 = torch.nn.Linear(in_features, hidden_features) + + def forward(self, x: torch.Tensor, optional_input: torch.Tensor | None = None) -> torch.Tensor: + x = self.linear_1(x) + x = self.activation(x) + if optional_input is not None: + # Add optional projections after linear_1 so dimensions match (both hidden_features) + x = x + self.optional_proj_1(optional_input) + x = x + self.optional_proj_2(optional_input) + for block in self.blocks: + x = block(x) + x = self.linear_2(x) + return x + + +class ConditionalModuleGroupOffloadTests(GroupOffloadTests): + """Tests for conditionally-executed modules under group offloading with streams. + + Regression tests for the case where a module is not executed during the first forward pass + (when the lazy prefetch execution order is traced), but IS executed on subsequent passes. + Without the fix, the weights of such modules remain on CPU while the input is on GPU, + causing a RuntimeError about tensor device mismatch. + """ + + def get_model(self): + torch.manual_seed(0) + return DummyModelWithConditionalModules( + in_features=self.in_features, + hidden_features=self.hidden_features, + out_features=self.out_features, + num_layers=self.num_layers, + ) + + @parameterized.expand([("leaf_level",), ("block_level",)]) + @unittest.skipIf( + torch.device(torch_device).type not in ["cuda", "xpu"], + "Test requires a CUDA or XPU device.", + ) + def test_conditional_modules_with_stream(self, offload_type: str): + """Regression test: conditionally-executed modules must not cause device mismatch when using streams. + + The model contains two optional Linear layers (optional_proj_1, optional_proj_2) that are only + executed when `optional_input` is provided. This simulates modules like patch_short/patch_mid/ + patch_long in HeliosTransformer3DModel, which are only called when history latents are present. + + When using streams, `LazyPrefetchGroupOffloadingHook` traces the execution order on the first + forward pass and sets up a prefetch chain so each module pre-loads the next one's weights. + Modules not executed during this tracing pass are excluded from the prefetch chain. + + The bug: if a module was absent from the first (tracing) pass, its `onload_self` flag gets set + to False (meaning "someone else will onload me"). But since it's not in the prefetch chain, + nobody ever does — so its weights remain on CPU. When the module is eventually called in a + subsequent pass, the input is on GPU but the weights are on CPU, causing a RuntimeError. + + We therefore must invoke the model multiple times: + 1. First pass WITHOUT optional_input: triggers the lazy prefetch tracing. optional_proj_1/2 + are absent, so they are excluded from the prefetch chain. + 2. Second pass WITH optional_input: the regression case. Without the fix, this raises a + RuntimeError because optional_proj_1/2 weights are still on CPU. + 3. Third pass WITHOUT optional_input: verifies the model remains stable after having seen + both code paths. + """ + + model = self.get_model() + model_ref = self.get_model() + model_ref.load_state_dict(model.state_dict(), strict=True) + model_ref.to(torch_device) + + model.enable_group_offload( + torch_device, + offload_type=offload_type, + num_blocks_per_group=1, + use_stream=True, + ) + + x = torch.randn(4, self.in_features).to(torch_device) + optional_input = torch.randn(4, self.in_features).to(torch_device) + + with torch.no_grad(): + # First forward pass WITHOUT optional_input — this is when the lazy prefetch + # execution order is traced. optional_proj_1/2 are NOT in the traced order. + out_ref_no_opt = model_ref(x, optional_input=None) + out_no_opt = model(x, optional_input=None) + self.assertTrue( + torch.allclose(out_ref_no_opt, out_no_opt, atol=1e-5), + f"[{offload_type}] Outputs do not match on first pass (no optional_input).", + ) + + # Second forward pass WITH optional_input — optional_proj_1/2 ARE now called. + out_ref_with_opt = model_ref(x, optional_input=optional_input) + out_with_opt = model(x, optional_input=optional_input) + self.assertTrue( + torch.allclose(out_ref_with_opt, out_with_opt, atol=1e-5), + f"[{offload_type}] Outputs do not match on second pass (with optional_input).", + ) + + # Third pass again without optional_input — verify stable behavior. + out_ref_no_opt2 = model_ref(x, optional_input=None) + out_no_opt2 = model(x, optional_input=None) + self.assertTrue( + torch.allclose(out_ref_no_opt2, out_no_opt2, atol=1e-5), + f"[{offload_type}] Outputs do not match on third pass (back to no optional_input).", + ) diff --git a/tests/hooks/test_mag_cache.py b/tests/hooks/test_mag_cache.py new file mode 100644 index 000000000000..a7e1b52d3b69 --- /dev/null +++ b/tests/hooks/test_mag_cache.py @@ -0,0 +1,244 @@ +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import numpy as np +import torch + +from diffusers import MagCacheConfig, apply_mag_cache +from diffusers.hooks._helpers import TransformerBlockMetadata, TransformerBlockRegistry +from diffusers.models import ModelMixin +from diffusers.utils import logging + + +logger = logging.get_logger(__name__) + + +class DummyBlock(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, hidden_states, encoder_hidden_states=None, **kwargs): + # Output is double input + # This ensures Residual = 2*Input - Input = Input + return hidden_states * 2.0 + + +class DummyTransformer(ModelMixin): + def __init__(self): + super().__init__() + self.transformer_blocks = torch.nn.ModuleList([DummyBlock(), DummyBlock()]) + + def forward(self, hidden_states, encoder_hidden_states=None): + for block in self.transformer_blocks: + hidden_states = block(hidden_states, encoder_hidden_states=encoder_hidden_states) + return hidden_states + + +class TupleOutputBlock(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, hidden_states, encoder_hidden_states=None, **kwargs): + # Returns a tuple + return hidden_states * 2.0, encoder_hidden_states + + +class TupleTransformer(ModelMixin): + def __init__(self): + super().__init__() + self.transformer_blocks = torch.nn.ModuleList([TupleOutputBlock()]) + + def forward(self, hidden_states, encoder_hidden_states=None): + for block in self.transformer_blocks: + # Emulate Flux-like behavior + output = block(hidden_states, encoder_hidden_states=encoder_hidden_states) + hidden_states = output[0] + encoder_hidden_states = output[1] + return hidden_states, encoder_hidden_states + + +class MagCacheTests(unittest.TestCase): + def setUp(self): + # Register standard dummy block + TransformerBlockRegistry.register( + DummyBlock, + TransformerBlockMetadata(return_hidden_states_index=None, return_encoder_hidden_states_index=None), + ) + # Register tuple block (Flux style) + TransformerBlockRegistry.register( + TupleOutputBlock, + TransformerBlockMetadata(return_hidden_states_index=0, return_encoder_hidden_states_index=1), + ) + + def _set_context(self, model, context_name): + """Helper to set context on all hooks in the model.""" + for module in model.modules(): + if hasattr(module, "_diffusers_hook"): + module._diffusers_hook._set_context(context_name) + + def _get_calibration_data(self, model): + for module in model.modules(): + if hasattr(module, "_diffusers_hook"): + hook = module._diffusers_hook.get_hook("mag_cache_block_hook") + if hook: + return hook.state_manager.get_state().calibration_ratios + return [] + + def test_mag_cache_validation(self): + """Test that missing mag_ratios raises ValueError.""" + with self.assertRaises(ValueError): + MagCacheConfig(num_inference_steps=10, calibrate=False) + + def test_mag_cache_skipping_logic(self): + """ + Tests that MagCache correctly calculates residuals and skips blocks when conditions are met. + """ + model = DummyTransformer() + + # Dummy ratios: [1.0, 1.0] implies 0 accumulated error if we skip + ratios = np.array([1.0, 1.0]) + + config = MagCacheConfig( + threshold=100.0, + num_inference_steps=2, + retention_ratio=0.0, # Enable immediate skipping + max_skip_steps=5, + mag_ratios=ratios, + ) + + apply_mag_cache(model, config) + self._set_context(model, "test_context") + + # Step 0: Input 10.0 -> Output 40.0 (2 blocks * 2x each) + # HeadInput=10. Output=40. Residual=30. + input_t0 = torch.tensor([[[10.0]]]) + output_t0 = model(input_t0) + self.assertTrue(torch.allclose(output_t0, torch.tensor([[[40.0]]])), "Step 0 failed") + + # Step 1: Input 11.0. + # If Skipped: Output = Input(11) + Residual(30) = 41.0 + # If Computed: Output = 11 * 4 = 44.0 + input_t1 = torch.tensor([[[11.0]]]) + output_t1 = model(input_t1) + + self.assertTrue( + torch.allclose(output_t1, torch.tensor([[[41.0]]])), f"Expected Skip (41.0), got {output_t1.item()}" + ) + + def test_mag_cache_retention(self): + """Test that retention_ratio prevents skipping even if error is low.""" + model = DummyTransformer() + # Ratios that imply 0 error, so it *would* skip if retention allowed it + ratios = np.array([1.0, 1.0]) + + config = MagCacheConfig( + threshold=100.0, + num_inference_steps=2, + retention_ratio=1.0, # Force retention for ALL steps + mag_ratios=ratios, + ) + + apply_mag_cache(model, config) + self._set_context(model, "test_context") + + # Step 0 + model(torch.tensor([[[10.0]]])) + + # Step 1: Should COMPUTE (44.0) not SKIP (41.0) because of retention + input_t1 = torch.tensor([[[11.0]]]) + output_t1 = model(input_t1) + + self.assertTrue( + torch.allclose(output_t1, torch.tensor([[[44.0]]])), + f"Expected Compute (44.0) due to retention, got {output_t1.item()}", + ) + + def test_mag_cache_tuple_outputs(self): + """Test compatibility with models returning (hidden, encoder_hidden) like Flux.""" + model = TupleTransformer() + ratios = np.array([1.0, 1.0]) + + config = MagCacheConfig(threshold=100.0, num_inference_steps=2, retention_ratio=0.0, mag_ratios=ratios) + + apply_mag_cache(model, config) + self._set_context(model, "test_context") + + # Step 0: Compute. Input 10.0 -> Output 20.0 (1 block * 2x) + # Residual = 10.0 + input_t0 = torch.tensor([[[10.0]]]) + enc_t0 = torch.tensor([[[1.0]]]) + out_0, _ = model(input_t0, encoder_hidden_states=enc_t0) + self.assertTrue(torch.allclose(out_0, torch.tensor([[[20.0]]]))) + + # Step 1: Skip. Input 11.0. + # Skipped Output = 11 + 10 = 21.0 + input_t1 = torch.tensor([[[11.0]]]) + out_1, _ = model(input_t1, encoder_hidden_states=enc_t0) + + self.assertTrue( + torch.allclose(out_1, torch.tensor([[[21.0]]])), f"Tuple skip failed. Expected 21.0, got {out_1.item()}" + ) + + def test_mag_cache_reset(self): + """Test that state resets correctly after num_inference_steps.""" + model = DummyTransformer() + config = MagCacheConfig( + threshold=100.0, num_inference_steps=2, retention_ratio=0.0, mag_ratios=np.array([1.0, 1.0]) + ) + apply_mag_cache(model, config) + self._set_context(model, "test_context") + + input_t = torch.ones(1, 1, 1) + + model(input_t) # Step 0 + model(input_t) # Step 1 (Skipped) + + # Step 2 (Reset -> Step 0) -> Should Compute + # Input 2.0 -> Output 8.0 + input_t2 = torch.tensor([[[2.0]]]) + output_t2 = model(input_t2) + + self.assertTrue(torch.allclose(output_t2, torch.tensor([[[8.0]]])), "State did not reset correctly") + + def test_mag_cache_calibration(self): + """Test that calibration mode records ratios.""" + model = DummyTransformer() + config = MagCacheConfig(num_inference_steps=2, calibrate=True) + apply_mag_cache(model, config) + self._set_context(model, "test_context") + + # Step 0 + # HeadInput = 10. Output = 40. Residual = 30. + # Ratio 0 is placeholder 1.0 + model(torch.tensor([[[10.0]]])) + + # Check intermediate state + ratios = self._get_calibration_data(model) + self.assertEqual(len(ratios), 1) + self.assertEqual(ratios[0], 1.0) + + # Step 1 + # HeadInput = 10. Output = 40. Residual = 30. + # PrevResidual = 30. CurrResidual = 30. + # Ratio = 30/30 = 1.0 + model(torch.tensor([[[10.0]]])) + + # Verify it computes fully (no skip) + # If it skipped, output would be 41.0. It should be 40.0 + # Actually in test setup, input is same (10.0) so output 40.0. + # Let's ensure list is empty after reset (end of step 1) + ratios_after = self._get_calibration_data(model) + self.assertEqual(ratios_after, []) diff --git a/tests/lora/test_lora_layers_auraflow.py b/tests/lora/test_lora_layers_auraflow.py index 91f63c4b56c4..78ef4ce151be 100644 --- a/tests/lora/test_lora_layers_auraflow.py +++ b/tests/lora/test_lora_layers_auraflow.py @@ -76,6 +76,8 @@ class AuraFlowLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): text_encoder_target_modules = ["q", "k", "v", "o"] denoiser_target_modules = ["to_q", "to_k", "to_v", "to_out.0", "linear_1"] + supports_text_encoder_loras = False + @property def output_shape(self): return (1, 8, 8, 3) @@ -114,23 +116,3 @@ def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(se @unittest.skip("Not supported in AuraFlow.") def test_modify_padding_mode(self): pass - - @unittest.skip("Text encoder LoRA is not supported in AuraFlow.") - def test_simple_inference_with_partial_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in AuraFlow.") - def test_simple_inference_with_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in AuraFlow.") - def test_simple_inference_with_text_lora_and_scale(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in AuraFlow.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in AuraFlow.") - def test_simple_inference_with_text_lora_save_load(self): - pass diff --git a/tests/lora/test_lora_layers_cogvideox.py b/tests/lora/test_lora_layers_cogvideox.py index fa57b4c9c2f9..7bd54b77ca35 100644 --- a/tests/lora/test_lora_layers_cogvideox.py +++ b/tests/lora/test_lora_layers_cogvideox.py @@ -87,6 +87,8 @@ class CogVideoXLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): text_encoder_target_modules = ["q", "k", "v", "o"] + supports_text_encoder_loras = False + @property def output_shape(self): return (1, 9, 16, 16, 3) @@ -147,26 +149,6 @@ def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(se def test_modify_padding_mode(self): pass - @unittest.skip("Text encoder LoRA is not supported in CogVideoX.") - def test_simple_inference_with_partial_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in CogVideoX.") - def test_simple_inference_with_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in CogVideoX.") - def test_simple_inference_with_text_lora_and_scale(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in CogVideoX.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in CogVideoX.") - def test_simple_inference_with_text_lora_save_load(self): - pass - @unittest.skip("Not supported in CogVideoX.") def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): pass diff --git a/tests/lora/test_lora_layers_cogview4.py b/tests/lora/test_lora_layers_cogview4.py index 30eb8fbb6367..e8ee6e7a7db6 100644 --- a/tests/lora/test_lora_layers_cogview4.py +++ b/tests/lora/test_lora_layers_cogview4.py @@ -85,6 +85,8 @@ class CogView4LoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): "text_encoder", ) + supports_text_encoder_loras = False + @property def output_shape(self): return (1, 32, 32, 3) @@ -162,23 +164,3 @@ def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(se @unittest.skip("Not supported in CogView4.") def test_modify_padding_mode(self): pass - - @unittest.skip("Text encoder LoRA is not supported in CogView4.") - def test_simple_inference_with_partial_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in CogView4.") - def test_simple_inference_with_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in CogView4.") - def test_simple_inference_with_text_lora_and_scale(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in CogView4.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in CogView4.") - def test_simple_inference_with_text_lora_save_load(self): - pass diff --git a/tests/lora/test_lora_layers_flux2.py b/tests/lora/test_lora_layers_flux2.py index 4ae189aceb66..d970b7d7847f 100644 --- a/tests/lora/test_lora_layers_flux2.py +++ b/tests/lora/test_lora_layers_flux2.py @@ -66,6 +66,8 @@ class Flux2LoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): text_encoder_cls, text_encoder_id = Mistral3ForConditionalGeneration, "hf-internal-testing/tiny-mistral3-diffusers" denoiser_target_modules = ["to_qkv_mlp_proj", "to_k"] + supports_text_encoder_loras = False + @property def output_shape(self): return (1, 8, 8, 3) @@ -146,23 +148,3 @@ def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(se @unittest.skip("Not supported in Flux2.") def test_modify_padding_mode(self): pass - - @unittest.skip("Text encoder LoRA is not supported in Flux2.") - def test_simple_inference_with_partial_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Flux2.") - def test_simple_inference_with_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Flux2.") - def test_simple_inference_with_text_lora_and_scale(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Flux2.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Flux2.") - def test_simple_inference_with_text_lora_save_load(self): - pass diff --git a/tests/lora/test_lora_layers_helios.py b/tests/lora/test_lora_layers_helios.py new file mode 100644 index 000000000000..fbcc3b808eee --- /dev/null +++ b/tests/lora/test_lora_layers_helios.py @@ -0,0 +1,120 @@ +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import sys +import unittest + +import torch +from transformers import AutoTokenizer, T5EncoderModel + +from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler, HeliosPipeline, HeliosTransformer3DModel + +from ..testing_utils import floats_tensor, require_peft_backend, skip_mps + + +sys.path.append(".") + +from .utils import PeftLoraLoaderMixinTests # noqa: E402 + + +@require_peft_backend +@skip_mps +class HeliosLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): + pipeline_class = HeliosPipeline + scheduler_cls = FlowMatchEulerDiscreteScheduler + scheduler_kwargs = {} + + transformer_kwargs = { + "patch_size": (1, 2, 2), + "num_attention_heads": 2, + "attention_head_dim": 12, + "in_channels": 16, + "out_channels": 16, + "text_dim": 32, + "freq_dim": 256, + "ffn_dim": 32, + "num_layers": 2, + "cross_attn_norm": True, + "qk_norm": "rms_norm_across_heads", + "rope_dim": (4, 4, 4), + "has_multi_term_memory_patch": True, + "guidance_cross_attn": True, + "zero_history_timestep": True, + "is_amplify_history": False, + } + transformer_cls = HeliosTransformer3DModel + vae_kwargs = { + "base_dim": 3, + "z_dim": 16, + "dim_mult": [1, 1, 1, 1], + "num_res_blocks": 1, + "temperal_downsample": [False, True, True], + } + vae_cls = AutoencoderKLWan + has_two_text_encoders = True + tokenizer_cls, tokenizer_id = AutoTokenizer, "hf-internal-testing/tiny-random-t5" + text_encoder_cls, text_encoder_id = T5EncoderModel, "hf-internal-testing/tiny-random-t5" + + text_encoder_target_modules = ["q", "k", "v", "o"] + + supports_text_encoder_loras = False + + @property + def output_shape(self): + return (1, 33, 32, 32, 3) + + def get_dummy_inputs(self, with_generator=True): + batch_size = 1 + sequence_length = 16 + num_channels = 4 + num_frames = 9 + num_latent_frames = 3 # (num_frames - 1) // temporal_compression_ratio + 1 + sizes = (4, 4) + + generator = torch.manual_seed(0) + noise = floats_tensor((batch_size, num_latent_frames, num_channels) + sizes) + input_ids = torch.randint(1, sequence_length, size=(batch_size, sequence_length), generator=generator) + + pipeline_inputs = { + "prompt": "", + "num_frames": num_frames, + "num_inference_steps": 1, + "guidance_scale": 6.0, + "height": 32, + "width": 32, + "max_sequence_length": sequence_length, + "output_type": "np", + } + if with_generator: + pipeline_inputs.update({"generator": generator}) + + return noise, input_ids, pipeline_inputs + + def test_simple_inference_with_text_lora_denoiser_fused_multi(self): + super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3) + + def test_simple_inference_with_text_denoiser_lora_unfused(self): + super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3) + + @unittest.skip("Not supported in Helios.") + def test_simple_inference_with_text_denoiser_block_scale(self): + pass + + @unittest.skip("Not supported in Helios.") + def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): + pass + + @unittest.skip("Not supported in Helios.") + def test_modify_padding_mode(self): + pass diff --git a/tests/lora/test_lora_layers_hunyuanvideo.py b/tests/lora/test_lora_layers_hunyuanvideo.py index cfd5d3146a91..e59bc5662fe1 100644 --- a/tests/lora/test_lora_layers_hunyuanvideo.py +++ b/tests/lora/test_lora_layers_hunyuanvideo.py @@ -117,6 +117,8 @@ class HunyuanVideoLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): "text_encoder_2", ) + supports_text_encoder_loras = False + @property def output_shape(self): return (1, 9, 32, 32, 3) @@ -172,26 +174,6 @@ def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(se def test_modify_padding_mode(self): pass - @unittest.skip("Text encoder LoRA is not supported in HunyuanVideo.") - def test_simple_inference_with_partial_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in HunyuanVideo.") - def test_simple_inference_with_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in HunyuanVideo.") - def test_simple_inference_with_text_lora_and_scale(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in HunyuanVideo.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in HunyuanVideo.") - def test_simple_inference_with_text_lora_save_load(self): - pass - @nightly @require_torch_accelerator diff --git a/tests/lora/test_lora_layers_ltx2.py b/tests/lora/test_lora_layers_ltx2.py new file mode 100644 index 000000000000..0a4b14454f5b --- /dev/null +++ b/tests/lora/test_lora_layers_ltx2.py @@ -0,0 +1,271 @@ +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import sys +import unittest + +import torch +from transformers import AutoTokenizer, Gemma3ForConditionalGeneration + +from diffusers import ( + AutoencoderKLLTX2Audio, + AutoencoderKLLTX2Video, + FlowMatchEulerDiscreteScheduler, + LTX2Pipeline, + LTX2VideoTransformer3DModel, +) +from diffusers.pipelines.ltx2 import LTX2TextConnectors +from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder +from diffusers.utils.import_utils import is_peft_available + +from ..testing_utils import floats_tensor, require_peft_backend + + +if is_peft_available(): + from peft import LoraConfig + + +sys.path.append(".") + +from .utils import PeftLoraLoaderMixinTests # noqa: E402 + + +@require_peft_backend +class LTX2LoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): + pipeline_class = LTX2Pipeline + scheduler_cls = FlowMatchEulerDiscreteScheduler + scheduler_kwargs = {} + + transformer_kwargs = { + "in_channels": 4, + "out_channels": 4, + "patch_size": 1, + "patch_size_t": 1, + "num_attention_heads": 2, + "attention_head_dim": 8, + "cross_attention_dim": 16, + "audio_in_channels": 4, + "audio_out_channels": 4, + "audio_num_attention_heads": 2, + "audio_attention_head_dim": 4, + "audio_cross_attention_dim": 8, + "num_layers": 1, + "qk_norm": "rms_norm_across_heads", + "caption_channels": 32, + "rope_double_precision": False, + "rope_type": "split", + } + transformer_cls = LTX2VideoTransformer3DModel + + vae_kwargs = { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 4, + "block_out_channels": (8,), + "decoder_block_out_channels": (8,), + "layers_per_block": (1,), + "decoder_layers_per_block": (1, 1), + "spatio_temporal_scaling": (True,), + "decoder_spatio_temporal_scaling": (True,), + "decoder_inject_noise": (False, False), + "downsample_type": ("spatial",), + "upsample_residual": (False,), + "upsample_factor": (1,), + "timestep_conditioning": False, + "patch_size": 1, + "patch_size_t": 1, + "encoder_causal": True, + "decoder_causal": False, + } + vae_cls = AutoencoderKLLTX2Video + + audio_vae_kwargs = { + "base_channels": 4, + "output_channels": 2, + "ch_mult": (1,), + "num_res_blocks": 1, + "attn_resolutions": None, + "in_channels": 2, + "resolution": 32, + "latent_channels": 2, + "norm_type": "pixel", + "causality_axis": "height", + "dropout": 0.0, + "mid_block_add_attention": False, + "sample_rate": 16000, + "mel_hop_length": 160, + "is_causal": True, + "mel_bins": 8, + } + audio_vae_cls = AutoencoderKLLTX2Audio + + vocoder_kwargs = { + "in_channels": 16, # output_channels * mel_bins = 2 * 8 + "hidden_channels": 32, + "out_channels": 2, + "upsample_kernel_sizes": [4, 4], + "upsample_factors": [2, 2], + "resnet_kernel_sizes": [3], + "resnet_dilations": [[1, 3, 5]], + "leaky_relu_negative_slope": 0.1, + "output_sampling_rate": 16000, + } + vocoder_cls = LTX2Vocoder + + connectors_kwargs = { + "caption_channels": 32, # Will be set dynamically from text_encoder + "text_proj_in_factor": 2, # Will be set dynamically from text_encoder + "video_connector_num_attention_heads": 4, + "video_connector_attention_head_dim": 8, + "video_connector_num_layers": 1, + "video_connector_num_learnable_registers": None, + "audio_connector_num_attention_heads": 4, + "audio_connector_attention_head_dim": 8, + "audio_connector_num_layers": 1, + "audio_connector_num_learnable_registers": None, + "connector_rope_base_seq_len": 32, + "rope_theta": 10000.0, + "rope_double_precision": False, + "causal_temporal_positioning": False, + "rope_type": "split", + } + connectors_cls = LTX2TextConnectors + + tokenizer_cls, tokenizer_id = AutoTokenizer, "hf-internal-testing/tiny-gemma3" + text_encoder_cls, text_encoder_id = ( + Gemma3ForConditionalGeneration, + "hf-internal-testing/tiny-gemma3", + ) + + denoiser_target_modules = ["to_q", "to_k", "to_out.0"] + + supports_text_encoder_loras = False + + @property + def output_shape(self): + return (1, 5, 32, 32, 3) + + def get_dummy_inputs(self, with_generator=True): + batch_size = 1 + sequence_length = 16 + num_channels = 4 + num_frames = 5 + num_latent_frames = 2 + latent_height = 8 + latent_width = 8 + + generator = torch.manual_seed(0) + noise = floats_tensor((batch_size, num_latent_frames, num_channels, latent_height, latent_width)) + input_ids = torch.randint(1, sequence_length, size=(batch_size, sequence_length), generator=generator) + + pipeline_inputs = { + "prompt": "a robot dancing", + "num_frames": num_frames, + "num_inference_steps": 2, + "guidance_scale": 1.0, + "height": 32, + "width": 32, + "frame_rate": 25.0, + "max_sequence_length": sequence_length, + "output_type": "np", + } + if with_generator: + pipeline_inputs.update({"generator": generator}) + + return noise, input_ids, pipeline_inputs + + def get_dummy_components(self, scheduler_cls=None, use_dora=False, lora_alpha=None): + # Override to instantiate LTX2-specific components (connectors, audio_vae, vocoder) + torch.manual_seed(0) + text_encoder = self.text_encoder_cls.from_pretrained(self.text_encoder_id) + tokenizer = self.tokenizer_cls.from_pretrained(self.tokenizer_id) + + # Update caption_channels and text_proj_in_factor based on text_encoder config + transformer_kwargs = self.transformer_kwargs.copy() + transformer_kwargs["caption_channels"] = text_encoder.config.text_config.hidden_size + + connectors_kwargs = self.connectors_kwargs.copy() + connectors_kwargs["caption_channels"] = text_encoder.config.text_config.hidden_size + connectors_kwargs["text_proj_in_factor"] = text_encoder.config.text_config.num_hidden_layers + 1 + + torch.manual_seed(0) + transformer = self.transformer_cls(**transformer_kwargs) + + torch.manual_seed(0) + vae = self.vae_cls(**self.vae_kwargs) + vae.use_framewise_encoding = False + vae.use_framewise_decoding = False + + torch.manual_seed(0) + audio_vae = self.audio_vae_cls(**self.audio_vae_kwargs) + + torch.manual_seed(0) + vocoder = self.vocoder_cls(**self.vocoder_kwargs) + + torch.manual_seed(0) + connectors = self.connectors_cls(**connectors_kwargs) + + if scheduler_cls is None: + scheduler_cls = self.scheduler_cls + scheduler = scheduler_cls(**self.scheduler_kwargs) + + rank = 4 + lora_alpha = rank if lora_alpha is None else lora_alpha + + text_lora_config = LoraConfig( + r=rank, + lora_alpha=lora_alpha, + target_modules=self.text_encoder_target_modules, + init_lora_weights=False, + use_dora=use_dora, + ) + + denoiser_lora_config = LoraConfig( + r=rank, + lora_alpha=lora_alpha, + target_modules=["to_q", "to_k", "to_v", "to_out.0"], + init_lora_weights=False, + use_dora=use_dora, + ) + + pipeline_components = { + "transformer": transformer, + "vae": vae, + "audio_vae": audio_vae, + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + "connectors": connectors, + "vocoder": vocoder, + } + + return pipeline_components, text_lora_config, denoiser_lora_config + + def test_simple_inference_with_text_lora_denoiser_fused_multi(self): + super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3) + + def test_simple_inference_with_text_denoiser_lora_unfused(self): + super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3) + + @unittest.skip("Not supported in LTX2.") + def test_simple_inference_with_text_denoiser_block_scale(self): + pass + + @unittest.skip("Not supported in LTX2.") + def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): + pass + + @unittest.skip("Not supported in LTX2.") + def test_modify_padding_mode(self): + pass diff --git a/tests/lora/test_lora_layers_ltx_video.py b/tests/lora/test_lora_layers_ltx_video.py index 6ab51a5e513f..095e5b577cf0 100644 --- a/tests/lora/test_lora_layers_ltx_video.py +++ b/tests/lora/test_lora_layers_ltx_video.py @@ -76,6 +76,8 @@ class LTXVideoLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): text_encoder_target_modules = ["q", "k", "v", "o"] + supports_text_encoder_loras = False + @property def output_shape(self): return (1, 9, 32, 32, 3) @@ -125,23 +127,3 @@ def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(se @unittest.skip("Not supported in LTXVideo.") def test_modify_padding_mode(self): pass - - @unittest.skip("Text encoder LoRA is not supported in LTXVideo.") - def test_simple_inference_with_partial_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in LTXVideo.") - def test_simple_inference_with_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in LTXVideo.") - def test_simple_inference_with_text_lora_and_scale(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in LTXVideo.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in LTXVideo.") - def test_simple_inference_with_text_lora_save_load(self): - pass diff --git a/tests/lora/test_lora_layers_lumina2.py b/tests/lora/test_lora_layers_lumina2.py index 0417b05b33a1..da032229a785 100644 --- a/tests/lora/test_lora_layers_lumina2.py +++ b/tests/lora/test_lora_layers_lumina2.py @@ -74,6 +74,8 @@ class Lumina2LoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): tokenizer_cls, tokenizer_id = AutoTokenizer, "hf-internal-testing/dummy-gemma" text_encoder_cls, text_encoder_id = GemmaForCausalLM, "hf-internal-testing/dummy-gemma-diffusers" + supports_text_encoder_loras = False + @property def output_shape(self): return (1, 4, 4, 3) @@ -113,26 +115,6 @@ def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(se def test_modify_padding_mode(self): pass - @unittest.skip("Text encoder LoRA is not supported in Lumina2.") - def test_simple_inference_with_partial_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Lumina2.") - def test_simple_inference_with_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Lumina2.") - def test_simple_inference_with_text_lora_and_scale(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Lumina2.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Lumina2.") - def test_simple_inference_with_text_lora_save_load(self): - pass - @skip_mps @pytest.mark.xfail( condition=torch.device(torch_device).type == "cpu" and is_torch_version(">=", "2.5"), diff --git a/tests/lora/test_lora_layers_mochi.py b/tests/lora/test_lora_layers_mochi.py index 7be81273db77..ee8254112924 100644 --- a/tests/lora/test_lora_layers_mochi.py +++ b/tests/lora/test_lora_layers_mochi.py @@ -67,6 +67,8 @@ class MochiLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): text_encoder_target_modules = ["q", "k", "v", "o"] + supports_text_encoder_loras = False + @property def output_shape(self): return (1, 7, 16, 16, 3) @@ -117,26 +119,6 @@ def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(se def test_modify_padding_mode(self): pass - @unittest.skip("Text encoder LoRA is not supported in Mochi.") - def test_simple_inference_with_partial_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Mochi.") - def test_simple_inference_with_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Mochi.") - def test_simple_inference_with_text_lora_and_scale(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Mochi.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Mochi.") - def test_simple_inference_with_text_lora_save_load(self): - pass - @unittest.skip("Not supported in CogVideoX.") def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): pass diff --git a/tests/lora/test_lora_layers_qwenimage.py b/tests/lora/test_lora_layers_qwenimage.py index 51de2f8e20e1..73fd026a670c 100644 --- a/tests/lora/test_lora_layers_qwenimage.py +++ b/tests/lora/test_lora_layers_qwenimage.py @@ -69,6 +69,8 @@ class QwenImageLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): ) denoiser_target_modules = ["to_q", "to_k", "to_v", "to_out.0"] + supports_text_encoder_loras = False + @property def output_shape(self): return (1, 8, 8, 3) @@ -107,23 +109,3 @@ def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(se @unittest.skip("Not supported in Qwen Image.") def test_modify_padding_mode(self): pass - - @unittest.skip("Text encoder LoRA is not supported in Qwen Image.") - def test_simple_inference_with_partial_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Qwen Image.") - def test_simple_inference_with_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Qwen Image.") - def test_simple_inference_with_text_lora_and_scale(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Qwen Image.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Qwen Image.") - def test_simple_inference_with_text_lora_save_load(self): - pass diff --git a/tests/lora/test_lora_layers_sana.py b/tests/lora/test_lora_layers_sana.py index a860b7b44f2c..97bf5cbba920 100644 --- a/tests/lora/test_lora_layers_sana.py +++ b/tests/lora/test_lora_layers_sana.py @@ -75,6 +75,8 @@ class SanaLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): tokenizer_cls, tokenizer_id = GemmaTokenizer, "hf-internal-testing/dummy-gemma" text_encoder_cls, text_encoder_id = Gemma2Model, "hf-internal-testing/dummy-gemma-for-diffusers" + supports_text_encoder_loras = False + @property def output_shape(self): return (1, 32, 32, 3) @@ -117,26 +119,6 @@ def test_simple_inference_with_text_denoiser_block_scale(self): def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): pass - @unittest.skip("Text encoder LoRA is not supported in SANA.") - def test_simple_inference_with_partial_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in SANA.") - def test_simple_inference_with_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in SANA.") - def test_simple_inference_with_text_lora_and_scale(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in SANA.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in SANA.") - def test_simple_inference_with_text_lora_save_load(self): - pass - @unittest.skipIf(IS_GITHUB_ACTIONS, reason="Skipping test inside GitHub Actions environment") def test_layerwise_casting_inference_denoiser(self): return super().test_layerwise_casting_inference_denoiser() diff --git a/tests/lora/test_lora_layers_wan.py b/tests/lora/test_lora_layers_wan.py index 5734509b410f..1a5b47ce4e5e 100644 --- a/tests/lora/test_lora_layers_wan.py +++ b/tests/lora/test_lora_layers_wan.py @@ -18,18 +18,9 @@ import torch from transformers import AutoTokenizer, T5EncoderModel -from diffusers import ( - AutoencoderKLWan, - FlowMatchEulerDiscreteScheduler, - WanPipeline, - WanTransformer3DModel, -) +from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler, WanPipeline, WanTransformer3DModel -from ..testing_utils import ( - floats_tensor, - require_peft_backend, - skip_mps, -) +from ..testing_utils import floats_tensor, require_peft_backend, skip_mps sys.path.append(".") @@ -73,6 +64,8 @@ class WanLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): text_encoder_target_modules = ["q", "k", "v", "o"] + supports_text_encoder_loras = False + @property def output_shape(self): return (1, 9, 32, 32, 3) @@ -121,23 +114,3 @@ def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(se @unittest.skip("Not supported in Wan.") def test_modify_padding_mode(self): pass - - @unittest.skip("Text encoder LoRA is not supported in Wan.") - def test_simple_inference_with_partial_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Wan.") - def test_simple_inference_with_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Wan.") - def test_simple_inference_with_text_lora_and_scale(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Wan.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Wan.") - def test_simple_inference_with_text_lora_save_load(self): - pass diff --git a/tests/lora/test_lora_layers_wanvace.py b/tests/lora/test_lora_layers_wanvace.py index ab1f57bfc9da..c8acaea9bef0 100644 --- a/tests/lora/test_lora_layers_wanvace.py +++ b/tests/lora/test_lora_layers_wanvace.py @@ -85,6 +85,8 @@ class WanVACELoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): text_encoder_target_modules = ["q", "k", "v", "o"] + supports_text_encoder_loras = False + @property def output_shape(self): return (1, 9, 16, 16, 3) @@ -139,26 +141,6 @@ def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(se def test_modify_padding_mode(self): pass - @unittest.skip("Text encoder LoRA is not supported in Wan VACE.") - def test_simple_inference_with_partial_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Wan VACE.") - def test_simple_inference_with_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Wan VACE.") - def test_simple_inference_with_text_lora_and_scale(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Wan VACE.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Wan VACE.") - def test_simple_inference_with_text_lora_save_load(self): - pass - def test_layerwise_casting_inference_denoiser(self): super().test_layerwise_casting_inference_denoiser() diff --git a/tests/lora/test_lora_layers_z_image.py b/tests/lora/test_lora_layers_z_image.py index 35d1389d9612..8432ea56a6fa 100644 --- a/tests/lora/test_lora_layers_z_image.py +++ b/tests/lora/test_lora_layers_z_image.py @@ -75,6 +75,8 @@ class ZImageLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): text_encoder_cls, text_encoder_id = Qwen3Model, None # Will be created inline denoiser_target_modules = ["to_q", "to_k", "to_v", "to_out.0"] + supports_text_encoder_loras = False + @property def output_shape(self): return (1, 32, 32, 3) @@ -263,23 +265,3 @@ def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(se @unittest.skip("Not supported in ZImage.") def test_modify_padding_mode(self): pass - - @unittest.skip("Text encoder LoRA is not supported in ZImage.") - def test_simple_inference_with_partial_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in ZImage.") - def test_simple_inference_with_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in ZImage.") - def test_simple_inference_with_text_lora_and_scale(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in ZImage.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in ZImage.") - def test_simple_inference_with_text_lora_save_load(self): - pass diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 5fae6cac0a7f..efa49b9f4838 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -117,6 +117,7 @@ class PeftLoraLoaderMixinTests: tokenizer_cls, tokenizer_id, tokenizer_subfolder = None, None, "" tokenizer_2_cls, tokenizer_2_id, tokenizer_2_subfolder = None, None, "" tokenizer_3_cls, tokenizer_3_id, tokenizer_3_subfolder = None, None, "" + supports_text_encoder_loras = True unet_kwargs = None transformer_cls = None @@ -333,6 +334,9 @@ def test_simple_inference_with_text_lora(self): Tests a simple inference with lora attached on the text encoder and makes sure it works as expected """ + if not self.supports_text_encoder_loras: + pytest.skip("Skipping test as text encoder LoRAs are not currently supported.") + components, text_lora_config, _ = self.get_dummy_components() pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) @@ -457,6 +461,9 @@ def test_simple_inference_with_text_lora_and_scale(self): Tests a simple inference with lora attached on the text encoder + scale argument and makes sure it works as expected """ + if not self.supports_text_encoder_loras: + pytest.skip("Skipping test as text encoder LoRAs are not currently supported.") + attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class) components, text_lora_config, _ = self.get_dummy_components() pipe = self.pipeline_class(**components) @@ -494,6 +501,9 @@ def test_simple_inference_with_text_lora_fused(self): Tests a simple inference with lora attached into text encoder + fuses the lora weights into base model and makes sure it works as expected """ + if not self.supports_text_encoder_loras: + pytest.skip("Skipping test as text encoder LoRAs are not currently supported.") + components, text_lora_config, _ = self.get_dummy_components() pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) @@ -555,6 +565,9 @@ def test_simple_inference_with_text_lora_save_load(self): """ Tests a simple usecase where users could use saving utilities for LoRA. """ + if not self.supports_text_encoder_loras: + pytest.skip("Skipping test as text encoder LoRAs are not currently supported.") + components, text_lora_config, _ = self.get_dummy_components() pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) @@ -593,6 +606,9 @@ def test_simple_inference_with_partial_text_lora(self): with different ranks and some adapters removed and makes sure it works as expected """ + if not self.supports_text_encoder_loras: + pytest.skip("Skipping test as text encoder LoRAs are not currently supported.") + components, _, _ = self.get_dummy_components() # Verify `StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder` handles different ranks per module (PR#8324). text_lora_config = LoraConfig( @@ -651,6 +667,9 @@ def test_simple_inference_save_pretrained_with_text_lora(self): """ Tests a simple usecase where users could use saving utilities for LoRA through save_pretrained """ + if not self.supports_text_encoder_loras: + pytest.skip("Skipping test as text encoder LoRAs are not currently supported.") + components, text_lora_config, _ = self.get_dummy_components() pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_ltx2_audio.py b/tests/models/autoencoders/test_models_autoencoder_kl_ltx2_audio.py new file mode 100644 index 000000000000..ce93dfb42afe --- /dev/null +++ b/tests/models/autoencoders/test_models_autoencoder_kl_ltx2_audio.py @@ -0,0 +1,88 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +from diffusers import AutoencoderKLLTX2Audio + +from ...testing_utils import ( + floats_tensor, + torch_device, +) +from ..test_modeling_common import ModelTesterMixin +from .testing_utils import AutoencoderTesterMixin + + +class AutoencoderKLLTX2AudioTests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase): + model_class = AutoencoderKLLTX2Audio + main_input_name = "sample" + base_precision = 1e-2 + + def get_autoencoder_kl_ltx_video_config(self): + return { + "in_channels": 2, # stereo, + "output_channels": 2, + "latent_channels": 4, + "base_channels": 16, + "ch_mult": (1, 2, 4), + "resolution": 16, + "attn_resolutions": None, + "num_res_blocks": 2, + "norm_type": "pixel", + "causality_axis": "height", + "mid_block_add_attention": False, + "sample_rate": 16000, + "mel_hop_length": 160, + "mel_bins": 16, + "is_causal": True, + "double_z": True, + } + + @property + def dummy_input(self): + batch_size = 2 + num_channels = 2 + num_frames = 8 + num_mel_bins = 16 + + spectrogram = floats_tensor((batch_size, num_channels, num_frames, num_mel_bins)).to(torch_device) + + input_dict = {"sample": spectrogram} + return input_dict + + @property + def input_shape(self): + return (2, 5, 16) + + @property + def output_shape(self): + return (2, 5, 16) + + def prepare_init_args_and_inputs_for_common(self): + init_dict = self.get_autoencoder_kl_ltx_video_config() + inputs_dict = self.dummy_input + return init_dict, inputs_dict + + # Overriding as output shape is not the same as input shape for LTX 2.0 audio VAE + def test_output(self): + super().test_output(expected_output_shape=(2, 2, 5, 16)) + + @unittest.skip("Unsupported test.") + def test_outputs_equivalence(self): + pass + + @unittest.skip("AutoencoderKLLTX2Audio does not support `norm_num_groups` because it does not use GroupNorm.") + def test_forward_with_norm_groups(self): + pass diff --git a/tests/models/autoencoders/test_models_autoencoder_ltx2_video.py b/tests/models/autoencoders/test_models_autoencoder_ltx2_video.py new file mode 100644 index 000000000000..146241361a82 --- /dev/null +++ b/tests/models/autoencoders/test_models_autoencoder_ltx2_video.py @@ -0,0 +1,103 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +from diffusers import AutoencoderKLLTX2Video + +from ...testing_utils import ( + enable_full_determinism, + floats_tensor, + torch_device, +) +from ..test_modeling_common import ModelTesterMixin +from .testing_utils import AutoencoderTesterMixin + + +enable_full_determinism() + + +class AutoencoderKLLTX2VideoTests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase): + model_class = AutoencoderKLLTX2Video + main_input_name = "sample" + base_precision = 1e-2 + + def get_autoencoder_kl_ltx_video_config(self): + return { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 8, + "block_out_channels": (8, 8, 8, 8), + "decoder_block_out_channels": (16, 32, 64), + "layers_per_block": (1, 1, 1, 1, 1), + "decoder_layers_per_block": (1, 1, 1, 1), + "spatio_temporal_scaling": (True, True, True, True), + "decoder_spatio_temporal_scaling": (True, True, True), + "decoder_inject_noise": (False, False, False, False), + "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"), + "upsample_residual": (True, True, True), + "upsample_factor": (2, 2, 2), + "timestep_conditioning": False, + "patch_size": 1, + "patch_size_t": 1, + "encoder_causal": True, + "decoder_causal": False, + "encoder_spatial_padding_mode": "zeros", + # Full model uses `reflect` but this does not have deterministic backward implementation, so use `zeros` + "decoder_spatial_padding_mode": "zeros", + } + + @property + def dummy_input(self): + batch_size = 2 + num_frames = 9 + num_channels = 3 + sizes = (16, 16) + + image = floats_tensor((batch_size, num_channels, num_frames) + sizes).to(torch_device) + + input_dict = {"sample": image} + return input_dict + + @property + def input_shape(self): + return (3, 9, 16, 16) + + @property + def output_shape(self): + return (3, 9, 16, 16) + + def prepare_init_args_and_inputs_for_common(self): + init_dict = self.get_autoencoder_kl_ltx_video_config() + inputs_dict = self.dummy_input + return init_dict, inputs_dict + + def test_gradient_checkpointing_is_applied(self): + expected_set = { + "LTX2VideoEncoder3d", + "LTX2VideoDecoder3d", + "LTX2VideoDownBlock3D", + "LTX2VideoMidBlock3d", + "LTX2VideoUpBlock3d", + } + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + + @unittest.skip("Unsupported test.") + def test_outputs_equivalence(self): + pass + + @unittest.skip("AutoencoderKLLTXVideo does not support `norm_num_groups` because it does not use GroupNorm.") + def test_forward_with_norm_groups(self): + pass diff --git a/tests/models/autoencoders/test_models_autoencoder_rae.py b/tests/models/autoencoders/test_models_autoencoder_rae.py new file mode 100644 index 000000000000..cc8869737bcc --- /dev/null +++ b/tests/models/autoencoders/test_models_autoencoder_rae.py @@ -0,0 +1,300 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import gc + +import pytest +import torch +import torch.nn.functional as F +from torchvision.transforms.functional import to_tensor + +import diffusers.models.autoencoders.autoencoder_rae as _rae_module +from diffusers.models.autoencoders.autoencoder_rae import ( + _ENCODER_FORWARD_FNS, + AutoencoderRAE, + _build_encoder, +) +from diffusers.utils import load_image + +from ...testing_utils import ( + backend_empty_cache, + enable_full_determinism, + slow, + torch_all_close, + torch_device, +) +from ..testing_utils import BaseModelTesterConfig, ModelTesterMixin +from .testing_utils import AutoencoderTesterMixin + + +enable_full_determinism() + + +# --------------------------------------------------------------------------- +# Tiny test encoder for fast unit tests (no transformers dependency) +# --------------------------------------------------------------------------- + + +class _TinyTestEncoderModule(torch.nn.Module): + """Minimal encoder that mimics the patch-token interface without any HF model.""" + + def __init__(self, hidden_size: int = 16, patch_size: int = 8, **kwargs): + super().__init__() + self.patch_size = patch_size + self.hidden_size = hidden_size + + def forward(self, images: torch.Tensor) -> torch.Tensor: + pooled = F.avg_pool2d(images.mean(dim=1, keepdim=True), kernel_size=self.patch_size, stride=self.patch_size) + tokens = pooled.flatten(2).transpose(1, 2).contiguous() + return tokens.repeat(1, 1, self.hidden_size) + + +def _tiny_test_encoder_forward(model, images): + return model(images) + + +def _build_tiny_test_encoder(encoder_type, hidden_size, patch_size, num_hidden_layers): + return _TinyTestEncoderModule(hidden_size=hidden_size, patch_size=patch_size) + + +# Monkey-patch the dispatch tables so "tiny_test" is recognised by AutoencoderRAE +_ENCODER_FORWARD_FNS["tiny_test"] = _tiny_test_encoder_forward +_original_build_encoder = _build_encoder + + +def _patched_build_encoder(encoder_type, hidden_size, patch_size, num_hidden_layers): + if encoder_type == "tiny_test": + return _build_tiny_test_encoder(encoder_type, hidden_size, patch_size, num_hidden_layers) + return _original_build_encoder(encoder_type, hidden_size, patch_size, num_hidden_layers) + + +_rae_module._build_encoder = _patched_build_encoder + + +# --------------------------------------------------------------------------- +# Test config +# --------------------------------------------------------------------------- + + +class AutoencoderRAETesterConfig(BaseModelTesterConfig): + @property + def model_class(self): + return AutoencoderRAE + + @property + def output_shape(self): + return (3, 16, 16) + + def get_init_dict(self): + return { + "encoder_type": "tiny_test", + "encoder_hidden_size": 16, + "encoder_patch_size": 8, + "encoder_input_size": 32, + "patch_size": 4, + "image_size": 16, + "decoder_hidden_size": 32, + "decoder_num_hidden_layers": 1, + "decoder_num_attention_heads": 4, + "decoder_intermediate_size": 64, + "num_channels": 3, + "encoder_norm_mean": [0.5, 0.5, 0.5], + "encoder_norm_std": [0.5, 0.5, 0.5], + "noise_tau": 0.0, + "reshape_to_2d": True, + "scaling_factor": 1.0, + } + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_dummy_inputs(self): + return {"sample": torch.randn(2, 3, 32, 32, generator=self.generator, device="cpu").to(torch_device)} + + # Bridge for AutoencoderTesterMixin which still uses the old interface + def prepare_init_args_and_inputs_for_common(self): + return self.get_init_dict(), self.get_dummy_inputs() + + def _make_model(self, **overrides) -> AutoencoderRAE: + config = self.get_init_dict() + config.update(overrides) + return AutoencoderRAE(**config).to(torch_device) + + +class TestAutoEncoderRAE(AutoencoderRAETesterConfig, ModelTesterMixin): + """Core model tests for AutoencoderRAE.""" + + @pytest.mark.skip(reason="AutoencoderRAE does not support torch dynamo yet") + def test_from_save_pretrained_dynamo(self): ... + + def test_fast_encode_decode_and_forward_shapes(self): + model = self._make_model().eval() + x = torch.rand(2, 3, 32, 32, device=torch_device) + + with torch.no_grad(): + z = model.encode(x).latent + decoded = model.decode(z).sample + recon = model(x).sample + + assert z.shape == (2, 16, 4, 4) + assert decoded.shape == (2, 3, 16, 16) + assert recon.shape == (2, 3, 16, 16) + assert torch.isfinite(recon).all().item() + + def test_fast_scaling_factor_encode_and_decode_consistency(self): + torch.manual_seed(0) + model_base = self._make_model(scaling_factor=1.0).eval() + torch.manual_seed(0) + model_scaled = self._make_model(scaling_factor=2.0).eval() + + x = torch.rand(2, 3, 32, 32, device=torch_device) + with torch.no_grad(): + z_base = model_base.encode(x).latent + z_scaled = model_scaled.encode(x).latent + recon_base = model_base.decode(z_base).sample + recon_scaled = model_scaled.decode(z_scaled).sample + + assert torch.allclose(z_scaled, z_base * 2.0, atol=1e-5, rtol=1e-4) + assert torch.allclose(recon_scaled, recon_base, atol=1e-5, rtol=1e-4) + + def test_fast_latents_normalization_matches_formula(self): + latents_mean = torch.full((1, 16, 1, 1), 0.25, dtype=torch.float32) + latents_std = torch.full((1, 16, 1, 1), 2.0, dtype=torch.float32) + + model_raw = self._make_model().eval() + model_norm = self._make_model(latents_mean=latents_mean, latents_std=latents_std).eval() + x = torch.rand(1, 3, 32, 32, device=torch_device) + + with torch.no_grad(): + z_raw = model_raw.encode(x).latent + z_norm = model_norm.encode(x).latent + + expected = (z_raw - latents_mean.to(z_raw.device, z_raw.dtype)) / ( + latents_std.to(z_raw.device, z_raw.dtype) + 1e-5 + ) + assert torch.allclose(z_norm, expected, atol=1e-5, rtol=1e-4) + + def test_fast_slicing_matches_non_slicing(self): + model = self._make_model().eval() + x = torch.rand(3, 3, 32, 32, device=torch_device) + + with torch.no_grad(): + model.use_slicing = False + z_no_slice = model.encode(x).latent + out_no_slice = model.decode(z_no_slice).sample + + model.use_slicing = True + z_slice = model.encode(x).latent + out_slice = model.decode(z_slice).sample + + assert torch.allclose(z_slice, z_no_slice, atol=1e-6, rtol=1e-5) + assert torch.allclose(out_slice, out_no_slice, atol=1e-6, rtol=1e-5) + + def test_fast_noise_tau_applies_only_in_train(self): + model = self._make_model(noise_tau=0.5).to(torch_device) + x = torch.rand(2, 3, 32, 32, device=torch_device) + + model.train() + torch.manual_seed(0) + z_train_1 = model.encode(x).latent + torch.manual_seed(1) + z_train_2 = model.encode(x).latent + + model.eval() + torch.manual_seed(0) + z_eval_1 = model.encode(x).latent + torch.manual_seed(1) + z_eval_2 = model.encode(x).latent + + assert z_train_1.shape == z_eval_1.shape + assert not torch.allclose(z_train_1, z_train_2) + assert torch.allclose(z_eval_1, z_eval_2, atol=1e-6, rtol=1e-5) + + +class TestAutoEncoderRAESlicingTiling(AutoencoderRAETesterConfig, AutoencoderTesterMixin): + """Slicing and tiling tests for AutoencoderRAE.""" + + +@slow +@pytest.mark.skip(reason="Not enough model usage to justify slow tests yet.") +class AutoencoderRAEEncoderIntegrationTests: + def teardown_method(self): + gc.collect() + backend_empty_cache(torch_device) + + def test_dinov2_encoder_forward_shape(self): + encoder = _build_encoder("dinov2", hidden_size=768, patch_size=14, num_hidden_layers=12).to(torch_device) + x = torch.rand(1, 3, 224, 224, device=torch_device) + y = _ENCODER_FORWARD_FNS["dinov2"](encoder, x) + + assert y.ndim == 3 + assert y.shape[0] == 1 + assert y.shape[1] == 256 # (224/14)^2 - 5 (CLS + 4 register) = 251? Actually dinov2 has 256 patches + assert y.shape[2] == 768 + + def test_siglip2_encoder_forward_shape(self): + encoder = _build_encoder("siglip2", hidden_size=768, patch_size=16, num_hidden_layers=12).to(torch_device) + x = torch.rand(1, 3, 224, 224, device=torch_device) + y = _ENCODER_FORWARD_FNS["siglip2"](encoder, x) + + assert y.ndim == 3 + assert y.shape[0] == 1 + assert y.shape[1] == 196 # (224/16)^2 + assert y.shape[2] == 768 + + def test_mae_encoder_forward_shape(self): + encoder = _build_encoder("mae", hidden_size=768, patch_size=16, num_hidden_layers=12).to(torch_device) + x = torch.rand(1, 3, 224, 224, device=torch_device) + y = _ENCODER_FORWARD_FNS["mae"](encoder, x, patch_size=16) + + assert y.ndim == 3 + assert y.shape[0] == 1 + assert y.shape[1] == 196 # (224/16)^2 + assert y.shape[2] == 768 + + +@slow +@pytest.mark.skip(reason="Not enough model usage to justify slow tests yet.") +class AutoencoderRAEIntegrationTests: + def teardown_method(self): + gc.collect() + backend_empty_cache(torch_device) + + def test_autoencoder_rae_from_pretrained_dinov2(self): + model = AutoencoderRAE.from_pretrained("nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08").to(torch_device) + model.eval() + + image = load_image( + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png" + ) + image = image.convert("RGB").resize((224, 224)) + x = to_tensor(image).unsqueeze(0).to(torch_device) + + with torch.no_grad(): + latents = model.encode(x).latent + assert latents.shape == (1, 768, 16, 16) + + recon = model.decode(latents).sample + assert recon.shape == (1, 3, 256, 256) + assert torch.isfinite(recon).all().item() + + # fmt: off + expected_latent_slice = torch.tensor([0.7617, 0.8824, -0.4891]) + expected_recon_slice = torch.tensor([0.1263, 0.1355, 0.1435]) + # fmt: on + + assert torch_all_close(latents[0, :3, 0, 0].float().cpu(), expected_latent_slice, atol=1e-3) + assert torch_all_close(recon[0, 0, 0, :3].float().cpu(), expected_recon_slice, atol=1e-3) diff --git a/tests/models/controlnets/__init__.py b/tests/models/controlnets/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/models/controlnets/test_models_controlnet_cosmos.py b/tests/models/controlnets/test_models_controlnet_cosmos.py new file mode 100644 index 000000000000..5241a49261b2 --- /dev/null +++ b/tests/models/controlnets/test_models_controlnet_cosmos.py @@ -0,0 +1,275 @@ +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import torch + +from diffusers import CosmosControlNetModel +from diffusers.models.controlnets.controlnet_cosmos import CosmosControlNetOutput + +from ...testing_utils import enable_full_determinism, torch_device +from ..test_modeling_common import ModelTesterMixin + + +enable_full_determinism() + + +class CosmosControlNetModelTests(ModelTesterMixin, unittest.TestCase): + model_class = CosmosControlNetModel + main_input_name = "controls_latents" + uses_custom_attn_processor = True + + @property + def dummy_input(self): + batch_size = 1 + num_channels = 16 + num_frames = 1 + height = 16 + width = 16 + text_embed_dim = 32 + sequence_length = 12 + img_context_dim_in = 32 + img_context_num_tokens = 4 + + # Raw latents (not patchified) - the controlnet computes embeddings internally + controls_latents = torch.randn((batch_size, num_channels, num_frames, height, width)).to(torch_device) + latents = torch.randn((batch_size, num_channels, num_frames, height, width)).to(torch_device) + timestep = torch.tensor([0.5]).to(torch_device) # Diffusion timestep + condition_mask = torch.ones(batch_size, 1, num_frames, height, width).to(torch_device) + padding_mask = torch.zeros(batch_size, 1, height, width).to(torch_device) + + # Text embeddings + text_context = torch.randn((batch_size, sequence_length, text_embed_dim)).to(torch_device) + # Image context for Cosmos 2.5 + img_context = torch.randn((batch_size, img_context_num_tokens, img_context_dim_in)).to(torch_device) + encoder_hidden_states = (text_context, img_context) + + return { + "controls_latents": controls_latents, + "latents": latents, + "timestep": timestep, + "encoder_hidden_states": encoder_hidden_states, + "condition_mask": condition_mask, + "conditioning_scale": 1.0, + "padding_mask": padding_mask, + } + + @property + def input_shape(self): + return (16, 1, 16, 16) + + @property + def output_shape(self): + # Output is tuple of n_controlnet_blocks tensors, each with shape (batch, num_patches, model_channels) + # After stacking by normalize_output: (n_blocks, batch, num_patches, model_channels) + # For test config: n_blocks=2, num_patches=64 (1*8*8), model_channels=32 + # output_shape is used as (batch_size,) + output_shape, so: (2, 64, 32) + return (2, 64, 32) + + def prepare_init_args_and_inputs_for_common(self): + init_dict = { + "n_controlnet_blocks": 2, + "in_channels": 16 + 1 + 1, # control_latent_channels + condition_mask + padding_mask + "latent_channels": 16 + 1 + 1, # base_latent_channels (16) + condition_mask (1) + padding_mask (1) = 18 + "model_channels": 32, + "num_attention_heads": 2, + "attention_head_dim": 16, + "mlp_ratio": 2, + "text_embed_dim": 32, + "adaln_lora_dim": 4, + "patch_size": (1, 2, 2), + "max_size": (4, 32, 32), + "rope_scale": (2.0, 1.0, 1.0), + "extra_pos_embed_type": None, + "img_context_dim_in": 32, + "img_context_dim_out": 32, + "use_crossattn_projection": False, # Test doesn't need this projection + } + inputs_dict = self.dummy_input + return init_dict, inputs_dict + + def test_output_format(self): + """Test that the model outputs CosmosControlNetOutput with correct structure.""" + init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + model = self.model_class(**init_dict) + model.to(torch_device) + model.eval() + + with torch.no_grad(): + output = model(**inputs_dict) + + self.assertIsInstance(output, CosmosControlNetOutput) + self.assertIsInstance(output.control_block_samples, list) + self.assertEqual(len(output.control_block_samples), init_dict["n_controlnet_blocks"]) + for tensor in output.control_block_samples: + self.assertIsInstance(tensor, torch.Tensor) + + def test_output_list_format(self): + """Test that return_dict=False returns a tuple containing a list.""" + init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + model = self.model_class(**init_dict) + model.to(torch_device) + model.eval() + + with torch.no_grad(): + output = model(**inputs_dict, return_dict=False) + + self.assertIsInstance(output, tuple) + self.assertEqual(len(output), 1) + self.assertIsInstance(output[0], list) + self.assertEqual(len(output[0]), init_dict["n_controlnet_blocks"]) + + def test_condition_mask_changes_output(self): + """Test that condition mask affects control outputs.""" + init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + model = self.model_class(**init_dict) + model.to(torch_device) + model.eval() + + inputs_no_mask = dict(inputs_dict) + inputs_no_mask["condition_mask"] = torch.zeros_like(inputs_dict["condition_mask"]) + + with torch.no_grad(): + output_no_mask = model(**inputs_no_mask) + output_with_mask = model(**inputs_dict) + + self.assertEqual(len(output_no_mask.control_block_samples), len(output_with_mask.control_block_samples)) + for no_mask_tensor, with_mask_tensor in zip( + output_no_mask.control_block_samples, output_with_mask.control_block_samples + ): + self.assertFalse(torch.allclose(no_mask_tensor, with_mask_tensor)) + + def test_conditioning_scale_single(self): + """Test that a single conditioning scale is broadcast to all blocks.""" + init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + model = self.model_class(**init_dict) + model.to(torch_device) + model.eval() + + inputs_dict["conditioning_scale"] = 0.5 + + with torch.no_grad(): + output = model(**inputs_dict) + + self.assertEqual(len(output.control_block_samples), init_dict["n_controlnet_blocks"]) + + def test_conditioning_scale_list(self): + """Test that a list of conditioning scales is applied per block.""" + init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + model = self.model_class(**init_dict) + model.to(torch_device) + model.eval() + + # Provide a scale for each block + inputs_dict["conditioning_scale"] = [0.5, 1.0] + + with torch.no_grad(): + output = model(**inputs_dict) + + self.assertEqual(len(output.control_block_samples), init_dict["n_controlnet_blocks"]) + + def test_forward_with_none_img_context(self): + """Test forward pass when img_context is None.""" + init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + model = self.model_class(**init_dict) + model.to(torch_device) + model.eval() + + # Set encoder_hidden_states to (text_context, None) + text_context = inputs_dict["encoder_hidden_states"][0] + inputs_dict["encoder_hidden_states"] = (text_context, None) + + with torch.no_grad(): + output = model(**inputs_dict) + + self.assertIsInstance(output, CosmosControlNetOutput) + self.assertEqual(len(output.control_block_samples), init_dict["n_controlnet_blocks"]) + + def test_forward_without_img_context_proj(self): + """Test forward pass when img_context_proj is not configured.""" + init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + # Disable img_context_proj + init_dict["img_context_dim_in"] = None + model = self.model_class(**init_dict) + model.to(torch_device) + model.eval() + + # When img_context is disabled, pass only text context (not a tuple) + text_context = inputs_dict["encoder_hidden_states"][0] + inputs_dict["encoder_hidden_states"] = text_context + + with torch.no_grad(): + output = model(**inputs_dict) + + self.assertIsInstance(output, CosmosControlNetOutput) + self.assertEqual(len(output.control_block_samples), init_dict["n_controlnet_blocks"]) + + def test_gradient_checkpointing_is_applied(self): + expected_set = {"CosmosControlNetModel"} + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + + # Note: test_set_attn_processor_for_determinism already handles uses_custom_attn_processor=True + # so no explicit skip needed for it + # Note: test_forward_signature and test_set_default_attn_processor don't exist in base class + + # Skip tests that don't apply to this architecture + @unittest.skip("CosmosControlNetModel doesn't use norm groups.") + def test_forward_with_norm_groups(self): + pass + + # Skip tests that expect .sample attribute - ControlNets don't have this + @unittest.skip("ControlNet output doesn't have .sample attribute") + def test_effective_gradient_checkpointing(self): + pass + + # Skip tests that compute MSE loss against single tensor output + @unittest.skip("ControlNet outputs list of control blocks, not single tensor for MSE loss") + def test_ema_training(self): + pass + + @unittest.skip("ControlNet outputs list of control blocks, not single tensor for MSE loss") + def test_training(self): + pass + + # Skip tests where output shape comparison doesn't apply to ControlNets + @unittest.skip("ControlNet output shape doesn't match input shape by design") + def test_output(self): + pass + + # Skip outputs_equivalence - dict/list comparison logic not compatible (recursive_check expects dict.values()) + @unittest.skip("ControlNet output structure not compatible with recursive dict check") + def test_outputs_equivalence(self): + pass + + # Skip model parallelism - base test uses torch.allclose(base_output[0], new_output[0]) which fails + # because output[0] is the list of control_block_samples, not a tensor + @unittest.skip("test_model_parallelism uses torch.allclose on output[0] which is a list, not a tensor") + def test_model_parallelism(self): + pass + + # Skip layerwise casting tests - these have two issues: + # 1. _inference and _memory: dtype compatibility issues with learnable_pos_embed and float8/bfloat16 + # 2. _training: same as test_training - mse_loss expects tensor, not list + @unittest.skip("Layerwise casting has dtype issues with learnable_pos_embed") + def test_layerwise_casting_inference(self): + pass + + @unittest.skip("Layerwise casting has dtype issues with learnable_pos_embed") + def test_layerwise_casting_memory(self): + pass + + @unittest.skip("test_layerwise_casting_training computes mse_loss on list output") + def test_layerwise_casting_training(self): + pass diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py index b9dfe932335c..1b1a51d1e26f 100644 --- a/tests/models/test_modeling_common.py +++ b/tests/models/test_modeling_common.py @@ -26,7 +26,7 @@ import unittest.mock as mock import uuid from collections import defaultdict -from typing import Dict, List, Optional, Tuple, Union +from typing import Dict, List, Tuple import numpy as np import pytest @@ -168,8 +168,8 @@ def named_persistent_module_tensors( def compute_module_persistent_sizes( model: nn.Module, - dtype: Optional[Union[str, torch.device]] = None, - special_dtypes: Optional[Dict[str, Union[str, torch.device]]] = None, + dtype: str | torch.device | None = None, + special_dtypes: dict[str, str | torch.device] | None = None, ): """ Compute the size of each submodule of a given model (parameters + persistent buffers). diff --git a/tests/models/test_models_auto.py b/tests/models/test_models_auto.py index a70754343f30..e35fb26518ef 100644 --- a/tests/models/test_models_auto.py +++ b/tests/models/test_models_auto.py @@ -1,9 +1,15 @@ +import json +import os +import tempfile import unittest -from unittest.mock import patch +from unittest.mock import MagicMock, patch +import torch from transformers import CLIPTextModel, LongformerModel +from diffusers import ConfigMixin from diffusers.models import AutoModel, UNet2DConditionModel +from diffusers.models.modeling_utils import ModelMixin class TestAutoModel(unittest.TestCase): @@ -20,7 +26,9 @@ def test_load_from_config_diffusers_with_subfolder(self, mock_load_config): side_effect=[EnvironmentError("File not found"), {"model_type": "clip_text_model"}], ) def test_load_from_config_transformers_with_subfolder(self, mock_load_config): - model = AutoModel.from_pretrained("hf-internal-testing/tiny-stable-diffusion-torch", subfolder="text_encoder") + model = AutoModel.from_pretrained( + "hf-internal-testing/tiny-stable-diffusion-torch", subfolder="text_encoder", use_safetensors=False + ) assert isinstance(model, CLIPTextModel) def test_load_from_config_without_subfolder(self): @@ -28,5 +36,160 @@ def test_load_from_config_without_subfolder(self): assert isinstance(model, LongformerModel) def test_load_from_model_index(self): - model = AutoModel.from_pretrained("hf-internal-testing/tiny-stable-diffusion-torch", subfolder="text_encoder") + model = AutoModel.from_pretrained( + "hf-internal-testing/tiny-stable-diffusion-torch", subfolder="text_encoder", use_safetensors=False + ) assert isinstance(model, CLIPTextModel) + + def test_load_dynamic_module_from_local_path_with_subfolder(self): + CUSTOM_MODEL_CODE = ( + "import torch\n" + "from diffusers import ModelMixin, ConfigMixin\n" + "from diffusers.configuration_utils import register_to_config\n" + "\n" + "class CustomModel(ModelMixin, ConfigMixin):\n" + " @register_to_config\n" + " def __init__(self, hidden_size=8):\n" + " super().__init__()\n" + " self.linear = torch.nn.Linear(hidden_size, hidden_size)\n" + "\n" + " def forward(self, x):\n" + " return self.linear(x)\n" + ) + + with tempfile.TemporaryDirectory() as tmpdir: + subfolder = "custom_model" + model_dir = os.path.join(tmpdir, subfolder) + os.makedirs(model_dir) + + with open(os.path.join(model_dir, "modeling.py"), "w") as f: + f.write(CUSTOM_MODEL_CODE) + + config = { + "_class_name": "CustomModel", + "_diffusers_version": "0.0.0", + "auto_map": {"AutoModel": "modeling.CustomModel"}, + "hidden_size": 8, + } + with open(os.path.join(model_dir, "config.json"), "w") as f: + json.dump(config, f) + + torch.save({}, os.path.join(model_dir, "diffusion_pytorch_model.bin")) + + model = AutoModel.from_pretrained(tmpdir, subfolder=subfolder, trust_remote_code=True) + assert model.__class__.__name__ == "CustomModel" + assert model.config["hidden_size"] == 8 + + +class TestAutoModelFromConfig(unittest.TestCase): + @patch( + "diffusers.pipelines.pipeline_loading_utils.get_class_obj_and_candidates", + return_value=(MagicMock(), None), + ) + def test_from_config_with_dict_diffusers_class(self, mock_get_class): + config = {"_class_name": "UNet2DConditionModel", "sample_size": 64} + mock_model = MagicMock() + mock_get_class.return_value[0].from_config.return_value = mock_model + + result = AutoModel.from_config(config) + + mock_get_class.assert_called_once_with( + library_name="diffusers", + class_name="UNet2DConditionModel", + importable_classes=unittest.mock.ANY, + pipelines=None, + is_pipeline_module=False, + ) + mock_get_class.return_value[0].from_config.assert_called_once_with(config) + assert result is mock_model + + @patch( + "diffusers.pipelines.pipeline_loading_utils.get_class_obj_and_candidates", + return_value=(MagicMock(), None), + ) + @patch("diffusers.models.AutoModel.load_config", return_value={"_class_name": "UNet2DConditionModel"}) + def test_from_config_with_string_path(self, mock_load_config, mock_get_class): + mock_model = MagicMock() + mock_get_class.return_value[0].from_config.return_value = mock_model + + result = AutoModel.from_config("hf-internal-testing/tiny-stable-diffusion-torch", subfolder="unet") + + mock_load_config.assert_called_once() + assert result is mock_model + + def test_from_config_raises_on_missing_class_info(self): + config = {"some_key": "some_value"} + with self.assertRaises(ValueError, msg="Couldn't find a model class"): + AutoModel.from_config(config) + + @patch( + "diffusers.pipelines.pipeline_loading_utils.get_class_obj_and_candidates", + return_value=(MagicMock(), None), + ) + def test_from_config_with_model_type_routes_to_transformers(self, mock_get_class): + config = {"model_type": "clip_text_model"} + mock_model = MagicMock() + mock_get_class.return_value[0].from_config.return_value = mock_model + + result = AutoModel.from_config(config) + + mock_get_class.assert_called_once_with( + library_name="transformers", + class_name="AutoModel", + importable_classes=unittest.mock.ANY, + pipelines=None, + is_pipeline_module=False, + ) + assert result is mock_model + + def test_from_config_raises_on_none(self): + with self.assertRaises(ValueError, msg="Please provide a `pretrained_model_name_or_path_or_dict`"): + AutoModel.from_config(None) + + +class TestRegisterForAutoClass(unittest.TestCase): + def test_register_for_auto_class_sets_attribute(self): + class DummyModel(ModelMixin, ConfigMixin): + config_name = "config.json" + + DummyModel.register_for_auto_class("AutoModel") + self.assertEqual(DummyModel._auto_class, "AutoModel") + + def test_register_for_auto_class_rejects_unsupported(self): + class DummyModel(ModelMixin, ConfigMixin): + config_name = "config.json" + + with self.assertRaises(ValueError, msg="Only 'AutoModel' is supported"): + DummyModel.register_for_auto_class("AutoPipeline") + + def test_auto_map_in_saved_config(self): + class DummyModel(ModelMixin, ConfigMixin): + config_name = "config.json" + + DummyModel.register_for_auto_class("AutoModel") + model = DummyModel() + + with tempfile.TemporaryDirectory() as tmpdir: + model.save_config(tmpdir) + config_path = os.path.join(tmpdir, "config.json") + with open(config_path, "r") as f: + config = json.load(f) + + self.assertIn("auto_map", config) + self.assertIn("AutoModel", config["auto_map"]) + module_name = DummyModel.__module__.split(".")[-1] + self.assertEqual(config["auto_map"]["AutoModel"], f"{module_name}.DummyModel") + + def test_no_auto_map_without_register(self): + class DummyModel(ModelMixin, ConfigMixin): + config_name = "config.json" + + model = DummyModel() + + with tempfile.TemporaryDirectory() as tmpdir: + model.save_config(tmpdir) + config_path = os.path.join(tmpdir, "config.json") + with open(config_path, "r") as f: + config = json.load(f) + + self.assertNotIn("auto_map", config) diff --git a/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py new file mode 100644 index 000000000000..ea076b3ec774 --- /dev/null +++ b/tests/models/testing_utils/__init__.py @@ -0,0 +1,79 @@ +from .attention import AttentionTesterMixin +from .cache import ( + CacheTesterMixin, + FasterCacheConfigMixin, + FasterCacheTesterMixin, + FirstBlockCacheConfigMixin, + FirstBlockCacheTesterMixin, + PyramidAttentionBroadcastConfigMixin, + PyramidAttentionBroadcastTesterMixin, +) +from .common import BaseModelTesterConfig, ModelTesterMixin +from .compile import TorchCompileTesterMixin +from .ip_adapter import IPAdapterTesterMixin +from .lora import LoraHotSwappingForModelTesterMixin, LoraTesterMixin +from .memory import CPUOffloadTesterMixin, GroupOffloadTesterMixin, LayerwiseCastingTesterMixin, MemoryTesterMixin +from .parallelism import ContextParallelTesterMixin +from .quantization import ( + BitsAndBytesCompileTesterMixin, + BitsAndBytesConfigMixin, + BitsAndBytesTesterMixin, + GGUFCompileTesterMixin, + GGUFConfigMixin, + GGUFTesterMixin, + ModelOptCompileTesterMixin, + ModelOptConfigMixin, + ModelOptTesterMixin, + QuantizationCompileTesterMixin, + QuantizationTesterMixin, + QuantoCompileTesterMixin, + QuantoConfigMixin, + QuantoTesterMixin, + TorchAoCompileTesterMixin, + TorchAoConfigMixin, + TorchAoTesterMixin, +) +from .single_file import SingleFileTesterMixin +from .training import TrainingTesterMixin + + +__all__ = [ + "AttentionTesterMixin", + "BaseModelTesterConfig", + "BitsAndBytesCompileTesterMixin", + "BitsAndBytesConfigMixin", + "BitsAndBytesTesterMixin", + "CacheTesterMixin", + "ContextParallelTesterMixin", + "CPUOffloadTesterMixin", + "FasterCacheConfigMixin", + "FasterCacheTesterMixin", + "FirstBlockCacheConfigMixin", + "FirstBlockCacheTesterMixin", + "GGUFCompileTesterMixin", + "GGUFConfigMixin", + "GGUFTesterMixin", + "GroupOffloadTesterMixin", + "IPAdapterTesterMixin", + "LayerwiseCastingTesterMixin", + "LoraHotSwappingForModelTesterMixin", + "LoraTesterMixin", + "MemoryTesterMixin", + "ModelOptCompileTesterMixin", + "ModelOptConfigMixin", + "ModelOptTesterMixin", + "ModelTesterMixin", + "PyramidAttentionBroadcastConfigMixin", + "PyramidAttentionBroadcastTesterMixin", + "QuantizationCompileTesterMixin", + "QuantizationTesterMixin", + "QuantoCompileTesterMixin", + "QuantoConfigMixin", + "QuantoTesterMixin", + "SingleFileTesterMixin", + "TorchAoCompileTesterMixin", + "TorchAoConfigMixin", + "TorchAoTesterMixin", + "TorchCompileTesterMixin", + "TrainingTesterMixin", +] diff --git a/tests/models/testing_utils/attention.py b/tests/models/testing_utils/attention.py new file mode 100644 index 000000000000..134b3fa33bfe --- /dev/null +++ b/tests/models/testing_utils/attention.py @@ -0,0 +1,181 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import gc + +import pytest +import torch + +from diffusers.models.attention import AttentionModuleMixin +from diffusers.models.attention_processor import ( + AttnProcessor, +) + +from ...testing_utils import ( + assert_tensors_close, + backend_empty_cache, + is_attention, + torch_device, +) + + +@is_attention +class AttentionTesterMixin: + """ + Mixin class for testing attention processor and module functionality on models. + + Tests functionality from AttentionModuleMixin including: + - Attention processor management (set/get) + - QKV projection fusion/unfusion + - Attention backends (XFormers, NPU, etc.) + + Expected from config mixin: + - model_class: The model class to test + + Expected methods from config mixin: + - get_init_dict(): Returns dict of arguments to initialize the model + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest mark: attention + Use `pytest -m "not attention"` to skip these tests + """ + + def setup_method(self): + gc.collect() + backend_empty_cache(torch_device) + + def teardown_method(self): + gc.collect() + backend_empty_cache(torch_device) + + @torch.no_grad() + def test_fuse_unfuse_qkv_projections(self, atol=1e-3, rtol=0): + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict) + model.to(torch_device) + model.eval() + + if not hasattr(model, "fuse_qkv_projections"): + pytest.skip("Model does not support QKV projection fusion.") + + output_before_fusion = model(**inputs_dict, return_dict=False)[0] + + model.fuse_qkv_projections() + + has_fused_projections = False + for module in model.modules(): + if isinstance(module, AttentionModuleMixin): + if hasattr(module, "to_qkv") or hasattr(module, "to_kv"): + has_fused_projections = True + assert module.fused_projections, "fused_projections flag should be True" + break + + if has_fused_projections: + output_after_fusion = model(**inputs_dict, return_dict=False)[0] + + assert_tensors_close( + output_before_fusion, + output_after_fusion, + atol=atol, + rtol=rtol, + msg="Output should not change after fusing projections", + ) + + model.unfuse_qkv_projections() + + for module in model.modules(): + if isinstance(module, AttentionModuleMixin): + assert not hasattr(module, "to_qkv"), "to_qkv should be removed after unfusing" + assert not hasattr(module, "to_kv"), "to_kv should be removed after unfusing" + assert not module.fused_projections, "fused_projections flag should be False" + + output_after_unfusion = model(**inputs_dict, return_dict=False)[0] + + assert_tensors_close( + output_before_fusion, + output_after_unfusion, + atol=atol, + rtol=rtol, + msg="Output should match original after unfusing projections", + ) + + def test_get_set_processor(self): + init_dict = self.get_init_dict() + model = self.model_class(**init_dict) + model.to(torch_device) + + # Check if model has attention processors + if not hasattr(model, "attn_processors"): + pytest.skip("Model does not have attention processors.") + + # Test getting processors + processors = model.attn_processors + assert isinstance(processors, dict), "attn_processors should return a dict" + assert len(processors) > 0, "Model should have at least one attention processor" + + # Test that all processors can be retrieved via get_processor + for module in model.modules(): + if isinstance(module, AttentionModuleMixin): + processor = module.get_processor() + assert processor is not None, "get_processor should return a processor" + + # Test setting a new processor + new_processor = AttnProcessor() + module.set_processor(new_processor) + retrieved_processor = module.get_processor() + assert retrieved_processor is new_processor, "Retrieved processor should be the same as the one set" + + def test_attention_processor_dict(self): + init_dict = self.get_init_dict() + model = self.model_class(**init_dict) + model.to(torch_device) + + if not hasattr(model, "set_attn_processor"): + pytest.skip("Model does not support setting attention processors.") + + # Get current processors + current_processors = model.attn_processors + + # Create a dict of new processors + new_processors = {key: AttnProcessor() for key in current_processors.keys()} + + # Set processors using dict + model.set_attn_processor(new_processors) + + # Verify all processors were set + updated_processors = model.attn_processors + for key in current_processors.keys(): + assert type(updated_processors[key]) == AttnProcessor, f"Processor {key} should be AttnProcessor" + + def test_attention_processor_count_mismatch_raises_error(self): + init_dict = self.get_init_dict() + model = self.model_class(**init_dict) + model.to(torch_device) + + if not hasattr(model, "set_attn_processor"): + pytest.skip("Model does not support setting attention processors.") + + # Get current processors + current_processors = model.attn_processors + + # Create a dict with wrong number of processors + wrong_processors = {list(current_processors.keys())[0]: AttnProcessor()} + + # Verify error is raised + with pytest.raises(ValueError) as exc_info: + model.set_attn_processor(wrong_processors) + + assert "number of processors" in str(exc_info.value).lower(), "Error should mention processor count mismatch" diff --git a/tests/models/testing_utils/cache.py b/tests/models/testing_utils/cache.py new file mode 100644 index 000000000000..f1c2ecba88a7 --- /dev/null +++ b/tests/models/testing_utils/cache.py @@ -0,0 +1,556 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import gc + +import pytest +import torch + +from diffusers.hooks import FasterCacheConfig, FirstBlockCacheConfig, PyramidAttentionBroadcastConfig +from diffusers.hooks.faster_cache import _FASTER_CACHE_BLOCK_HOOK, _FASTER_CACHE_DENOISER_HOOK +from diffusers.hooks.first_block_cache import _FBC_BLOCK_HOOK, _FBC_LEADER_BLOCK_HOOK +from diffusers.hooks.pyramid_attention_broadcast import _PYRAMID_ATTENTION_BROADCAST_HOOK +from diffusers.models.cache_utils import CacheMixin + +from ...testing_utils import assert_tensors_close, backend_empty_cache, is_cache, torch_device + + +def require_cache_mixin(func): + """Decorator to skip tests if model doesn't use CacheMixin.""" + + def wrapper(self, *args, **kwargs): + if not issubclass(self.model_class, CacheMixin): + pytest.skip(f"{self.model_class.__name__} does not use CacheMixin.") + return func(self, *args, **kwargs) + + return wrapper + + +class CacheTesterMixin: + """ + Base mixin class providing common test implementations for cache testing. + + Cache-specific mixins should: + 1. Inherit from their respective config mixin (e.g., PyramidAttentionBroadcastConfigMixin) + 2. Inherit from this mixin + 3. Define the cache config to use for tests + + Expected class attributes: + - model_class: The model class to test (must use CacheMixin) + + Expected methods in test classes: + - get_init_dict(): Returns dict of arguments to initialize the model + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Optional overrides: + - cache_input_key: Property returning the input tensor key to vary between passes (default: "hidden_states") + """ + + @property + def cache_input_key(self): + return "hidden_states" + + def setup_method(self): + gc.collect() + backend_empty_cache(torch_device) + + def teardown_method(self): + gc.collect() + backend_empty_cache(torch_device) + + def _get_cache_config(self): + """ + Get the cache config for testing. + Should be implemented by subclasses. + """ + raise NotImplementedError("Subclass must implement _get_cache_config") + + def _get_hook_names(self): + """ + Get the hook names to check for this cache type. + Should be implemented by subclasses. + Returns a list of hook name strings. + """ + raise NotImplementedError("Subclass must implement _get_hook_names") + + def _test_cache_enable_disable_state(self): + """Test that cache enable/disable updates the is_cache_enabled state correctly.""" + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device) + + # Initially cache should not be enabled + assert not model.is_cache_enabled, "Cache should not be enabled initially." + + config = self._get_cache_config() + + # Enable cache + model.enable_cache(config) + assert model.is_cache_enabled, "Cache should be enabled after enable_cache()." + + # Disable cache + model.disable_cache() + assert not model.is_cache_enabled, "Cache should not be enabled after disable_cache()." + + def _test_cache_double_enable_raises_error(self): + """Test that enabling cache twice raises an error.""" + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device) + + config = self._get_cache_config() + + model.enable_cache(config) + + # Trying to enable again should raise ValueError + with pytest.raises(ValueError, match="Caching has already been enabled"): + model.enable_cache(config) + + # Cleanup + model.disable_cache() + + def _test_cache_hooks_registered(self): + """Test that cache hooks are properly registered and removed.""" + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device) + + config = self._get_cache_config() + hook_names = self._get_hook_names() + + model.enable_cache(config) + + # Check that at least one hook was registered + hook_count = 0 + for module in model.modules(): + if hasattr(module, "_diffusers_hook"): + for hook_name in hook_names: + hook = module._diffusers_hook.get_hook(hook_name) + if hook is not None: + hook_count += 1 + + assert hook_count > 0, f"At least one cache hook should be registered. Hook names: {hook_names}" + + # Disable and verify hooks are removed + model.disable_cache() + + hook_count_after = 0 + for module in model.modules(): + if hasattr(module, "_diffusers_hook"): + for hook_name in hook_names: + hook = module._diffusers_hook.get_hook(hook_name) + if hook is not None: + hook_count_after += 1 + + assert hook_count_after == 0, "Cache hooks should be removed after disable_cache()." + + @torch.no_grad() + def _test_cache_inference(self): + """Test that model can run inference with cache enabled.""" + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict).to(torch_device) + model.eval() + + config = self._get_cache_config() + + model.enable_cache(config) + + # First pass populates the cache + _ = model(**inputs_dict, return_dict=False)[0] + + # Create modified inputs for second pass (vary input tensor to simulate denoising) + inputs_dict_step2 = inputs_dict.copy() + if self.cache_input_key in inputs_dict_step2: + inputs_dict_step2[self.cache_input_key] = inputs_dict_step2[self.cache_input_key] + torch.randn_like( + inputs_dict_step2[self.cache_input_key] + ) + + # Second pass uses cached attention with different inputs (produces approximated output) + output_with_cache = model(**inputs_dict_step2, return_dict=False)[0] + + assert output_with_cache is not None, "Model output should not be None with cache enabled." + assert not torch.isnan(output_with_cache).any(), "Model output contains NaN with cache enabled." + + # Run same inputs without cache to compare + model.disable_cache() + output_without_cache = model(**inputs_dict_step2, return_dict=False)[0] + + # Cached output should be different from non-cached output (due to approximation) + assert not torch.allclose(output_without_cache, output_with_cache, atol=1e-5), ( + "Cached output should be different from non-cached output due to cache approximation." + ) + + @torch.no_grad() + def _test_cache_context_manager(self, atol=1e-5, rtol=0): + """Test the cache_context context manager properly isolates cache state.""" + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict).to(torch_device) + model.eval() + + config = self._get_cache_config() + model.enable_cache(config) + + # Run inference in first context + with model.cache_context("context_1"): + output_ctx1 = model(**inputs_dict, return_dict=False)[0] + + # Run same inference in second context (cache should be reset) + with model.cache_context("context_2"): + output_ctx2 = model(**inputs_dict, return_dict=False)[0] + + # Both contexts should produce the same output (first pass in each) + assert_tensors_close( + output_ctx1, + output_ctx2, + atol=atol, + rtol=rtol, + msg="First pass in different cache contexts should produce the same output.", + ) + + model.disable_cache() + + @torch.no_grad() + def _test_reset_stateful_cache(self): + """Test that _reset_stateful_cache resets the cache state.""" + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict).to(torch_device) + model.eval() + + config = self._get_cache_config() + + model.enable_cache(config) + + _ = model(**inputs_dict, return_dict=False)[0] + + model._reset_stateful_cache() + + model.disable_cache() + + +@is_cache +class PyramidAttentionBroadcastConfigMixin: + """ + Base mixin providing PyramidAttentionBroadcast cache config. + + Expected class attributes: + - model_class: The model class to test (must use CacheMixin) + """ + + # Default PAB config - can be overridden by subclasses + PAB_CONFIG = { + "spatial_attention_block_skip_range": 2, + } + + # Store timestep for callback (must be within default range (100, 800) for skipping to trigger) + _current_timestep = 500 + + def _get_cache_config(self): + config_kwargs = self.PAB_CONFIG.copy() + config_kwargs["current_timestep_callback"] = lambda: self._current_timestep + return PyramidAttentionBroadcastConfig(**config_kwargs) + + def _get_hook_names(self): + return [_PYRAMID_ATTENTION_BROADCAST_HOOK] + + +@is_cache +class PyramidAttentionBroadcastTesterMixin(PyramidAttentionBroadcastConfigMixin, CacheTesterMixin): + """ + Mixin class for testing PyramidAttentionBroadcast caching on models. + + Expected class attributes: + - model_class: The model class to test (must use CacheMixin) + + Expected methods to be implemented by subclasses: + - get_init_dict(): Returns dict of arguments to initialize the model + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest mark: cache + Use `pytest -m "not cache"` to skip these tests + """ + + @require_cache_mixin + def test_pab_cache_enable_disable_state(self): + self._test_cache_enable_disable_state() + + @require_cache_mixin + def test_pab_cache_double_enable_raises_error(self): + self._test_cache_double_enable_raises_error() + + @require_cache_mixin + def test_pab_cache_hooks_registered(self): + self._test_cache_hooks_registered() + + @require_cache_mixin + def test_pab_cache_inference(self): + self._test_cache_inference() + + @require_cache_mixin + def test_pab_cache_context_manager(self): + self._test_cache_context_manager() + + @require_cache_mixin + def test_pab_reset_stateful_cache(self): + self._test_reset_stateful_cache() + + +@is_cache +class FirstBlockCacheConfigMixin: + """ + Base mixin providing FirstBlockCache config. + + Expected class attributes: + - model_class: The model class to test (must use CacheMixin) + """ + + # Default FBC config - can be overridden by subclasses + # Higher threshold makes FBC more aggressive about caching (skips more often) + FBC_CONFIG = { + "threshold": 1.0, + } + + def _get_cache_config(self): + return FirstBlockCacheConfig(**self.FBC_CONFIG) + + def _get_hook_names(self): + return [_FBC_LEADER_BLOCK_HOOK, _FBC_BLOCK_HOOK] + + +@is_cache +class FirstBlockCacheTesterMixin(FirstBlockCacheConfigMixin, CacheTesterMixin): + """ + Mixin class for testing FirstBlockCache on models. + + Expected class attributes: + - model_class: The model class to test (must use CacheMixin) + + Expected methods to be implemented by subclasses: + - get_init_dict(): Returns dict of arguments to initialize the model + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest mark: cache + Use `pytest -m "not cache"` to skip these tests + """ + + @torch.no_grad() + def _test_cache_inference(self): + """Test that model can run inference with FBC cache enabled (requires cache_context).""" + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict).to(torch_device) + model.eval() + + config = self._get_cache_config() + model.enable_cache(config) + + # FBC requires cache_context to be set for inference + with model.cache_context("fbc_test"): + # First pass populates the cache + _ = model(**inputs_dict, return_dict=False)[0] + + # Create modified inputs for second pass + inputs_dict_step2 = inputs_dict.copy() + if self.cache_input_key in inputs_dict_step2: + inputs_dict_step2[self.cache_input_key] = inputs_dict_step2[self.cache_input_key] + torch.randn_like( + inputs_dict_step2[self.cache_input_key] + ) + + # Second pass - FBC should skip remaining blocks and use cached residuals + output_with_cache = model(**inputs_dict_step2, return_dict=False)[0] + + assert output_with_cache is not None, "Model output should not be None with cache enabled." + assert not torch.isnan(output_with_cache).any(), "Model output contains NaN with cache enabled." + + # Run same inputs without cache to compare + model.disable_cache() + output_without_cache = model(**inputs_dict_step2, return_dict=False)[0] + + # Cached output should be different from non-cached output (due to approximation) + assert not torch.allclose(output_without_cache, output_with_cache, atol=1e-5), ( + "Cached output should be different from non-cached output due to cache approximation." + ) + + @torch.no_grad() + def _test_reset_stateful_cache(self): + """Test that _reset_stateful_cache resets the FBC cache state (requires cache_context).""" + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict).to(torch_device) + model.eval() + + config = self._get_cache_config() + model.enable_cache(config) + + with model.cache_context("fbc_test"): + _ = model(**inputs_dict, return_dict=False)[0] + + model._reset_stateful_cache() + + model.disable_cache() + + @require_cache_mixin + def test_fbc_cache_enable_disable_state(self): + self._test_cache_enable_disable_state() + + @require_cache_mixin + def test_fbc_cache_double_enable_raises_error(self): + self._test_cache_double_enable_raises_error() + + @require_cache_mixin + def test_fbc_cache_hooks_registered(self): + self._test_cache_hooks_registered() + + @require_cache_mixin + def test_fbc_cache_inference(self): + self._test_cache_inference() + + @require_cache_mixin + def test_fbc_cache_context_manager(self): + self._test_cache_context_manager() + + @require_cache_mixin + def test_fbc_reset_stateful_cache(self): + self._test_reset_stateful_cache() + + +@is_cache +class FasterCacheConfigMixin: + """ + Base mixin providing FasterCache config. + + Expected class attributes: + - model_class: The model class to test (must use CacheMixin) + """ + + # Default FasterCache config - can be overridden by subclasses + FASTER_CACHE_CONFIG = { + "spatial_attention_block_skip_range": 2, + "spatial_attention_timestep_skip_range": (-1, 901), + "tensor_format": "BCHW", + } + + def _get_cache_config(self, current_timestep_callback=None): + config_kwargs = self.FASTER_CACHE_CONFIG.copy() + if current_timestep_callback is None: + current_timestep_callback = lambda: 1000 # noqa: E731 + config_kwargs["current_timestep_callback"] = current_timestep_callback + return FasterCacheConfig(**config_kwargs) + + def _get_hook_names(self): + return [_FASTER_CACHE_DENOISER_HOOK, _FASTER_CACHE_BLOCK_HOOK] + + +@is_cache +class FasterCacheTesterMixin(FasterCacheConfigMixin, CacheTesterMixin): + """ + Mixin class for testing FasterCache on models. + + Note: FasterCache is designed for pipeline-level inference with proper CFG batch handling + and timestep management. Inference tests are skipped at model level - FasterCache should + be tested via pipeline tests (e.g., FluxPipeline, HunyuanVideoPipeline). + + Expected class attributes: + - model_class: The model class to test (must use CacheMixin) + + Expected methods to be implemented by subclasses: + - get_init_dict(): Returns dict of arguments to initialize the model + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest mark: cache + Use `pytest -m "not cache"` to skip these tests + """ + + @torch.no_grad() + def _test_cache_inference(self): + """Test that model can run inference with FasterCache enabled.""" + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict).to(torch_device) + model.eval() + + current_timestep = [1000] + config = self._get_cache_config(current_timestep_callback=lambda: current_timestep[0]) + + model.enable_cache(config) + + # First pass with timestep outside skip range - computes and populates cache + current_timestep[0] = 1000 + _ = model(**inputs_dict, return_dict=False)[0] + + # Move timestep inside skip range so subsequent passes use cache + current_timestep[0] = 500 + + # Create modified inputs for second pass + inputs_dict_step2 = inputs_dict.copy() + if self.cache_input_key in inputs_dict_step2: + inputs_dict_step2[self.cache_input_key] = inputs_dict_step2[self.cache_input_key] + torch.randn_like( + inputs_dict_step2[self.cache_input_key] + ) + + # Second pass uses cached attention with different inputs + output_with_cache = model(**inputs_dict_step2, return_dict=False)[0] + + assert output_with_cache is not None, "Model output should not be None with cache enabled." + assert not torch.isnan(output_with_cache).any(), "Model output contains NaN with cache enabled." + + # Run same inputs without cache to compare + model.disable_cache() + output_without_cache = model(**inputs_dict_step2, return_dict=False)[0] + + # Cached output should be different from non-cached output (due to approximation) + assert not torch.allclose(output_without_cache, output_with_cache, atol=1e-5), ( + "Cached output should be different from non-cached output due to cache approximation." + ) + + @torch.no_grad() + def _test_reset_stateful_cache(self): + """Test that _reset_stateful_cache resets the FasterCache state.""" + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict).to(torch_device) + model.eval() + + config = self._get_cache_config() + model.enable_cache(config) + + _ = model(**inputs_dict, return_dict=False)[0] + + model._reset_stateful_cache() + + model.disable_cache() + + @require_cache_mixin + def test_faster_cache_enable_disable_state(self): + self._test_cache_enable_disable_state() + + @require_cache_mixin + def test_faster_cache_double_enable_raises_error(self): + self._test_cache_double_enable_raises_error() + + @require_cache_mixin + def test_faster_cache_hooks_registered(self): + self._test_cache_hooks_registered() + + @require_cache_mixin + def test_faster_cache_inference(self): + self._test_cache_inference() + + @require_cache_mixin + def test_faster_cache_context_manager(self): + self._test_cache_context_manager() + + @require_cache_mixin + def test_faster_cache_reset_stateful_cache(self): + self._test_reset_stateful_cache() diff --git a/tests/models/testing_utils/common.py b/tests/models/testing_utils/common.py new file mode 100644 index 000000000000..7036bb16203d --- /dev/null +++ b/tests/models/testing_utils/common.py @@ -0,0 +1,663 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os +from collections import defaultdict +from typing import Any, Dict, Optional, Type + +import pytest +import torch +import torch.nn as nn +from accelerate.utils.modeling import _get_proper_dtype, compute_module_sizes, dtype_byte_size + +from diffusers.utils import SAFE_WEIGHTS_INDEX_NAME, _add_variant, logging +from diffusers.utils.testing_utils import require_accelerator, require_torch_multi_accelerator + +from ...testing_utils import assert_tensors_close, torch_device + + +def named_persistent_module_tensors( + module: nn.Module, + recurse: bool = False, +): + """ + A helper function that gathers all the tensors (parameters + persistent buffers) of a given module. + + Args: + module (`torch.nn.Module`): + The module we want the tensors on. + recurse (`bool`, *optional`, defaults to `False`): + Whether or not to go look in every submodule or just return the direct parameters and buffers. + """ + yield from module.named_parameters(recurse=recurse) + + for named_buffer in module.named_buffers(recurse=recurse): + name, _ = named_buffer + # Get parent by splitting on dots and traversing the model + parent = module + if "." in name: + parent_name = name.rsplit(".", 1)[0] + for part in parent_name.split("."): + parent = getattr(parent, part) + name = name.split(".")[-1] + if name not in parent._non_persistent_buffers_set: + yield named_buffer + + +def compute_module_persistent_sizes( + model: nn.Module, + dtype: str | torch.device | None = None, + special_dtypes: dict[str, str | torch.device] | None = None, +): + """ + Compute the size of each submodule of a given model (parameters + persistent buffers). + """ + if dtype is not None: + dtype = _get_proper_dtype(dtype) + dtype_size = dtype_byte_size(dtype) + if special_dtypes is not None: + special_dtypes = {key: _get_proper_dtype(dtyp) for key, dtyp in special_dtypes.items()} + special_dtypes_size = {key: dtype_byte_size(dtyp) for key, dtyp in special_dtypes.items()} + module_sizes = defaultdict(int) + + module_list = [] + + module_list = named_persistent_module_tensors(model, recurse=True) + + for name, tensor in module_list: + if special_dtypes is not None and name in special_dtypes: + size = tensor.numel() * special_dtypes_size[name] + elif dtype is None: + size = tensor.numel() * dtype_byte_size(tensor.dtype) + elif str(tensor.dtype).startswith(("torch.uint", "torch.int", "torch.bool")): + # According to the code in set_module_tensor_to_device, these types won't be converted + # so use their original size here + size = tensor.numel() * dtype_byte_size(tensor.dtype) + else: + size = tensor.numel() * min(dtype_size, dtype_byte_size(tensor.dtype)) + name_parts = name.split(".") + for idx in range(len(name_parts) + 1): + module_sizes[".".join(name_parts[:idx])] += size + + return module_sizes + + +def calculate_expected_num_shards(index_map_path): + """ + Calculate expected number of shards from index file. + + Args: + index_map_path: Path to the sharded checkpoint index file + + Returns: + int: Expected number of shards + """ + with open(index_map_path) as f: + weight_map_dict = json.load(f)["weight_map"] + first_key = list(weight_map_dict.keys())[0] + weight_loc = weight_map_dict[first_key] # e.g., diffusion_pytorch_model-00001-of-00002.safetensors + expected_num_shards = int(weight_loc.split("-")[-1].split(".")[0]) + return expected_num_shards + + +def check_device_map_is_respected(model, device_map): + for param_name, param in model.named_parameters(): + # Find device in device_map + while len(param_name) > 0 and param_name not in device_map: + param_name = ".".join(param_name.split(".")[:-1]) + if param_name not in device_map: + raise ValueError("device map is incomplete, it does not contain any device for `param_name`.") + + param_device = device_map[param_name] + if param_device in ["cpu", "disk"]: + assert param.device == torch.device("meta"), f"Expected device 'meta' for {param_name}, got {param.device}" + else: + assert param.device == torch.device(param_device), ( + f"Expected device {param_device} for {param_name}, got {param.device}" + ) + + +def cast_inputs_to_dtype(inputs, current_dtype, target_dtype): + if torch.is_tensor(inputs): + return inputs.to(target_dtype) if inputs.dtype == current_dtype else inputs + if isinstance(inputs, dict): + return {k: cast_inputs_to_dtype(v, current_dtype, target_dtype) for k, v in inputs.items()} + if isinstance(inputs, list): + return [cast_inputs_to_dtype(v, current_dtype, target_dtype) for v in inputs] + + return inputs + + +class BaseModelTesterConfig: + """ + Base class defining the configuration interface for model testing. + + This class defines the contract that all model test classes must implement. + It provides a consistent interface for accessing model configuration, initialization + parameters, and test inputs across all testing mixins. + + Required properties (must be implemented by subclasses): + - model_class: The model class to test + + Optional properties (can be overridden, have sensible defaults): + - pretrained_model_name_or_path: Hub repository ID for pretrained model (default: None) + - pretrained_model_kwargs: Additional kwargs for from_pretrained (default: {}) + - output_shape: Expected output shape for output validation tests (default: None) + - model_split_percents: Percentages for model parallelism tests (default: [0.5, 0.7]) + + Required methods (must be implemented by subclasses): + - get_init_dict(): Returns dict of arguments to initialize the model + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Example usage: + class MyModelTestConfig(BaseModelTesterConfig): + @property + def model_class(self): + return MyModel + + @property + def pretrained_model_name_or_path(self): + return "org/my-model" + + @property + def output_shape(self): + return (1, 3, 32, 32) + + def get_init_dict(self): + return {"in_channels": 3, "out_channels": 3} + + def get_dummy_inputs(self): + return {"sample": torch.randn(1, 3, 32, 32, device=torch_device)} + + class TestMyModel(MyModelTestConfig, ModelTesterMixin, QuantizationTesterMixin): + pass + """ + + # ==================== Required Properties ==================== + + @property + def model_class(self) -> Type[nn.Module]: + """The model class to test. Must be implemented by subclasses.""" + raise NotImplementedError("Subclasses must implement the `model_class` property.") + + # ==================== Optional Properties ==================== + + @property + def pretrained_model_name_or_path(self) -> Optional[str]: + """Hub repository ID for the pretrained model (used for quantization and hub tests).""" + return None + + @property + def pretrained_model_kwargs(self) -> Dict[str, Any]: + """Additional kwargs to pass to from_pretrained (e.g., subfolder, variant).""" + return {} + + @property + def output_shape(self) -> Optional[tuple]: + """Expected output shape for output validation tests.""" + return None + + @property + def model_split_percents(self) -> list: + """Percentages for model parallelism tests.""" + return [0.9] + + # ==================== Required Methods ==================== + + def get_init_dict(self) -> Dict[str, Any]: + """ + Returns dict of arguments to initialize the model. + + Returns: + Dict[str, Any]: Initialization arguments for the model constructor. + + Example: + return { + "in_channels": 3, + "out_channels": 3, + "sample_size": 32, + } + """ + raise NotImplementedError("Subclasses must implement `get_init_dict()`.") + + def get_dummy_inputs(self) -> Dict[str, Any]: + """ + Returns dict of inputs to pass to the model forward pass. + + Returns: + Dict[str, Any]: Input tensors/values for model.forward(). + + Example: + return { + "sample": torch.randn(1, 3, 32, 32, device=torch_device), + "timestep": torch.tensor([1], device=torch_device), + } + """ + raise NotImplementedError("Subclasses must implement `get_dummy_inputs()`.") + + +class ModelTesterMixin: + """ + Base mixin class for model testing with common test methods. + + This mixin expects the test class to also inherit from BaseModelTesterConfig + (or implement its interface) which provides: + - model_class: The model class to test + - get_init_dict(): Returns dict of arguments to initialize the model + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Example: + class MyModelTestConfig(BaseModelTesterConfig): + model_class = MyModel + def get_init_dict(self): ... + def get_dummy_inputs(self): ... + + class TestMyModel(MyModelTestConfig, ModelTesterMixin): + pass + """ + + @torch.no_grad() + def test_from_save_pretrained(self, tmp_path, atol=5e-5, rtol=5e-5): + torch.manual_seed(0) + model = self.model_class(**self.get_init_dict()) + model.to(torch_device) + model.eval() + + model.save_pretrained(tmp_path) + new_model = self.model_class.from_pretrained(tmp_path) + new_model.to(torch_device) + + for param_name in model.state_dict().keys(): + param_1 = model.state_dict()[param_name] + param_2 = new_model.state_dict()[param_name] + assert param_1.shape == param_2.shape, ( + f"Parameter shape mismatch for {param_name}. Original: {param_1.shape}, loaded: {param_2.shape}" + ) + + image = model(**self.get_dummy_inputs(), return_dict=False)[0] + new_image = new_model(**self.get_dummy_inputs(), return_dict=False)[0] + + assert_tensors_close(image, new_image, atol=atol, rtol=rtol, msg="Models give different forward passes.") + + @torch.no_grad() + def test_from_save_pretrained_variant(self, tmp_path, atol=5e-5, rtol=0): + model = self.model_class(**self.get_init_dict()) + model.to(torch_device) + model.eval() + + model.save_pretrained(tmp_path, variant="fp16") + new_model = self.model_class.from_pretrained(tmp_path, variant="fp16") + + with pytest.raises(OSError) as exc_info: + self.model_class.from_pretrained(tmp_path) + + assert "Error no file named diffusion_pytorch_model.bin found in directory" in str(exc_info.value) + + new_model.to(torch_device) + + image = model(**self.get_dummy_inputs(), return_dict=False)[0] + new_image = new_model(**self.get_dummy_inputs(), return_dict=False)[0] + + assert_tensors_close(image, new_image, atol=atol, rtol=rtol, msg="Models give different forward passes.") + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16], ids=["fp32", "fp16", "bf16"]) + def test_from_save_pretrained_dtype(self, tmp_path, dtype): + model = self.model_class(**self.get_init_dict()) + model.to(torch_device) + model.eval() + + if torch_device == "mps" and dtype == torch.bfloat16: + pytest.skip(reason=f"{dtype} is not supported on {torch_device}") + + model.to(dtype) + model.save_pretrained(tmp_path) + new_model = self.model_class.from_pretrained(tmp_path, low_cpu_mem_usage=True, torch_dtype=dtype) + assert new_model.dtype == dtype + if hasattr(self.model_class, "_keep_in_fp32_modules") and self.model_class._keep_in_fp32_modules is None: + # When loading without accelerate dtype == torch.float32 if _keep_in_fp32_modules is not None + new_model = self.model_class.from_pretrained(tmp_path, low_cpu_mem_usage=False, torch_dtype=dtype) + assert new_model.dtype == dtype + + @torch.no_grad() + def test_determinism(self, atol=1e-5, rtol=0): + model = self.model_class(**self.get_init_dict()) + model.to(torch_device) + model.eval() + + first = model(**self.get_dummy_inputs(), return_dict=False)[0] + second = model(**self.get_dummy_inputs(), return_dict=False)[0] + + first_flat = first.flatten() + second_flat = second.flatten() + mask = ~(torch.isnan(first_flat) | torch.isnan(second_flat)) + first_filtered = first_flat[mask] + second_filtered = second_flat[mask] + + assert_tensors_close( + first_filtered, second_filtered, atol=atol, rtol=rtol, msg="Model outputs are not deterministic" + ) + + @torch.no_grad() + def test_output(self, expected_output_shape=None): + model = self.model_class(**self.get_init_dict()) + model.to(torch_device) + model.eval() + + inputs_dict = self.get_dummy_inputs() + output = model(**inputs_dict, return_dict=False)[0] + + assert output is not None, "Model output is None" + assert output[0].shape == expected_output_shape or self.output_shape, ( + f"Output shape does not match expected. Expected {expected_output_shape}, got {output.shape}" + ) + + @torch.no_grad() + def test_outputs_equivalence(self, atol=1e-5, rtol=0): + def set_nan_tensor_to_zero(t): + device = t.device + if device.type == "mps": + t = t.to("cpu") + t[t != t] = 0 + return t.to(device) + + def recursive_check(tuple_object, dict_object): + if isinstance(tuple_object, (list, tuple)): + for tuple_iterable_value, dict_iterable_value in zip(tuple_object, dict_object.values()): + recursive_check(tuple_iterable_value, dict_iterable_value) + elif isinstance(tuple_object, dict): + for tuple_iterable_value, dict_iterable_value in zip(tuple_object.values(), dict_object.values()): + recursive_check(tuple_iterable_value, dict_iterable_value) + elif tuple_object is None: + return + else: + assert_tensors_close( + set_nan_tensor_to_zero(tuple_object), + set_nan_tensor_to_zero(dict_object), + atol=atol, + rtol=rtol, + msg="Tuple and dict output are not equal", + ) + + model = self.model_class(**self.get_init_dict()) + model.to(torch_device) + model.eval() + + outputs_dict = model(**self.get_dummy_inputs()) + outputs_tuple = model(**self.get_dummy_inputs(), return_dict=False) + + recursive_check(outputs_tuple, outputs_dict) + + def test_getattr_is_correct(self, caplog): + init_dict = self.get_init_dict() + model = self.model_class(**init_dict) + + # save some things to test + model.dummy_attribute = 5 + model.register_to_config(test_attribute=5) + + logger_name = "diffusers.models.modeling_utils" + with caplog.at_level(logging.WARNING, logger=logger_name): + caplog.clear() + assert hasattr(model, "dummy_attribute") + assert getattr(model, "dummy_attribute") == 5 + assert model.dummy_attribute == 5 + + # no warning should be thrown + assert caplog.text == "" + + with caplog.at_level(logging.WARNING, logger=logger_name): + caplog.clear() + assert hasattr(model, "save_pretrained") + fn = model.save_pretrained + fn_1 = getattr(model, "save_pretrained") + + assert fn == fn_1 + + # no warning should be thrown + assert caplog.text == "" + + # warning should be thrown for config attributes accessed directly + with pytest.warns(FutureWarning): + assert model.test_attribute == 5 + + with pytest.warns(FutureWarning): + assert getattr(model, "test_attribute") == 5 + + with pytest.raises(AttributeError) as error: + model.does_not_exist + + assert str(error.value) == f"'{type(model).__name__}' object has no attribute 'does_not_exist'" + + @require_accelerator + @pytest.mark.skipif( + torch_device not in ["cuda", "xpu"], + reason="float16 and bfloat16 can only be used with an accelerator", + ) + def test_keep_in_fp32_modules(self, tmp_path): + model = self.model_class(**self.get_init_dict()) + fp32_modules = model._keep_in_fp32_modules + + if fp32_modules is None or len(fp32_modules) == 0: + pytest.skip("Model does not have _keep_in_fp32_modules defined.") + + # Save the model and reload with float16 dtype + # _keep_in_fp32_modules is only enforced during from_pretrained loading + model.save_pretrained(tmp_path) + model = self.model_class.from_pretrained(tmp_path, torch_dtype=torch.float16).to(torch_device) + + for name, param in model.named_parameters(): + if any(module_to_keep_in_fp32 in name.split(".") for module_to_keep_in_fp32 in fp32_modules): + assert param.dtype == torch.float32, f"Parameter {name} should be float32 but got {param.dtype}" + else: + assert param.dtype == torch.float16, f"Parameter {name} should be float16 but got {param.dtype}" + + @require_accelerator + @pytest.mark.skipif( + torch_device not in ["cuda", "xpu"], + reason="float16 and bfloat16 can only be use for inference with an accelerator", + ) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) + @torch.no_grad() + def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype, atol=1e-4, rtol=0): + model = self.model_class(**self.get_init_dict()) + model.to(torch_device) + fp32_modules = model._keep_in_fp32_modules or [] + + model.to(dtype).save_pretrained(tmp_path) + model_loaded = self.model_class.from_pretrained(tmp_path, torch_dtype=dtype).to(torch_device) + + for name, param in model_loaded.named_parameters(): + if fp32_modules and any( + module_to_keep_in_fp32 in name.split(".") for module_to_keep_in_fp32 in fp32_modules + ): + assert param.data.dtype == torch.float32 + else: + assert param.data.dtype == dtype + + inputs = cast_inputs_to_dtype(self.get_dummy_inputs(), torch.float32, dtype) + output = model(**inputs, return_dict=False)[0] + output_loaded = model_loaded(**inputs, return_dict=False)[0] + + assert_tensors_close( + output, output_loaded, atol=atol, rtol=rtol, msg=f"Loaded model output differs for {dtype}" + ) + + @require_accelerator + @torch.no_grad() + def test_sharded_checkpoints(self, tmp_path, atol=1e-5, rtol=0): + torch.manual_seed(0) + config = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**config).eval() + model = model.to(torch_device) + + base_output = model(**inputs_dict, return_dict=False)[0] + + model_size = compute_module_persistent_sizes(model)[""] + max_shard_size = int((model_size * 0.75) / (2**10)) # Convert to KB as these test models are small + + model.cpu().save_pretrained(tmp_path, max_shard_size=f"{max_shard_size}KB") + assert os.path.exists(os.path.join(tmp_path, SAFE_WEIGHTS_INDEX_NAME)), "Index file should exist" + + # Check if the right number of shards exists + expected_num_shards = calculate_expected_num_shards(os.path.join(tmp_path, SAFE_WEIGHTS_INDEX_NAME)) + actual_num_shards = len([file for file in os.listdir(tmp_path) if file.endswith(".safetensors")]) + assert actual_num_shards == expected_num_shards, ( + f"Expected {expected_num_shards} shards, got {actual_num_shards}" + ) + + new_model = self.model_class.from_pretrained(tmp_path).eval() + new_model = new_model.to(torch_device) + + torch.manual_seed(0) + inputs_dict_new = self.get_dummy_inputs() + new_output = new_model(**inputs_dict_new, return_dict=False)[0] + + assert_tensors_close( + base_output, new_output, atol=atol, rtol=rtol, msg="Output should match after sharded save/load" + ) + + @require_accelerator + @torch.no_grad() + def test_sharded_checkpoints_with_variant(self, tmp_path, atol=1e-5, rtol=0): + torch.manual_seed(0) + config = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**config).eval() + model = model.to(torch_device) + + base_output = model(**inputs_dict, return_dict=False)[0] + + model_size = compute_module_persistent_sizes(model)[""] + max_shard_size = int((model_size * 0.75) / (2**10)) # Convert to KB as these test models are small + variant = "fp16" + + model.cpu().save_pretrained(tmp_path, max_shard_size=f"{max_shard_size}KB", variant=variant) + + index_filename = _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant) + assert os.path.exists(os.path.join(tmp_path, index_filename)), ( + f"Variant index file {index_filename} should exist" + ) + + # Check if the right number of shards exists + expected_num_shards = calculate_expected_num_shards(os.path.join(tmp_path, index_filename)) + actual_num_shards = len([file for file in os.listdir(tmp_path) if file.endswith(".safetensors")]) + assert actual_num_shards == expected_num_shards, ( + f"Expected {expected_num_shards} shards, got {actual_num_shards}" + ) + + new_model = self.model_class.from_pretrained(tmp_path, variant=variant).eval() + new_model = new_model.to(torch_device) + + torch.manual_seed(0) + inputs_dict_new = self.get_dummy_inputs() + new_output = new_model(**inputs_dict_new, return_dict=False)[0] + + assert_tensors_close( + base_output, new_output, atol=atol, rtol=rtol, msg="Output should match after variant sharded save/load" + ) + + @torch.no_grad() + def test_sharded_checkpoints_with_parallel_loading(self, tmp_path, atol=1e-5, rtol=0): + from diffusers.utils import constants + + torch.manual_seed(0) + config = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**config).eval() + model = model.to(torch_device) + + base_output = model(**inputs_dict, return_dict=False)[0] + + model_size = compute_module_persistent_sizes(model)[""] + max_shard_size = int((model_size * 0.75) / (2**10)) # Convert to KB as these test models are small + + # Save original values to restore after test + original_parallel_loading = constants.HF_ENABLE_PARALLEL_LOADING + original_parallel_workers = getattr(constants, "HF_PARALLEL_WORKERS", None) + + try: + model.cpu().save_pretrained(tmp_path, max_shard_size=f"{max_shard_size}KB") + assert os.path.exists(os.path.join(tmp_path, SAFE_WEIGHTS_INDEX_NAME)), "Index file should exist" + + # Check if the right number of shards exists + expected_num_shards = calculate_expected_num_shards(os.path.join(tmp_path, SAFE_WEIGHTS_INDEX_NAME)) + actual_num_shards = len([file for file in os.listdir(tmp_path) if file.endswith(".safetensors")]) + assert actual_num_shards == expected_num_shards, ( + f"Expected {expected_num_shards} shards, got {actual_num_shards}" + ) + + # Load without parallel loading + constants.HF_ENABLE_PARALLEL_LOADING = False + model_sequential = self.model_class.from_pretrained(tmp_path).eval() + model_sequential = model_sequential.to(torch_device) + + # Load with parallel loading + constants.HF_ENABLE_PARALLEL_LOADING = True + constants.DEFAULT_HF_PARALLEL_LOADING_WORKERS = 2 + + torch.manual_seed(0) + model_parallel = self.model_class.from_pretrained(tmp_path).eval() + model_parallel = model_parallel.to(torch_device) + + torch.manual_seed(0) + inputs_dict_parallel = self.get_dummy_inputs() + output_parallel = model_parallel(**inputs_dict_parallel, return_dict=False)[0] + + assert_tensors_close( + base_output, output_parallel, atol=atol, rtol=rtol, msg="Output should match with parallel loading" + ) + + finally: + # Restore original values + constants.HF_ENABLE_PARALLEL_LOADING = original_parallel_loading + if original_parallel_workers is not None: + constants.HF_PARALLEL_WORKERS = original_parallel_workers + + @require_torch_multi_accelerator + @torch.no_grad() + def test_model_parallelism(self, tmp_path, atol=1e-5, rtol=0): + if self.model_class._no_split_modules is None: + pytest.skip("Test not supported for this model as `_no_split_modules` is not set.") + + config = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**config).eval() + + model = model.to(torch_device) + + torch.manual_seed(0) + base_output = model(**inputs_dict, return_dict=False)[0] + + model_size = compute_module_sizes(model)[""] + max_gpu_sizes = [int(p * model_size) for p in self.model_split_percents] + + model.cpu().save_pretrained(tmp_path) + + for max_size in max_gpu_sizes: + max_memory = {0: max_size, 1: model_size * 2, "cpu": model_size * 2} + new_model = self.model_class.from_pretrained(tmp_path, device_map="auto", max_memory=max_memory) + # Making sure part of the model will be on GPU 0 and GPU 1 + assert set(new_model.hf_device_map.values()) == {0, 1}, "Model should be split across GPUs" + + check_device_map_is_respected(new_model, new_model.hf_device_map) + + torch.manual_seed(0) + new_output = new_model(**inputs_dict, return_dict=False)[0] + + assert_tensors_close( + base_output, new_output, atol=atol, rtol=rtol, msg="Output should match with model parallelism" + ) diff --git a/tests/models/testing_utils/compile.py b/tests/models/testing_utils/compile.py new file mode 100644 index 000000000000..998b88fb469e --- /dev/null +++ b/tests/models/testing_utils/compile.py @@ -0,0 +1,165 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import gc +import os + +import pytest +import torch + +from ...testing_utils import ( + backend_empty_cache, + is_torch_compile, + require_accelerator, + require_torch_version_greater, + torch_device, +) + + +@is_torch_compile +@require_accelerator +@require_torch_version_greater("2.7.1") +class TorchCompileTesterMixin: + """ + Mixin class for testing torch.compile functionality on models. + + Expected from config mixin: + - model_class: The model class to test + + Optional properties: + - different_shapes_for_compilation: List of (height, width) tuples for dynamic shape testing (default: None) + + Expected methods from config mixin: + - get_init_dict(): Returns dict of arguments to initialize the model + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest mark: compile + Use `pytest -m "not compile"` to skip these tests + """ + + @property + def different_shapes_for_compilation(self) -> list[tuple[int, int]] | None: + """Optional list of (height, width) tuples for dynamic shape testing.""" + return None + + def setup_method(self): + torch.compiler.reset() + gc.collect() + backend_empty_cache(torch_device) + + def teardown_method(self): + torch.compiler.reset() + gc.collect() + backend_empty_cache(torch_device) + + @torch.no_grad() + def test_torch_compile_recompilation_and_graph_break(self): + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + + model = self.model_class(**init_dict).to(torch_device) + model.eval() + model = torch.compile(model, fullgraph=True) + + with ( + torch._inductor.utils.fresh_inductor_cache(), + torch._dynamo.config.patch(error_on_recompile=True), + ): + _ = model(**inputs_dict) + _ = model(**inputs_dict) + + @torch.no_grad() + def test_torch_compile_repeated_blocks(self, recompile_limit=1): + if self.model_class._repeated_blocks is None: + pytest.skip("Skipping test as the model class doesn't have `_repeated_blocks` set.") + + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + + model = self.model_class(**init_dict).to(torch_device) + model.eval() + model.compile_repeated_blocks(fullgraph=True) + + if self.model_class.__name__ == "UNet2DConditionModel": + recompile_limit = 2 + + with ( + torch._inductor.utils.fresh_inductor_cache(), + torch._dynamo.config.patch(recompile_limit=recompile_limit), + ): + _ = model(**inputs_dict) + _ = model(**inputs_dict) + + @torch.no_grad() + def test_compile_with_group_offloading(self): + if not self.model_class._supports_group_offloading: + pytest.skip("Model does not support group offloading.") + + torch._dynamo.config.cache_size_limit = 10000 + + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict) + model.eval() + + group_offload_kwargs = { + "onload_device": torch_device, + "offload_device": "cpu", + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": True, + "non_blocking": True, + } + model.enable_group_offload(**group_offload_kwargs) + model.compile() + + _ = model(**inputs_dict) + _ = model(**inputs_dict) + + @torch.no_grad() + def test_compile_on_different_shapes(self): + if self.different_shapes_for_compilation is None: + pytest.skip(f"Skipping as `different_shapes_for_compilation` is not set for {self.__class__.__name__}.") + torch.fx.experimental._config.use_duck_shape = False + + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device) + model.eval() + model = torch.compile(model, fullgraph=True, dynamic=True) + + for height, width in self.different_shapes_for_compilation: + with torch._dynamo.config.patch(error_on_recompile=True): + inputs_dict = self.get_dummy_inputs(height=height, width=width) + _ = model(**inputs_dict) + + @torch.no_grad() + def test_compile_works_with_aot(self, tmp_path): + from torch._inductor.package import load_package + + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + + model = self.model_class(**init_dict).to(torch_device) + exported_model = torch.export.export(model, args=(), kwargs=inputs_dict) + + package_path = os.path.join(str(tmp_path), f"{self.model_class.__name__}.pt2") + _ = torch._inductor.aoti_compile_and_package(exported_model, package_path=package_path) + assert os.path.exists(package_path), f"Package file not created at {package_path}" + loaded_binary = load_package(package_path, run_single_threaded=True) + + model.forward = loaded_binary + + _ = model(**inputs_dict) + _ = model(**inputs_dict) diff --git a/tests/models/testing_utils/ip_adapter.py b/tests/models/testing_utils/ip_adapter.py new file mode 100644 index 000000000000..632019c87499 --- /dev/null +++ b/tests/models/testing_utils/ip_adapter.py @@ -0,0 +1,158 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import gc + +import pytest +import torch + +from ...testing_utils import backend_empty_cache, is_ip_adapter, torch_device + + +def check_if_ip_adapter_correctly_set(model, processor_cls) -> bool: + """ + Check if IP Adapter processors are correctly set in the model. + + Args: + model: The model to check + + Returns: + bool: True if IP Adapter is correctly set, False otherwise + """ + for module in model.attn_processors.values(): + if isinstance(module, processor_cls): + return True + return False + + +@is_ip_adapter +class IPAdapterTesterMixin: + """ + Mixin class for testing IP Adapter functionality on models. + + Expected from config mixin: + - model_class: The model class to test + + Required properties (must be implemented by subclasses): + - ip_adapter_processor_cls: The IP Adapter processor class to use + + Required methods (must be implemented by subclasses): + - create_ip_adapter_state_dict(): Creates IP Adapter state dict for testing + - modify_inputs_for_ip_adapter(): Modifies inputs to include IP Adapter data + + Expected methods from config mixin: + - get_init_dict(): Returns dict of arguments to initialize the model + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest mark: ip_adapter + Use `pytest -m "not ip_adapter"` to skip these tests + """ + + def setup_method(self): + gc.collect() + backend_empty_cache(torch_device) + + def teardown_method(self): + gc.collect() + backend_empty_cache(torch_device) + + @property + def ip_adapter_processor_cls(self): + """IP Adapter processor class to use for testing. Must be implemented by subclasses.""" + raise NotImplementedError("Subclasses must implement the `ip_adapter_processor_cls` property.") + + def create_ip_adapter_state_dict(self, model): + raise NotImplementedError("child class must implement method to create IPAdapter State Dict") + + def modify_inputs_for_ip_adapter(self, model, inputs_dict): + raise NotImplementedError("child class must implement method to create IPAdapter model inputs") + + @torch.no_grad() + def test_load_ip_adapter(self): + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict).to(torch_device) + + torch.manual_seed(0) + output_no_adapter = model(**inputs_dict, return_dict=False)[0] + + ip_adapter_state_dict = self.create_ip_adapter_state_dict(model) + + model._load_ip_adapter_weights([ip_adapter_state_dict]) + assert check_if_ip_adapter_correctly_set(model, self.ip_adapter_processor_cls), ( + "IP Adapter processors not set correctly" + ) + + inputs_dict_with_adapter = self.modify_inputs_for_ip_adapter(model, inputs_dict.copy()) + outputs_with_adapter = model(**inputs_dict_with_adapter, return_dict=False)[0] + + assert not torch.allclose(output_no_adapter, outputs_with_adapter, atol=1e-4, rtol=1e-4), ( + "Output should differ with IP Adapter enabled" + ) + + @pytest.mark.skip( + reason="Setting IP Adapter scale is not defined at the model level. Enable this test after refactoring" + ) + def test_ip_adapter_scale(self): + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict).to(torch_device) + + ip_adapter_state_dict = self.create_ip_adapter_state_dict(model) + model._load_ip_adapter_weights([ip_adapter_state_dict]) + + inputs_dict_with_adapter = self.modify_inputs_for_ip_adapter(model, inputs_dict.copy()) + + # Test scale = 0.0 (no effect) + model.set_ip_adapter_scale(0.0) + torch.manual_seed(0) + output_scale_zero = model(**inputs_dict_with_adapter, return_dict=False)[0] + + # Test scale = 1.0 (full effect) + model.set_ip_adapter_scale(1.0) + torch.manual_seed(0) + output_scale_one = model(**inputs_dict_with_adapter, return_dict=False)[0] + + # Outputs should differ with different scales + assert not torch.allclose(output_scale_zero, output_scale_one, atol=1e-4, rtol=1e-4), ( + "Output should differ with different IP Adapter scales" + ) + + @pytest.mark.skip( + reason="Unloading IP Adapter is not defined at the model level. Enable this test after refactoring" + ) + def test_unload_ip_adapter(self): + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device) + + # Save original processors + original_processors = {k: type(v).__name__ for k, v in model.attn_processors.items()} + + # Create and load IP adapter + ip_adapter_state_dict = self.create_ip_adapter_state_dict(model) + model._load_ip_adapter_weights([ip_adapter_state_dict]) + + assert check_if_ip_adapter_correctly_set(model, self.ip_adapter_processor_cls), "IP Adapter should be set" + + # Unload IP adapter + model.unload_ip_adapter() + + assert not check_if_ip_adapter_correctly_set(model, self.ip_adapter_processor_cls), ( + "IP Adapter should be unloaded" + ) + + # Verify processors are restored + current_processors = {k: type(v).__name__ for k, v in model.attn_processors.items()} + assert original_processors == current_processors, "Processors should be restored after unload" diff --git a/tests/models/testing_utils/lora.py b/tests/models/testing_utils/lora.py new file mode 100644 index 000000000000..dfdc4835ee88 --- /dev/null +++ b/tests/models/testing_utils/lora.py @@ -0,0 +1,555 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import gc +import json +import os +import re + +import pytest +import safetensors.torch +import torch +import torch.nn as nn + +from diffusers.utils.import_utils import is_peft_available +from diffusers.utils.testing_utils import check_if_dicts_are_equal + +from ...testing_utils import ( + assert_tensors_close, + backend_empty_cache, + is_lora, + is_torch_compile, + require_peft_backend, + require_peft_version_greater, + require_torch_accelerator, + require_torch_version_greater, + torch_device, +) + + +if is_peft_available(): + from diffusers.loaders.peft import PeftAdapterMixin + + +def check_if_lora_correctly_set(model) -> bool: + """ + Check if LoRA layers are correctly set in the model. + + Args: + model: The model to check + + Returns: + bool: True if LoRA is correctly set, False otherwise + """ + from peft.tuners.tuners_utils import BaseTunerLayer + + for module in model.modules(): + if isinstance(module, BaseTunerLayer): + return True + return False + + +@is_lora +@require_peft_backend +class LoraTesterMixin: + """ + Mixin class for testing LoRA/PEFT functionality on models. + + Expected from config mixin: + - model_class: The model class to test + + Expected methods from config mixin: + - get_init_dict(): Returns dict of arguments to initialize the model + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest mark: lora + Use `pytest -m "not lora"` to skip these tests + """ + + def setup_method(self): + if not issubclass(self.model_class, PeftAdapterMixin): + pytest.skip(f"PEFT is not supported for this model ({self.model_class.__name__}).") + + @torch.no_grad() + def test_save_load_lora_adapter(self, tmp_path, rank=4, lora_alpha=4, use_dora=False, atol=1e-4, rtol=1e-4): + from peft import LoraConfig + from peft.utils import get_peft_model_state_dict + + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict).to(torch_device) + + torch.manual_seed(0) + output_no_lora = model(**inputs_dict, return_dict=False)[0] + + denoiser_lora_config = LoraConfig( + r=rank, + lora_alpha=lora_alpha, + target_modules=["to_q", "to_k", "to_v", "to_out.0"], + init_lora_weights=False, + use_dora=use_dora, + ) + model.add_adapter(denoiser_lora_config) + assert check_if_lora_correctly_set(model), "LoRA layers not set correctly" + + torch.manual_seed(0) + outputs_with_lora = model(**inputs_dict, return_dict=False)[0] + + assert not torch.allclose(output_no_lora, outputs_with_lora, atol=atol, rtol=rtol), ( + "Output should differ with LoRA enabled" + ) + + model.save_lora_adapter(tmp_path) + assert os.path.isfile(os.path.join(tmp_path, "pytorch_lora_weights.safetensors")), ( + "LoRA weights file not created" + ) + + state_dict_loaded = safetensors.torch.load_file(os.path.join(tmp_path, "pytorch_lora_weights.safetensors")) + + model.unload_lora() + assert not check_if_lora_correctly_set(model), "LoRA should be unloaded" + + model.load_lora_adapter(tmp_path, prefix=None, use_safetensors=True) + state_dict_retrieved = get_peft_model_state_dict(model, adapter_name="default_0") + + for k in state_dict_loaded: + loaded_v = state_dict_loaded[k] + retrieved_v = state_dict_retrieved[k].to(loaded_v.device) + assert_tensors_close(loaded_v, retrieved_v, atol=atol, rtol=rtol, msg=f"Mismatch in LoRA weight {k}") + + assert check_if_lora_correctly_set(model), "LoRA layers not set correctly after reload" + + torch.manual_seed(0) + outputs_with_lora_2 = model(**inputs_dict, return_dict=False)[0] + + assert not torch.allclose(output_no_lora, outputs_with_lora_2, atol=atol, rtol=rtol), ( + "Output should differ with LoRA enabled" + ) + assert_tensors_close( + outputs_with_lora, + outputs_with_lora_2, + atol=atol, + rtol=rtol, + msg="Outputs should match before and after save/load", + ) + + def test_lora_wrong_adapter_name_raises_error(self, tmp_path): + from peft import LoraConfig + + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device) + + denoiser_lora_config = LoraConfig( + r=4, + lora_alpha=4, + target_modules=["to_q", "to_k", "to_v", "to_out.0"], + init_lora_weights=False, + use_dora=False, + ) + model.add_adapter(denoiser_lora_config) + assert check_if_lora_correctly_set(model), "LoRA layers not set correctly" + + wrong_name = "foo" + with pytest.raises(ValueError) as exc_info: + model.save_lora_adapter(tmp_path, adapter_name=wrong_name) + + assert f"Adapter name {wrong_name} not found in the model." in str(exc_info.value) + + def test_lora_adapter_metadata_is_loaded_correctly(self, tmp_path, rank=4, lora_alpha=4, use_dora=False): + from peft import LoraConfig + + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device) + + denoiser_lora_config = LoraConfig( + r=rank, + lora_alpha=lora_alpha, + target_modules=["to_q", "to_k", "to_v", "to_out.0"], + init_lora_weights=False, + use_dora=use_dora, + ) + model.add_adapter(denoiser_lora_config) + metadata = model.peft_config["default"].to_dict() + assert check_if_lora_correctly_set(model), "LoRA layers not set correctly" + + model.save_lora_adapter(tmp_path) + model_file = os.path.join(tmp_path, "pytorch_lora_weights.safetensors") + assert os.path.isfile(model_file), "LoRA weights file not created" + + model.unload_lora() + assert not check_if_lora_correctly_set(model), "LoRA should be unloaded" + + model.load_lora_adapter(tmp_path, prefix=None, use_safetensors=True) + parsed_metadata = model.peft_config["default_0"].to_dict() + check_if_dicts_are_equal(metadata, parsed_metadata) + + def test_lora_adapter_wrong_metadata_raises_error(self, tmp_path): + from peft import LoraConfig + + from diffusers.loaders.lora_base import LORA_ADAPTER_METADATA_KEY + + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device) + + denoiser_lora_config = LoraConfig( + r=4, + lora_alpha=4, + target_modules=["to_q", "to_k", "to_v", "to_out.0"], + init_lora_weights=False, + use_dora=False, + ) + model.add_adapter(denoiser_lora_config) + assert check_if_lora_correctly_set(model), "LoRA layers not set correctly" + + model.save_lora_adapter(tmp_path) + model_file = os.path.join(tmp_path, "pytorch_lora_weights.safetensors") + assert os.path.isfile(model_file), "LoRA weights file not created" + + # Perturb the metadata in the state dict + loaded_state_dict = safetensors.torch.load_file(model_file) + metadata = {"format": "pt"} + lora_adapter_metadata = denoiser_lora_config.to_dict() + lora_adapter_metadata.update({"foo": 1, "bar": 2}) + for key, value in lora_adapter_metadata.items(): + if isinstance(value, set): + lora_adapter_metadata[key] = list(value) + metadata[LORA_ADAPTER_METADATA_KEY] = json.dumps(lora_adapter_metadata, indent=2, sort_keys=True) + safetensors.torch.save_file(loaded_state_dict, model_file, metadata=metadata) + + model.unload_lora() + assert not check_if_lora_correctly_set(model), "LoRA should be unloaded" + + with pytest.raises(TypeError) as exc_info: + model.load_lora_adapter(tmp_path, prefix=None, use_safetensors=True) + assert "`LoraConfig` class could not be instantiated" in str(exc_info.value) + + +@is_lora +@is_torch_compile +@require_peft_backend +@require_peft_version_greater("0.14.0") +@require_torch_version_greater("2.7.1") +@require_torch_accelerator +class LoraHotSwappingForModelTesterMixin: + """ + Mixin class for testing LoRA hot swapping functionality on models. + + Test that hotswapping does not result in recompilation on the model directly. + We're not extensively testing the hotswapping functionality since it is implemented in PEFT + and is extensively tested there. The goal of this test is specifically to ensure that + hotswapping with diffusers does not require recompilation. + + See https://github.com/huggingface/peft/blob/eaab05e18d51fb4cce20a73c9acd82a00c013b83/tests/test_gpu_examples.py#L4252 + for the analogous PEFT test. + + Expected from config mixin: + - model_class: The model class to test + + Optional properties: + - different_shapes_for_compilation: List of (height, width) tuples for dynamic compilation tests (default: None) + + Expected methods from config mixin: + - get_init_dict(): Returns dict of arguments to initialize the model + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest marks: lora, torch_compile + Use `pytest -m "not lora"` or `pytest -m "not torch_compile"` to skip these tests + """ + + @property + def different_shapes_for_compilation(self) -> list[tuple[int, int]] | None: + """Optional list of (height, width) tuples for dynamic compilation tests.""" + return None + + def setup_method(self): + if not issubclass(self.model_class, PeftAdapterMixin): + pytest.skip(f"PEFT is not supported for this model ({self.model_class.__name__}).") + + def teardown_method(self): + # It is critical that the dynamo cache is reset for each test. Otherwise, if the test re-uses the same model, + # there will be recompilation errors, as torch caches the model when run in the same process. + torch.compiler.reset() + gc.collect() + backend_empty_cache(torch_device) + + def _get_lora_config(self, lora_rank, lora_alpha, target_modules): + from peft import LoraConfig + + lora_config = LoraConfig( + r=lora_rank, + lora_alpha=lora_alpha, + target_modules=target_modules, + init_lora_weights=False, + use_dora=False, + ) + return lora_config + + def _get_linear_module_name_other_than_attn(self, model): + linear_names = [ + name for name, module in model.named_modules() if isinstance(module, nn.Linear) and "to_" not in name + ] + return linear_names[0] + + def _check_model_hotswap( + self, tmp_path, do_compile, rank0, rank1, target_modules0, target_modules1=None, atol=5e-3, rtol=5e-3 + ): + """ + Check that hotswapping works on a model. + + Steps: + - create 2 LoRA adapters and save them + - load the first adapter + - hotswap the second adapter + - check that the outputs are correct + - optionally compile the model + - optionally check if recompilations happen on different shapes + + Note: We set rank == alpha here because save_lora_adapter does not save the alpha scalings, thus the test would + fail if the values are different. Since rank != alpha does not matter for the purpose of this test, this is + fine. + """ + different_shapes = self.different_shapes_for_compilation + # create 2 adapters with different ranks and alphas + torch.manual_seed(0) + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict).to(torch_device) + + alpha0, alpha1 = rank0, rank1 + max_rank = max([rank0, rank1]) + if target_modules1 is None: + target_modules1 = target_modules0[:] + lora_config0 = self._get_lora_config(rank0, alpha0, target_modules0) + lora_config1 = self._get_lora_config(rank1, alpha1, target_modules1) + + model.add_adapter(lora_config0, adapter_name="adapter0") + with torch.inference_mode(): + torch.manual_seed(0) + output0_before = model(**inputs_dict)["sample"] + + model.add_adapter(lora_config1, adapter_name="adapter1") + model.set_adapter("adapter1") + with torch.inference_mode(): + torch.manual_seed(0) + output1_before = model(**inputs_dict)["sample"] + + # sanity checks: + assert not torch.allclose(output0_before, output1_before, atol=atol, rtol=rtol) + assert not (output0_before == 0).all() + assert not (output1_before == 0).all() + + # save the adapter checkpoints + model.save_lora_adapter(os.path.join(tmp_path, "0"), safe_serialization=True, adapter_name="adapter0") + model.save_lora_adapter(os.path.join(tmp_path, "1"), safe_serialization=True, adapter_name="adapter1") + del model + + # load the first adapter + torch.manual_seed(0) + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device) + + if do_compile or (rank0 != rank1): + # no need to prepare if the model is not compiled or if the ranks are identical + model.enable_lora_hotswap(target_rank=max_rank) + + file_name0 = os.path.join(os.path.join(tmp_path, "0"), "pytorch_lora_weights.safetensors") + file_name1 = os.path.join(os.path.join(tmp_path, "1"), "pytorch_lora_weights.safetensors") + model.load_lora_adapter(file_name0, safe_serialization=True, adapter_name="adapter0", prefix=None) + + if do_compile: + model = torch.compile(model, mode="reduce-overhead", dynamic=different_shapes is not None) + + with torch.inference_mode(): + # additionally check if dynamic compilation works. + if different_shapes is not None: + for height, width in different_shapes: + new_inputs_dict = self.get_dummy_inputs(height=height, width=width) + _ = model(**new_inputs_dict) + else: + output0_after = model(**inputs_dict)["sample"] + assert_tensors_close( + output0_before, output0_after, atol=atol, rtol=rtol, msg="Output mismatch after loading adapter0" + ) + + # hotswap the 2nd adapter + model.load_lora_adapter(file_name1, adapter_name="adapter0", hotswap=True, prefix=None) + + # we need to call forward to potentially trigger recompilation + with torch.inference_mode(): + if different_shapes is not None: + for height, width in different_shapes: + new_inputs_dict = self.get_dummy_inputs(height=height, width=width) + _ = model(**new_inputs_dict) + else: + output1_after = model(**inputs_dict)["sample"] + assert_tensors_close( + output1_before, + output1_after, + atol=atol, + rtol=rtol, + msg="Output mismatch after hotswapping to adapter1", + ) + + # check error when not passing valid adapter name + name = "does-not-exist" + msg = f"Trying to hotswap LoRA adapter '{name}' but there is no existing adapter by that name" + with pytest.raises(ValueError, match=re.escape(msg)): + model.load_lora_adapter(file_name1, adapter_name=name, hotswap=True, prefix=None) + + @pytest.mark.parametrize("rank0,rank1", [(11, 11), (7, 13), (13, 7)]) + def test_hotswapping_model(self, tmp_path, rank0, rank1): + self._check_model_hotswap( + tmp_path, do_compile=False, rank0=rank0, rank1=rank1, target_modules0=["to_q", "to_k", "to_v", "to_out.0"] + ) + + @pytest.mark.parametrize("rank0,rank1", [(11, 11), (7, 13), (13, 7)]) + def test_hotswapping_compiled_model_linear(self, tmp_path, rank0, rank1): + # It's important to add this context to raise an error on recompilation + target_modules = ["to_q", "to_k", "to_v", "to_out.0"] + with torch._dynamo.config.patch(error_on_recompile=True), torch._inductor.utils.fresh_inductor_cache(): + self._check_model_hotswap( + tmp_path, do_compile=True, rank0=rank0, rank1=rank1, target_modules0=target_modules + ) + + @pytest.mark.parametrize("rank0,rank1", [(11, 11), (7, 13), (13, 7)]) + def test_hotswapping_compiled_model_conv2d(self, tmp_path, rank0, rank1): + if "unet" not in self.model_class.__name__.lower(): + pytest.skip("Test only applies to UNet.") + + # It's important to add this context to raise an error on recompilation + target_modules = ["conv", "conv1", "conv2"] + with torch._dynamo.config.patch(error_on_recompile=True), torch._inductor.utils.fresh_inductor_cache(): + self._check_model_hotswap( + tmp_path, do_compile=True, rank0=rank0, rank1=rank1, target_modules0=target_modules + ) + + @pytest.mark.parametrize("rank0,rank1", [(11, 11), (7, 13), (13, 7)]) + def test_hotswapping_compiled_model_both_linear_and_conv2d(self, tmp_path, rank0, rank1): + if "unet" not in self.model_class.__name__.lower(): + pytest.skip("Test only applies to UNet.") + + # It's important to add this context to raise an error on recompilation + target_modules = ["to_q", "conv"] + with torch._dynamo.config.patch(error_on_recompile=True), torch._inductor.utils.fresh_inductor_cache(): + self._check_model_hotswap( + tmp_path, do_compile=True, rank0=rank0, rank1=rank1, target_modules0=target_modules + ) + + @pytest.mark.parametrize("rank0,rank1", [(11, 11), (7, 13), (13, 7)]) + def test_hotswapping_compiled_model_both_linear_and_other(self, tmp_path, rank0, rank1): + # In `test_hotswapping_compiled_model_both_linear_and_conv2d()`, we check if we can do hotswapping + # with `torch.compile()` for models that have both linear and conv layers. In this test, we check + # if we can target a linear layer from the transformer blocks and another linear layer from non-attention + # block. + target_modules = ["to_q"] + init_dict = self.get_init_dict() + model = self.model_class(**init_dict) + + target_modules.append(self._get_linear_module_name_other_than_attn(model)) + del model + + # It's important to add this context to raise an error on recompilation + with torch._dynamo.config.patch(error_on_recompile=True): + self._check_model_hotswap( + tmp_path, do_compile=True, rank0=rank0, rank1=rank1, target_modules0=target_modules + ) + + def test_enable_lora_hotswap_called_after_adapter_added_raises(self): + # ensure that enable_lora_hotswap is called before loading the first adapter + lora_config = self._get_lora_config(8, 8, target_modules=["to_q"]) + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device) + model.add_adapter(lora_config) + + msg = re.escape("Call `enable_lora_hotswap` before loading the first adapter.") + with pytest.raises(RuntimeError, match=msg): + model.enable_lora_hotswap(target_rank=32) + + def test_enable_lora_hotswap_called_after_adapter_added_warning(self, caplog): + # ensure that enable_lora_hotswap is called before loading the first adapter + import logging + + lora_config = self._get_lora_config(8, 8, target_modules=["to_q"]) + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device) + model.add_adapter(lora_config) + msg = ( + "It is recommended to call `enable_lora_hotswap` before loading the first adapter to avoid recompilation." + ) + with caplog.at_level(logging.WARNING): + model.enable_lora_hotswap(target_rank=32, check_compiled="warn") + assert any(msg in record.message for record in caplog.records) + + def test_enable_lora_hotswap_called_after_adapter_added_ignore(self, caplog): + # check possibility to ignore the error/warning + import logging + + lora_config = self._get_lora_config(8, 8, target_modules=["to_q"]) + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device) + model.add_adapter(lora_config) + with caplog.at_level(logging.WARNING): + model.enable_lora_hotswap(target_rank=32, check_compiled="ignore") + assert len(caplog.records) == 0 + + def test_enable_lora_hotswap_wrong_check_compiled_argument_raises(self): + # check that wrong argument value raises an error + lora_config = self._get_lora_config(8, 8, target_modules=["to_q"]) + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device) + model.add_adapter(lora_config) + msg = re.escape("check_compiles should be one of 'error', 'warn', or 'ignore', got 'wrong-argument' instead.") + with pytest.raises(ValueError, match=msg): + model.enable_lora_hotswap(target_rank=32, check_compiled="wrong-argument") + + def test_hotswap_second_adapter_targets_more_layers_raises(self, tmp_path, caplog): + # check the error and log + import logging + + # at the moment, PEFT requires the 2nd adapter to target the same or a subset of layers + target_modules0 = ["to_q"] + target_modules1 = ["to_q", "to_k"] + with pytest.raises(RuntimeError): # peft raises RuntimeError + with caplog.at_level(logging.ERROR): + self._check_model_hotswap( + tmp_path, + do_compile=True, + rank0=8, + rank1=8, + target_modules0=target_modules0, + target_modules1=target_modules1, + ) + assert any("Hotswapping adapter0 was unsuccessful" in record.message for record in caplog.records) + + @pytest.mark.parametrize("rank0,rank1", [(11, 11), (7, 13), (13, 7)]) + @require_torch_version_greater("2.7.1") + def test_hotswapping_compile_on_different_shapes(self, tmp_path, rank0, rank1): + different_shapes_for_compilation = self.different_shapes_for_compilation + if different_shapes_for_compilation is None: + pytest.skip(f"Skipping as `different_shapes_for_compilation` is not set for {self.__class__.__name__}.") + # Specifying `use_duck_shape=False` instructs the compiler if it should use the same symbolic + # variable to represent input sizes that are the same. For more details, + # check out this [comment](https://github.com/huggingface/diffusers/pull/11327#discussion_r2047659790). + torch.fx.experimental._config.use_duck_shape = False + + target_modules = ["to_q", "to_k", "to_v", "to_out.0"] + with torch._dynamo.config.patch(error_on_recompile=True): + self._check_model_hotswap( + tmp_path, + do_compile=True, + rank0=rank0, + rank1=rank1, + target_modules0=target_modules, + ) diff --git a/tests/models/testing_utils/memory.py b/tests/models/testing_utils/memory.py new file mode 100644 index 000000000000..68480bddd39c --- /dev/null +++ b/tests/models/testing_utils/memory.py @@ -0,0 +1,498 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import gc +import glob +import inspect +from functools import wraps + +import pytest +import torch +from accelerate.utils.modeling import compute_module_sizes + +from diffusers.utils.testing_utils import _check_safetensors_serialization +from diffusers.utils.torch_utils import get_torch_cuda_device_capability + +from ...testing_utils import ( + assert_tensors_close, + backend_empty_cache, + backend_max_memory_allocated, + backend_reset_peak_memory_stats, + backend_synchronize, + is_cpu_offload, + is_group_offload, + is_memory, + require_accelerator, + torch_device, +) +from .common import cast_inputs_to_dtype, check_device_map_is_respected + + +def require_offload_support(func): + """ + Decorator to skip tests if model doesn't support offloading (requires _no_split_modules). + """ + + @wraps(func) + def wrapper(self, *args, **kwargs): + if self.model_class._no_split_modules is None: + pytest.skip("Test not supported for this model as `_no_split_modules` is not set.") + return func(self, *args, **kwargs) + + return wrapper + + +def require_group_offload_support(func): + """ + Decorator to skip tests if model doesn't support group offloading. + """ + + @wraps(func) + def wrapper(self, *args, **kwargs): + if not self.model_class._supports_group_offloading: + pytest.skip("Model does not support group offloading.") + return func(self, *args, **kwargs) + + return wrapper + + +@is_cpu_offload +class CPUOffloadTesterMixin: + """ + Mixin class for testing CPU offloading functionality. + + Expected from config mixin: + - model_class: The model class to test + + Optional properties: + - model_split_percents: List of percentages for splitting model across devices (default: [0.5, 0.7]) + + Expected methods from config mixin: + - get_init_dict(): Returns dict of arguments to initialize the model + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest mark: cpu_offload + Use `pytest -m "not cpu_offload"` to skip these tests + """ + + @property + def model_split_percents(self) -> list[float]: + """List of percentages for splitting model across devices during offloading tests.""" + return [0.5, 0.7] + + @require_offload_support + @torch.no_grad() + def test_cpu_offload(self, tmp_path, atol=1e-5, rtol=0): + config = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**config).eval() + + model = model.to(torch_device) + + torch.manual_seed(0) + base_output = model(**inputs_dict) + + model_size = compute_module_sizes(model)[""] + # We test several splits of sizes to make sure it works + max_gpu_sizes = [int(p * model_size) for p in self.model_split_percents] + model.cpu().save_pretrained(str(tmp_path)) + + for max_size in max_gpu_sizes: + max_memory = {0: max_size, "cpu": model_size * 2} + new_model = self.model_class.from_pretrained(str(tmp_path), device_map="auto", max_memory=max_memory) + # Making sure part of the model will actually end up offloaded + assert set(new_model.hf_device_map.values()) == {0, "cpu"}, "Model should be split between GPU and CPU" + + check_device_map_is_respected(new_model, new_model.hf_device_map) + torch.manual_seed(0) + new_output = new_model(**inputs_dict) + + assert_tensors_close( + base_output[0], new_output[0], atol=atol, rtol=rtol, msg="Output should match with CPU offloading" + ) + + @require_offload_support + @torch.no_grad() + def test_disk_offload_without_safetensors(self, tmp_path, atol=1e-5, rtol=0): + config = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**config).eval() + + model = model.to(torch_device) + + torch.manual_seed(0) + base_output = model(**inputs_dict) + + model_size = compute_module_sizes(model)[""] + max_size = int(self.model_split_percents[0] * model_size) + # Force disk offload by setting very small CPU memory + max_memory = {0: max_size, "cpu": int(0.1 * max_size)} + + model.cpu().save_pretrained(str(tmp_path), safe_serialization=False) + # This errors out because it's missing an offload folder + with pytest.raises(ValueError): + new_model = self.model_class.from_pretrained(str(tmp_path), device_map="auto", max_memory=max_memory) + + new_model = self.model_class.from_pretrained( + str(tmp_path), device_map="auto", max_memory=max_memory, offload_folder=str(tmp_path) + ) + + check_device_map_is_respected(new_model, new_model.hf_device_map) + torch.manual_seed(0) + new_output = new_model(**inputs_dict) + + assert_tensors_close( + base_output[0], new_output[0], atol=atol, rtol=rtol, msg="Output should match with disk offloading" + ) + + @require_offload_support + @torch.no_grad() + def test_disk_offload_with_safetensors(self, tmp_path, atol=1e-5, rtol=0): + config = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**config).eval() + + model = model.to(torch_device) + + torch.manual_seed(0) + base_output = model(**inputs_dict) + + model_size = compute_module_sizes(model)[""] + model.cpu().save_pretrained(str(tmp_path)) + + max_size = int(self.model_split_percents[0] * model_size) + max_memory = {0: max_size, "cpu": max_size} + new_model = self.model_class.from_pretrained( + str(tmp_path), device_map="auto", offload_folder=str(tmp_path), max_memory=max_memory + ) + + check_device_map_is_respected(new_model, new_model.hf_device_map) + torch.manual_seed(0) + new_output = new_model(**inputs_dict) + + assert_tensors_close( + base_output[0], + new_output[0], + atol=atol, + rtol=rtol, + msg="Output should match with disk offloading (safetensors)", + ) + + +@is_group_offload +class GroupOffloadTesterMixin: + """ + Mixin class for testing group offloading functionality. + + Expected from config mixin: + - model_class: The model class to test + + Expected methods from config mixin: + - get_init_dict(): Returns dict of arguments to initialize the model + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest mark: group_offload + Use `pytest -m "not group_offload"` to skip these tests + """ + + @require_group_offload_support + @pytest.mark.parametrize("record_stream", [False, True]) + def test_group_offloading(self, record_stream, atol=1e-5, rtol=0): + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + torch.manual_seed(0) + + @torch.no_grad() + def run_forward(model): + assert all( + module._diffusers_hook.get_hook("group_offloading") is not None + for module in model.modules() + if hasattr(module, "_diffusers_hook") + ), "Group offloading hook should be set" + model.eval() + return model(**inputs_dict)[0] + + model = self.model_class(**init_dict) + + model.to(torch_device) + output_without_group_offloading = run_forward(model) + + torch.manual_seed(0) + model = self.model_class(**init_dict) + model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=1) + output_with_group_offloading1 = run_forward(model) + + torch.manual_seed(0) + model = self.model_class(**init_dict) + model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=1, non_blocking=True) + output_with_group_offloading2 = run_forward(model) + + torch.manual_seed(0) + model = self.model_class(**init_dict) + model.enable_group_offload(torch_device, offload_type="leaf_level") + output_with_group_offloading3 = run_forward(model) + + torch.manual_seed(0) + model = self.model_class(**init_dict) + model.enable_group_offload( + torch_device, offload_type="leaf_level", use_stream=True, record_stream=record_stream + ) + output_with_group_offloading4 = run_forward(model) + + assert_tensors_close( + output_without_group_offloading, + output_with_group_offloading1, + atol=atol, + rtol=rtol, + msg="Output should match with block-level offloading", + ) + assert_tensors_close( + output_without_group_offloading, + output_with_group_offloading2, + atol=atol, + rtol=rtol, + msg="Output should match with non-blocking block-level offloading", + ) + assert_tensors_close( + output_without_group_offloading, + output_with_group_offloading3, + atol=atol, + rtol=rtol, + msg="Output should match with leaf-level offloading", + ) + assert_tensors_close( + output_without_group_offloading, + output_with_group_offloading4, + atol=atol, + rtol=rtol, + msg="Output should match with leaf-level offloading with stream", + ) + + @require_group_offload_support + @pytest.mark.parametrize("record_stream", [False, True]) + @pytest.mark.parametrize("offload_type", ["block_level", "leaf_level"]) + @torch.no_grad() + def test_group_offloading_with_layerwise_casting(self, record_stream, offload_type): + torch.manual_seed(0) + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict) + + model.to(torch_device) + model.eval() + _ = model(**inputs_dict)[0] + + torch.manual_seed(0) + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + storage_dtype, compute_dtype = torch.float16, torch.float32 + inputs_dict = cast_inputs_to_dtype(inputs_dict, torch.float32, compute_dtype) + model = self.model_class(**init_dict) + model.eval() + additional_kwargs = {} if offload_type == "leaf_level" else {"num_blocks_per_group": 1} + model.enable_group_offload( + torch_device, offload_type=offload_type, use_stream=True, record_stream=record_stream, **additional_kwargs + ) + model.enable_layerwise_casting(storage_dtype=storage_dtype, compute_dtype=compute_dtype) + _ = model(**inputs_dict)[0] + + @require_group_offload_support + @pytest.mark.parametrize("record_stream", [False, True]) + @pytest.mark.parametrize("offload_type", ["block_level", "leaf_level"]) + @torch.no_grad() + @torch.inference_mode() + def test_group_offloading_with_disk(self, tmp_path, record_stream, offload_type, atol=1e-5, rtol=0): + def _has_generator_arg(model): + sig = inspect.signature(model.forward) + params = sig.parameters + return "generator" in params + + def _run_forward(model, inputs_dict): + accepts_generator = _has_generator_arg(model) + if accepts_generator: + inputs_dict["generator"] = torch.manual_seed(0) + torch.manual_seed(0) + return model(**inputs_dict)[0] + + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + torch.manual_seed(0) + model = self.model_class(**init_dict) + + model.eval() + model.to(torch_device) + output_without_group_offloading = _run_forward(model, inputs_dict) + + torch.manual_seed(0) + model = self.model_class(**init_dict) + model.eval() + + num_blocks_per_group = None if offload_type == "leaf_level" else 1 + additional_kwargs = {} if offload_type == "leaf_level" else {"num_blocks_per_group": num_blocks_per_group} + tmpdir = str(tmp_path) + model.enable_group_offload( + torch_device, + offload_type=offload_type, + offload_to_disk_path=tmpdir, + use_stream=True, + record_stream=record_stream, + **additional_kwargs, + ) + has_safetensors = glob.glob(f"{tmpdir}/*.safetensors") + assert has_safetensors, "No safetensors found in the directory." + + # For "leaf-level", there is a prefetching hook which makes this check a bit non-deterministic + # in nature. So, skip it. + if offload_type != "leaf_level": + is_correct, extra_files, missing_files = _check_safetensors_serialization( + module=model, + offload_to_disk_path=tmpdir, + offload_type=offload_type, + num_blocks_per_group=num_blocks_per_group, + ) + if not is_correct: + if extra_files: + raise ValueError(f"Found extra files: {', '.join(extra_files)}") + elif missing_files: + raise ValueError(f"Following files are missing: {', '.join(missing_files)}") + + output_with_group_offloading = _run_forward(model, inputs_dict) + assert_tensors_close( + output_without_group_offloading, + output_with_group_offloading, + atol=atol, + rtol=rtol, + msg="Output should match with disk-based group offloading", + ) + + +class LayerwiseCastingTesterMixin: + """ + Mixin class for testing layerwise dtype casting for memory optimization. + + Expected from config mixin: + - model_class: The model class to test + + Expected methods from config mixin: + - get_init_dict(): Returns dict of arguments to initialize the model + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + """ + + @torch.no_grad() + def test_layerwise_casting_memory(self): + MB_TOLERANCE = 0.2 + LEAST_COMPUTE_CAPABILITY = 8.0 + + def reset_memory_stats(): + gc.collect() + backend_synchronize(torch_device) + backend_empty_cache(torch_device) + backend_reset_peak_memory_stats(torch_device) + + def get_memory_usage(storage_dtype, compute_dtype): + torch.manual_seed(0) + config = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + inputs_dict = cast_inputs_to_dtype(inputs_dict, torch.float32, compute_dtype) + model = self.model_class(**config).eval() + model = model.to(torch_device, dtype=compute_dtype) + model.enable_layerwise_casting(storage_dtype=storage_dtype, compute_dtype=compute_dtype) + + reset_memory_stats() + model(**inputs_dict) + model_memory_footprint = model.get_memory_footprint() + peak_inference_memory_allocated_mb = backend_max_memory_allocated(torch_device) / 1024**2 + + return model_memory_footprint, peak_inference_memory_allocated_mb + + fp32_memory_footprint, fp32_max_memory = get_memory_usage(torch.float32, torch.float32) + fp8_e4m3_fp32_memory_footprint, fp8_e4m3_fp32_max_memory = get_memory_usage(torch.float8_e4m3fn, torch.float32) + fp8_e4m3_bf16_memory_footprint, fp8_e4m3_bf16_max_memory = get_memory_usage( + torch.float8_e4m3fn, torch.bfloat16 + ) + + compute_capability = get_torch_cuda_device_capability() if torch_device == "cuda" else None + assert fp8_e4m3_bf16_memory_footprint < fp8_e4m3_fp32_memory_footprint < fp32_memory_footprint, ( + "Memory footprint should decrease with lower precision storage" + ) + + # NOTE: the following assertion would fail on our CI (running Tesla T4) due to bf16 using more memory than fp32. + # On other devices, such as DGX (Ampere) and Audace (Ada), the test passes. So, we conditionally check it. + if compute_capability and compute_capability >= LEAST_COMPUTE_CAPABILITY: + assert fp8_e4m3_bf16_max_memory < fp8_e4m3_fp32_max_memory, ( + "Peak memory should be lower with bf16 compute on newer GPUs" + ) + + # On this dummy test case with a small model, sometimes fp8_e4m3_fp32 max memory usage is higher than fp32 by a few + # bytes. This only happens for some models, so we allow a small tolerance. + # For any real model being tested, the order would be fp8_e4m3_bf16 < fp8_e4m3_fp32 < fp32. + assert ( + fp8_e4m3_fp32_max_memory < fp32_max_memory + or abs(fp8_e4m3_fp32_max_memory - fp32_max_memory) < MB_TOLERANCE + ), "Peak memory should be lower or within tolerance with fp8 storage" + + def test_layerwise_casting_training(self): + def test_fn(storage_dtype, compute_dtype): + if torch.device(torch_device).type == "cpu" and compute_dtype == torch.bfloat16: + pytest.skip("Skipping test because CPU doesn't go well with bfloat16.") + + model = self.model_class(**self.get_init_dict()) + model = model.to(torch_device, dtype=compute_dtype) + model.enable_layerwise_casting(storage_dtype=storage_dtype, compute_dtype=compute_dtype) + model.train() + + inputs_dict = self.get_dummy_inputs() + inputs_dict = cast_inputs_to_dtype(inputs_dict, torch.float32, compute_dtype) + with torch.amp.autocast(device_type=torch.device(torch_device).type): + output = model(**inputs_dict, return_dict=False)[0] + + input_tensor = inputs_dict[self.main_input_name] + noise = torch.randn((input_tensor.shape[0],) + self.output_shape).to(torch_device) + noise = cast_inputs_to_dtype(noise, torch.float32, compute_dtype) + loss = torch.nn.functional.mse_loss(output, noise) + + loss.backward() + + test_fn(torch.float16, torch.float32) + test_fn(torch.float8_e4m3fn, torch.float32) + test_fn(torch.float8_e5m2, torch.float32) + test_fn(torch.float8_e4m3fn, torch.bfloat16) + + +@is_memory +@require_accelerator +class MemoryTesterMixin(CPUOffloadTesterMixin, GroupOffloadTesterMixin, LayerwiseCastingTesterMixin): + """ + Combined mixin class for all memory optimization tests including CPU/disk offloading, + group offloading, and layerwise dtype casting. + + This mixin inherits from: + - CPUOffloadTesterMixin: CPU and disk offloading tests + - GroupOffloadTesterMixin: Group offloading tests (block-level and leaf-level) + - LayerwiseCastingTesterMixin: Layerwise dtype casting tests + + Expected from config mixin: + - model_class: The model class to test + + Optional properties: + - model_split_percents: List of percentages for splitting model across devices (default: [0.5, 0.7]) + + Expected methods from config mixin: + - get_init_dict(): Returns dict of arguments to initialize the model + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest mark: memory + Use `pytest -m "not memory"` to skip these tests + """ diff --git a/tests/models/testing_utils/parallelism.py b/tests/models/testing_utils/parallelism.py new file mode 100644 index 000000000000..e05b36799e66 --- /dev/null +++ b/tests/models/testing_utils/parallelism.py @@ -0,0 +1,128 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import socket + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from diffusers.models._modeling_parallel import ContextParallelConfig + +from ...testing_utils import ( + is_context_parallel, + require_torch_multi_accelerator, +) + + +def _find_free_port(): + """Find a free port on localhost.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + s.listen(1) + port = s.getsockname()[1] + return port + + +def _context_parallel_worker(rank, world_size, master_port, model_class, init_dict, cp_dict, inputs_dict, return_dict): + """Worker function for context parallel testing.""" + try: + # Set up distributed environment + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(master_port) + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world_size) + + # Initialize process group + dist.init_process_group(backend="nccl", rank=rank, world_size=world_size) + + # Set device for this process + torch.cuda.set_device(rank) + device = torch.device(f"cuda:{rank}") + + # Create model + model = model_class(**init_dict) + model.to(device) + model.eval() + + # Move inputs to device + inputs_on_device = {} + for key, value in inputs_dict.items(): + if isinstance(value, torch.Tensor): + inputs_on_device[key] = value.to(device) + else: + inputs_on_device[key] = value + + # Enable context parallelism + cp_config = ContextParallelConfig(**cp_dict) + model.enable_parallelism(config=cp_config) + + # Run forward pass + with torch.no_grad(): + output = model(**inputs_on_device, return_dict=False)[0] + + # Only rank 0 reports results + if rank == 0: + return_dict["status"] = "success" + return_dict["output_shape"] = list(output.shape) + + except Exception as e: + if rank == 0: + return_dict["status"] = "error" + return_dict["error"] = str(e) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +@is_context_parallel +@require_torch_multi_accelerator +class ContextParallelTesterMixin: + @pytest.mark.parametrize("cp_type", ["ulysses_degree", "ring_degree"], ids=["ulysses", "ring"]) + def test_context_parallel_inference(self, cp_type): + if not torch.distributed.is_available(): + pytest.skip("torch.distributed is not available.") + + if not hasattr(self.model_class, "_cp_plan") or self.model_class._cp_plan is None: + pytest.skip("Model does not have a _cp_plan defined for context parallel inference.") + + world_size = 2 + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + + # Move all tensors to CPU for multiprocessing + inputs_dict = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()} + cp_dict = {cp_type: world_size} + + # Find a free port for distributed communication + master_port = _find_free_port() + + # Use multiprocessing manager for cross-process communication + manager = mp.Manager() + return_dict = manager.dict() + + # Spawn worker processes + mp.spawn( + _context_parallel_worker, + args=(world_size, master_port, self.model_class, init_dict, cp_dict, inputs_dict, return_dict), + nprocs=world_size, + join=True, + ) + + assert return_dict.get("status") == "success", ( + f"Context parallel inference failed: {return_dict.get('error', 'Unknown error')}" + ) diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py new file mode 100644 index 000000000000..0f1fbde72485 --- /dev/null +++ b/tests/models/testing_utils/quantization.py @@ -0,0 +1,1374 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import gc + +import pytest +import torch + +from diffusers import BitsAndBytesConfig, GGUFQuantizationConfig, NVIDIAModelOptConfig, QuantoConfig, TorchAoConfig +from diffusers.utils.import_utils import ( + is_bitsandbytes_available, + is_gguf_available, + is_nvidia_modelopt_available, + is_optimum_quanto_available, + is_torchao_available, + is_torchao_version, +) + +from ...testing_utils import ( + backend_empty_cache, + backend_max_memory_allocated, + backend_reset_peak_memory_stats, + is_bitsandbytes, + is_gguf, + is_modelopt, + is_quantization, + is_quanto, + is_torch_compile, + is_torchao, + require_accelerate, + require_accelerator, + require_bitsandbytes_version_greater, + require_gguf_version_greater_or_equal, + require_modelopt_version_greater_or_equal, + require_quanto, + require_torchao_version_greater_or_equal, + torch_device, +) + + +if is_nvidia_modelopt_available(): + import modelopt.torch.quantization as mtq + +if is_bitsandbytes_available(): + import bitsandbytes as bnb + +if is_optimum_quanto_available(): + from optimum.quanto import QLinear + +if is_gguf_available(): + pass + +if is_torchao_available(): + if is_torchao_version(">=", "0.9.0"): + pass + + +class LoRALayer(torch.nn.Module): + """Wraps a linear layer with LoRA-like adapter - Used for testing purposes only. + + Taken from + https://github.com/huggingface/transformers/blob/566302686a71de14125717dea9a6a45b24d42b37/tests/quantization/bnb/test_4bit.py#L62C5-L78C77 + """ + + def __init__(self, module: torch.nn.Module, rank: int): + super().__init__() + self.module = module + self.adapter = torch.nn.Sequential( + torch.nn.Linear(module.in_features, rank, bias=False), + torch.nn.Linear(rank, module.out_features, bias=False), + ) + small_std = (2.0 / (5 * min(module.in_features, module.out_features))) ** 0.5 + torch.nn.init.normal_(self.adapter[0].weight, std=small_std) + torch.nn.init.zeros_(self.adapter[1].weight) + self.adapter.to(module.weight.device) + + def forward(self, input, *args, **kwargs): + return self.module(input, *args, **kwargs) + self.adapter(input) + + +@is_quantization +@require_accelerator +class QuantizationTesterMixin: + """ + Base mixin class providing common test implementations for quantization testing. + + Backend-specific mixins should: + 1. Implement _create_quantized_model(config_kwargs) + 2. Implement _verify_if_layer_quantized(name, module, config_kwargs) + 3. Define their config dict (e.g., BNB_CONFIGS, QUANTO_WEIGHT_TYPES, etc.) + 4. Use @pytest.mark.parametrize to create tests that call the common test methods below + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained (e.g., {"subfolder": "transformer"}) + + Expected methods in test classes: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + """ + + def setup_method(self): + gc.collect() + backend_empty_cache(torch_device) + + def teardown_method(self): + gc.collect() + backend_empty_cache(torch_device) + + def _create_quantized_model(self, config_kwargs, **extra_kwargs): + """ + Create a quantized model with the given config kwargs. + + Args: + config_kwargs: Quantization config parameters + **extra_kwargs: Additional kwargs to pass to from_pretrained (e.g., device_map, offload_folder) + """ + raise NotImplementedError("Subclass must implement _create_quantized_model") + + def _verify_if_layer_quantized(self, name, module, config_kwargs): + raise NotImplementedError("Subclass must implement _verify_if_layer_quantized") + + def _is_module_quantized(self, module): + """ + Check if a module is quantized. Returns True if quantized, False otherwise. + Default implementation tries _verify_if_layer_quantized and catches exceptions. + Subclasses can override for more efficient checking. + """ + try: + self._verify_if_layer_quantized("", module, {}) + return True + except (AssertionError, AttributeError): + return False + + def _load_unquantized_model(self): + kwargs = getattr(self, "pretrained_model_kwargs", {}) + return self.model_class.from_pretrained(self.pretrained_model_name_or_path, **kwargs) + + def _test_quantization_num_parameters(self, config_kwargs): + model = self._load_unquantized_model() + num_params = model.num_parameters() + + model_quantized = self._create_quantized_model(config_kwargs) + num_params_quantized = model_quantized.num_parameters() + + assert num_params == num_params_quantized, ( + f"Parameter count mismatch: unquantized={num_params}, quantized={num_params_quantized}" + ) + + def _test_quantization_memory_footprint(self, config_kwargs, expected_memory_reduction=1.2): + model = self._load_unquantized_model() + mem = model.get_memory_footprint() + + model_quantized = self._create_quantized_model(config_kwargs) + mem_quantized = model_quantized.get_memory_footprint() + + ratio = mem / mem_quantized + assert ratio >= expected_memory_reduction, ( + f"Memory ratio {ratio:.2f} is less than expected ({expected_memory_reduction}x). unquantized={mem}, quantized={mem_quantized}" + ) + + @torch.no_grad() + def _test_quantization_inference(self, config_kwargs): + model_quantized = self._create_quantized_model(config_kwargs) + model_quantized.to(torch_device) + + inputs = self.get_dummy_inputs() + output = model_quantized(**inputs, return_dict=False)[0] + + assert output is not None, "Model output is None" + assert not torch.isnan(output).any(), "Model output contains NaN" + + def _test_quantization_dtype_assignment(self, config_kwargs): + model = self._create_quantized_model(config_kwargs) + + with pytest.raises(ValueError): + model.to(torch.float16) + + with pytest.raises(ValueError): + device_0 = f"{torch_device}:0" + model.to(device=device_0, dtype=torch.float16) + + with pytest.raises(ValueError): + model.float() + + with pytest.raises(ValueError): + model.half() + + model.to(torch_device) + + @torch.no_grad() + def _test_quantization_lora_inference(self, config_kwargs): + try: + from peft import LoraConfig + except ImportError: + pytest.skip("peft is not available") + + from diffusers.loaders.peft import PeftAdapterMixin + + if not issubclass(self.model_class, PeftAdapterMixin): + pytest.skip(f"PEFT is not supported for this model ({self.model_class.__name__})") + + model = self._create_quantized_model(config_kwargs) + + lora_config = LoraConfig( + r=4, + lora_alpha=4, + target_modules=["to_q", "to_k", "to_v", "to_out.0"], + init_lora_weights=False, + ) + model.add_adapter(lora_config) + # Move LoRA adapter weights to device (they default to CPU) + model.to(torch_device) + + inputs = self.get_dummy_inputs() + output = model(**inputs, return_dict=False)[0] + + assert output is not None, "Model output is None with LoRA" + assert not torch.isnan(output).any(), "Model output contains NaN with LoRA" + + @torch.no_grad() + def _test_quantization_serialization(self, config_kwargs, tmp_path): + model = self._create_quantized_model(config_kwargs) + + model.save_pretrained(str(tmp_path), safe_serialization=True) + + model_loaded = self.model_class.from_pretrained(str(tmp_path)) + + inputs = self.get_dummy_inputs() + output = model_loaded(**inputs, return_dict=False)[0] + assert not torch.isnan(output).any(), "Loaded model output contains NaN" + + def _test_quantized_layers(self, config_kwargs): + model_fp = self._load_unquantized_model() + num_linear_layers = sum(1 for module in model_fp.modules() if isinstance(module, torch.nn.Linear)) + + model_quantized = self._create_quantized_model(config_kwargs) + + num_fp32_modules = 0 + if hasattr(model_quantized, "_keep_in_fp32_modules") and model_quantized._keep_in_fp32_modules: + for name, module in model_quantized.named_modules(): + if isinstance(module, torch.nn.Linear): + if any(fp32_name in name for fp32_name in model_quantized._keep_in_fp32_modules): + num_fp32_modules += 1 + + expected_quantized_layers = num_linear_layers - num_fp32_modules + + num_quantized_layers = 0 + for name, module in model_quantized.named_modules(): + if isinstance(module, torch.nn.Linear): + if hasattr(model_quantized, "_keep_in_fp32_modules") and model_quantized._keep_in_fp32_modules: + if any(fp32_name in name for fp32_name in model_quantized._keep_in_fp32_modules): + continue + self._verify_if_layer_quantized(name, module, config_kwargs) + num_quantized_layers += 1 + + assert num_quantized_layers > 0, ( + f"No quantized layers found in model (expected {expected_quantized_layers} linear layers, {num_fp32_modules} kept in FP32)" + ) + assert num_quantized_layers == expected_quantized_layers, ( + f"Quantized layer count mismatch: expected {expected_quantized_layers}, got {num_quantized_layers} (total linear layers: {num_linear_layers}, FP32 modules: {num_fp32_modules})" + ) + + def _test_quantization_modules_to_not_convert(self, config_kwargs, modules_to_not_convert): + """ + Test that modules specified in modules_to_not_convert are not quantized. + + Args: + config_kwargs: Base quantization config kwargs + modules_to_not_convert: List of module names to exclude from quantization + """ + # Create config with modules_to_not_convert + config_kwargs_with_exclusion = config_kwargs.copy() + config_kwargs_with_exclusion["modules_to_not_convert"] = modules_to_not_convert + + model_with_exclusion = self._create_quantized_model(config_kwargs_with_exclusion) + + # Find a module that should NOT be quantized + found_excluded = False + for name, module in model_with_exclusion.named_modules(): + if isinstance(module, torch.nn.Linear): + # Check if this module is in the exclusion list + if any(excluded in name for excluded in modules_to_not_convert): + found_excluded = True + # This module should NOT be quantized + assert not self._is_module_quantized(module), ( + f"Module {name} should not be quantized but was found to be quantized" + ) + + assert found_excluded, f"No linear layers found in excluded modules: {modules_to_not_convert}" + + # Find a module that SHOULD be quantized (not in exclusion list) + found_quantized = False + for name, module in model_with_exclusion.named_modules(): + if isinstance(module, torch.nn.Linear): + # Check if this module is NOT in the exclusion list + if not any(excluded in name for excluded in modules_to_not_convert): + if self._is_module_quantized(module): + found_quantized = True + break + + assert found_quantized, "No quantized layers found outside of excluded modules" + + # Compare memory footprint with fully quantized model + model_fully_quantized = self._create_quantized_model(config_kwargs) + + mem_with_exclusion = model_with_exclusion.get_memory_footprint() + mem_fully_quantized = model_fully_quantized.get_memory_footprint() + + assert mem_with_exclusion > mem_fully_quantized, ( + f"Model with exclusions should be larger. With exclusion: {mem_with_exclusion}, fully quantized: {mem_fully_quantized}" + ) + + @torch.no_grad() + def _test_quantization_device_map(self, config_kwargs): + """ + Test that quantized models work correctly with device_map="auto". + + Args: + config_kwargs: Base quantization config kwargs + """ + model = self._create_quantized_model(config_kwargs, device_map="auto") + + assert hasattr(model, "hf_device_map"), "Model should have hf_device_map attribute" + assert model.hf_device_map is not None, "hf_device_map should not be None" + + inputs = self.get_dummy_inputs() + output = model(**inputs, return_dict=False)[0] + assert output is not None, "Model output is None" + assert not torch.isnan(output).any(), "Model output contains NaN" + + @torch.no_grad() + def _test_dequantize(self, config_kwargs): + """ + Test that dequantize() converts quantized model back to standard linear layers. + + Args: + config_kwargs: Quantization config parameters + """ + model = self._create_quantized_model(config_kwargs) + model.to(torch_device) + + if not hasattr(model, "dequantize"): + pytest.skip("Model does not have dequantize method") + + model.dequantize() + + for name, module in model.named_modules(): + if isinstance(module, torch.nn.Linear): + assert not self._is_module_quantized(module), f"Module {name} is still quantized after dequantize()" + + # Get model dtype from first parameter + model_dtype = next(model.parameters()).dtype + + inputs = self.get_dummy_inputs() + # Cast inputs to model dtype + inputs = { + k: v.to(model_dtype) if isinstance(v, torch.Tensor) and v.is_floating_point() else v + for k, v in inputs.items() + } + output = model(**inputs, return_dict=False)[0] + assert output is not None, "Model output is None after dequantization" + assert not torch.isnan(output).any(), "Model output contains NaN after dequantization" + + def _test_quantization_training(self, config_kwargs): + """ + Test that quantized models can be used for training with LoRA-like adapters. + + This test: + 1. Freezes all model parameters + 2. Casts small parameters (e.g., layernorm) to fp32 for stability + 3. Adds LoRA adapters to attention layers + 4. Runs forward and backward passes + 5. Verifies gradients are computed correctly + + Args: + config_kwargs: Quantization config parameters + """ + model = self._create_quantized_model(config_kwargs) + + # Step 1: freeze all parameters + for param in model.parameters(): + param.requires_grad = False + if param.ndim == 1: + # cast small parameters (e.g. layernorm) to fp32 for stability + param.data = param.data.to(torch.float32) + + # Step 2: add adapters to attention layers + adapter_count = 0 + for _, module in model.named_modules(): + if "Attention" in repr(type(module)): + if hasattr(module, "to_k"): + module.to_k = LoRALayer(module.to_k, rank=4) + adapter_count += 1 + if hasattr(module, "to_q"): + module.to_q = LoRALayer(module.to_q, rank=4) + adapter_count += 1 + if hasattr(module, "to_v"): + module.to_v = LoRALayer(module.to_v, rank=4) + adapter_count += 1 + + if adapter_count == 0: + pytest.skip("No attention layers found in model for adapter training test") + + # Step 3: run forward and backward pass + inputs = self.get_dummy_inputs() + + with torch.amp.autocast(torch_device, dtype=torch.float16): + out = model(**inputs, return_dict=False)[0] + out.norm().backward() + + # Step 4: verify gradients are computed + for module in model.modules(): + if isinstance(module, LoRALayer): + assert module.adapter[1].weight.grad is not None, "LoRA adapter gradient is None" + assert module.adapter[1].weight.grad.norm().item() > 0, "LoRA adapter gradient norm is zero" + + +@is_quantization +@is_bitsandbytes +@require_accelerator +@require_bitsandbytes_version_greater("0.43.2") +@require_accelerate +class BitsAndBytesConfigMixin: + """ + Base mixin providing BitsAndBytes quantization config and model creation. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained + """ + + BNB_CONFIGS = { + "4bit_nf4": { + "load_in_4bit": True, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_compute_dtype": torch.float16, + }, + "4bit_fp4": { + "load_in_4bit": True, + "bnb_4bit_quant_type": "fp4", + "bnb_4bit_compute_dtype": torch.float16, + }, + "8bit": { + "load_in_8bit": True, + }, + } + + BNB_EXPECTED_MEMORY_REDUCTIONS = { + "4bit_nf4": 3.0, + "4bit_fp4": 3.0, + "8bit": 1.5, + } + + def _create_quantized_model(self, config_kwargs, **extra_kwargs): + config = BitsAndBytesConfig(**config_kwargs) + kwargs = getattr(self, "pretrained_model_kwargs", {}).copy() + kwargs["quantization_config"] = config + kwargs.update(extra_kwargs) + return self.model_class.from_pretrained(self.pretrained_model_name_or_path, **kwargs) + + def _verify_if_layer_quantized(self, name, module, config_kwargs): + expected_weight_class = bnb.nn.Params4bit if config_kwargs.get("load_in_4bit") else bnb.nn.Int8Params + assert module.weight.__class__ == expected_weight_class, ( + f"Layer {name} has weight type {module.weight.__class__}, expected {expected_weight_class}" + ) + + +@is_bitsandbytes +@require_accelerator +@require_bitsandbytes_version_greater("0.43.2") +@require_accelerate +class BitsAndBytesTesterMixin(BitsAndBytesConfigMixin, QuantizationTesterMixin): + """ + Mixin class for testing BitsAndBytes quantization on models. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained (e.g., {"subfolder": "transformer"}) + + Expected methods to be implemented by subclasses: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Optional class attributes: + - BNB_CONFIGS: Dict of config name -> BitsAndBytesConfig kwargs to test + + Pytest mark: bitsandbytes + Use `pytest -m "not bitsandbytes"` to skip these tests + """ + + @pytest.mark.parametrize( + "config_name", + list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), + ids=list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), + ) + def test_bnb_quantization_num_parameters(self, config_name): + self._test_quantization_num_parameters(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name]) + + @pytest.mark.parametrize( + "config_name", + list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), + ids=list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), + ) + def test_bnb_quantization_memory_footprint(self, config_name): + expected = BitsAndBytesConfigMixin.BNB_EXPECTED_MEMORY_REDUCTIONS.get(config_name, 1.2) + self._test_quantization_memory_footprint( + BitsAndBytesConfigMixin.BNB_CONFIGS[config_name], expected_memory_reduction=expected + ) + + @pytest.mark.parametrize( + "config_name", + list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), + ids=list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), + ) + def test_bnb_quantization_inference(self, config_name): + self._test_quantization_inference(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", ["4bit_nf4"], ids=["4bit_nf4"]) + def test_bnb_quantization_dtype_assignment(self, config_name): + self._test_quantization_dtype_assignment(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", ["4bit_nf4"], ids=["4bit_nf4"]) + def test_bnb_quantization_lora_inference(self, config_name): + self._test_quantization_lora_inference(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", ["4bit_nf4"], ids=["4bit_nf4"]) + def test_bnb_quantization_serialization(self, config_name, tmp_path): + self._test_quantization_serialization(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name], tmp_path) + + @pytest.mark.parametrize( + "config_name", + list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), + ids=list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), + ) + def test_bnb_quantized_layers(self, config_name): + self._test_quantized_layers(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name]) + + @pytest.mark.parametrize( + "config_name", + list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), + ids=list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), + ) + def test_bnb_quantization_config_serialization(self, config_name): + model = self._create_quantized_model(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name]) + + assert "quantization_config" in model.config, "Missing quantization_config" + _ = model.config["quantization_config"].to_dict() + _ = model.config["quantization_config"].to_diff_dict() + _ = model.config["quantization_config"].to_json_string() + + def test_bnb_original_dtype(self): + config_name = list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys())[0] + config_kwargs = BitsAndBytesConfigMixin.BNB_CONFIGS[config_name] + + model = self._create_quantized_model(config_kwargs) + + assert "_pre_quantization_dtype" in model.config, "Missing _pre_quantization_dtype" + assert model.config["_pre_quantization_dtype"] in [ + torch.float16, + torch.float32, + torch.bfloat16, + ], f"Unexpected dtype: {model.config['_pre_quantization_dtype']}" + + @torch.no_grad() + def test_bnb_keep_modules_in_fp32(self): + if not hasattr(self.model_class, "_keep_in_fp32_modules"): + pytest.skip(f"{self.model_class.__name__} does not have _keep_in_fp32_modules") + + config_kwargs = BitsAndBytesConfigMixin.BNB_CONFIGS["4bit_nf4"] + + original_fp32_modules = getattr(self.model_class, "_keep_in_fp32_modules", None) + self.model_class._keep_in_fp32_modules = ["proj_out"] + + try: + model = self._create_quantized_model(config_kwargs) + + for name, module in model.named_modules(): + if isinstance(module, torch.nn.Linear): + if any(fp32_name in name for fp32_name in model._keep_in_fp32_modules): + assert module.weight.dtype == torch.float32, ( + f"Module {name} should be FP32 but is {module.weight.dtype}" + ) + else: + assert module.weight.dtype == torch.uint8, ( + f"Module {name} should be uint8 but is {module.weight.dtype}" + ) + + inputs = self.get_dummy_inputs() + _ = model(**inputs) + finally: + if original_fp32_modules is not None: + self.model_class._keep_in_fp32_modules = original_fp32_modules + + def test_bnb_modules_to_not_convert(self): + """Test that modules_to_not_convert parameter works correctly.""" + modules_to_exclude = getattr(self, "modules_to_not_convert_for_test", None) + if modules_to_exclude is None: + pytest.skip("modules_to_not_convert_for_test not defined for this model") + + self._test_quantization_modules_to_not_convert( + BitsAndBytesConfigMixin.BNB_CONFIGS["4bit_nf4"], modules_to_exclude + ) + + @pytest.mark.parametrize("config_name", ["4bit_nf4", "8bit"], ids=["4bit_nf4", "8bit"]) + def test_bnb_device_map(self, config_name): + """Test that device_map='auto' works correctly with quantization.""" + self._test_quantization_device_map(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name]) + + def test_bnb_dequantize(self): + """Test that dequantize() works correctly.""" + self._test_dequantize(BitsAndBytesConfigMixin.BNB_CONFIGS["4bit_nf4"]) + + def test_bnb_training(self): + """Test that quantized models can be used for training with adapters.""" + self._test_quantization_training(BitsAndBytesConfigMixin.BNB_CONFIGS["4bit_nf4"]) + + @pytest.mark.parametrize( + "config_name", + list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), + ids=list(BitsAndBytesConfigMixin.BNB_CONFIGS.keys()), + ) + def test_cpu_device_map(self, config_name): + config_kwargs = BitsAndBytesConfigMixin.BNB_CONFIGS[config_name] + model_quantized = self._create_quantized_model(config_kwargs, device_map="cpu") + + assert hasattr(model_quantized, "hf_device_map"), "Model should have hf_device_map attribute" + assert model_quantized.hf_device_map is not None, "hf_device_map should not be None" + assert model_quantized.device == torch.device("cpu"), ( + f"Model should be on CPU, but is on {model_quantized.device}" + ) + + +@is_quantization +@is_quanto +@require_quanto +@require_accelerate +@require_accelerator +class QuantoConfigMixin: + """ + Base mixin providing Quanto quantization config and model creation. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained + """ + + QUANTO_WEIGHT_TYPES = { + "float8": {"weights_dtype": "float8"}, + "int8": {"weights_dtype": "int8"}, + "int4": {"weights_dtype": "int4"}, + "int2": {"weights_dtype": "int2"}, + } + + QUANTO_EXPECTED_MEMORY_REDUCTIONS = { + "float8": 1.5, + "int8": 1.5, + "int4": 3.0, + "int2": 7.0, + } + + def _create_quantized_model(self, config_kwargs, **extra_kwargs): + config = QuantoConfig(**config_kwargs) + kwargs = getattr(self, "pretrained_model_kwargs", {}).copy() + kwargs["quantization_config"] = config + kwargs.update(extra_kwargs) + return self.model_class.from_pretrained(self.pretrained_model_name_or_path, **kwargs) + + def _verify_if_layer_quantized(self, name, module, config_kwargs): + assert isinstance(module, QLinear), f"Layer {name} is not QLinear, got {type(module)}" + + def _test_quantization_memory_footprint(self, config_kwargs, expected_memory_reduction=1.2): + """Override to use max_memory_allocated for Quanto (get_memory_footprint doesn't reflect quantized _data).""" + # Measure unquantized model memory + backend_reset_peak_memory_stats(torch_device) + backend_empty_cache(torch_device) + + model = self._load_unquantized_model() + model.to(torch_device) + mem = backend_max_memory_allocated(torch_device) + + del model + gc.collect() + backend_empty_cache(torch_device) + + # Measure quantized model memory + backend_reset_peak_memory_stats(torch_device) + + model_quantized = self._create_quantized_model(config_kwargs) + model_quantized.to(torch_device) + mem_quantized = backend_max_memory_allocated(torch_device) + + ratio = mem / mem_quantized + assert ratio >= expected_memory_reduction, ( + f"Memory ratio {ratio:.2f} is less than expected ({expected_memory_reduction}x). unquantized={mem}, quantized={mem_quantized}" + ) + + +@is_quanto +@require_quanto +@require_accelerate +@require_accelerator +class QuantoTesterMixin(QuantoConfigMixin, QuantizationTesterMixin): + """ + Mixin class for testing Quanto quantization on models. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained (e.g., {"subfolder": "transformer"}) + + Expected methods to be implemented by subclasses: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Optional class attributes: + - QUANTO_WEIGHT_TYPES: Dict of weight_type_name -> qtype + + Pytest mark: quanto + Use `pytest -m "not quanto"` to skip these tests + """ + + @pytest.mark.parametrize( + "weight_type_name", + list(QuantoConfigMixin.QUANTO_WEIGHT_TYPES.keys()), + ids=list(QuantoConfigMixin.QUANTO_WEIGHT_TYPES.keys()), + ) + def test_quanto_quantization_num_parameters(self, weight_type_name): + self._test_quantization_num_parameters(QuantoConfigMixin.QUANTO_WEIGHT_TYPES[weight_type_name]) + + @pytest.mark.parametrize( + "weight_type_name", + list(QuantoConfigMixin.QUANTO_WEIGHT_TYPES.keys()), + ids=list(QuantoConfigMixin.QUANTO_WEIGHT_TYPES.keys()), + ) + def test_quanto_quantization_memory_footprint(self, weight_type_name): + expected = QuantoConfigMixin.QUANTO_EXPECTED_MEMORY_REDUCTIONS.get(weight_type_name, 1.2) + self._test_quantization_memory_footprint( + QuantoConfigMixin.QUANTO_WEIGHT_TYPES[weight_type_name], expected_memory_reduction=expected + ) + + @pytest.mark.parametrize( + "weight_type_name", + list(QuantoConfigMixin.QUANTO_WEIGHT_TYPES.keys()), + ids=list(QuantoConfigMixin.QUANTO_WEIGHT_TYPES.keys()), + ) + def test_quanto_quantization_inference(self, weight_type_name): + self._test_quantization_inference(QuantoConfigMixin.QUANTO_WEIGHT_TYPES[weight_type_name]) + + @pytest.mark.parametrize("weight_type_name", ["int8"], ids=["int8"]) + def test_quanto_quantized_layers(self, weight_type_name): + self._test_quantized_layers(QuantoConfigMixin.QUANTO_WEIGHT_TYPES[weight_type_name]) + + @pytest.mark.parametrize("weight_type_name", ["int8"], ids=["int8"]) + def test_quanto_quantization_lora_inference(self, weight_type_name): + self._test_quantization_lora_inference(QuantoConfigMixin.QUANTO_WEIGHT_TYPES[weight_type_name]) + + @pytest.mark.parametrize("weight_type_name", ["int8"], ids=["int8"]) + def test_quanto_quantization_serialization(self, weight_type_name, tmp_path): + self._test_quantization_serialization(QuantoConfigMixin.QUANTO_WEIGHT_TYPES[weight_type_name], tmp_path) + + def test_quanto_modules_to_not_convert(self): + """Test that modules_to_not_convert parameter works correctly.""" + modules_to_exclude = getattr(self, "modules_to_not_convert_for_test", None) + if modules_to_exclude is None: + pytest.skip("modules_to_not_convert_for_test not defined for this model") + + self._test_quantization_modules_to_not_convert( + QuantoConfigMixin.QUANTO_WEIGHT_TYPES["int8"], modules_to_exclude + ) + + def test_quanto_device_map(self): + """Test that device_map='auto' works correctly with quantization.""" + self._test_quantization_device_map(QuantoConfigMixin.QUANTO_WEIGHT_TYPES["int8"]) + + def test_quanto_dequantize(self): + """Test that dequantize() works correctly.""" + self._test_dequantize(QuantoConfigMixin.QUANTO_WEIGHT_TYPES["int8"]) + + +@is_quantization +@is_torchao +@require_accelerator +@require_torchao_version_greater_or_equal("0.7.0") +class TorchAoConfigMixin: + """ + Base mixin providing TorchAO quantization config and model creation. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained + """ + + TORCHAO_QUANT_TYPES = { + "int4wo": {"quant_type": "int4_weight_only"}, + "int8wo": {"quant_type": "int8_weight_only"}, + "int8dq": {"quant_type": "int8_dynamic_activation_int8_weight"}, + } + + TORCHAO_EXPECTED_MEMORY_REDUCTIONS = { + "int4wo": 1.8, + "int8wo": 1.5, + "int8dq": 1.5, + } + + def _create_quantized_model(self, config_kwargs, **extra_kwargs): + config = TorchAoConfig(**config_kwargs) + kwargs = getattr(self, "pretrained_model_kwargs", {}).copy() + kwargs["quantization_config"] = config + kwargs["device_map"] = str(torch_device) + kwargs.update(extra_kwargs) + return self.model_class.from_pretrained(self.pretrained_model_name_or_path, **kwargs) + + def _verify_if_layer_quantized(self, name, module, config_kwargs): + assert isinstance(module, torch.nn.Linear), f"Layer {name} is not Linear, got {type(module)}" + + +# int4wo requires CUDA-specific ops (_convert_weight_to_int4pack) +_int4wo_skip = pytest.mark.skipif(torch_device != "cuda", reason="int4wo quantization requires CUDA") + + +@is_torchao +@require_accelerator +@require_torchao_version_greater_or_equal("0.7.0") +class TorchAoTesterMixin(TorchAoConfigMixin, QuantizationTesterMixin): + """ + Mixin class for testing TorchAO quantization on models. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained (e.g., {"subfolder": "transformer"}) + + Expected methods to be implemented by subclasses: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Optional class attributes: + - TORCHAO_QUANT_TYPES: Dict of quantization type strings to test + + Pytest mark: torchao + Use `pytest -m "not torchao"` to skip these tests + """ + + @pytest.mark.parametrize( + "quant_type", + [ + pytest.param("int4wo", marks=_int4wo_skip), + "int8wo", + "int8dq", + ], + ids=["int4wo", "int8wo", "int8dq"], + ) + def test_torchao_quantization_num_parameters(self, quant_type): + self._test_quantization_num_parameters(TorchAoConfigMixin.TORCHAO_QUANT_TYPES[quant_type]) + + @pytest.mark.parametrize( + "quant_type", + [ + pytest.param("int4wo", marks=_int4wo_skip), + "int8wo", + "int8dq", + ], + ids=["int4wo", "int8wo", "int8dq"], + ) + def test_torchao_quantization_memory_footprint(self, quant_type): + expected = TorchAoConfigMixin.TORCHAO_EXPECTED_MEMORY_REDUCTIONS.get(quant_type, 1.2) + self._test_quantization_memory_footprint( + TorchAoConfigMixin.TORCHAO_QUANT_TYPES[quant_type], expected_memory_reduction=expected + ) + + @pytest.mark.parametrize( + "quant_type", + [ + pytest.param("int4wo", marks=_int4wo_skip), + "int8wo", + "int8dq", + ], + ids=["int4wo", "int8wo", "int8dq"], + ) + def test_torchao_quantization_inference(self, quant_type): + self._test_quantization_inference(TorchAoConfigMixin.TORCHAO_QUANT_TYPES[quant_type]) + + @pytest.mark.parametrize("quant_type", ["int8wo"], ids=["int8wo"]) + def test_torchao_quantized_layers(self, quant_type): + self._test_quantized_layers(TorchAoConfigMixin.TORCHAO_QUANT_TYPES[quant_type]) + + @pytest.mark.parametrize("quant_type", ["int8wo"], ids=["int8wo"]) + def test_torchao_quantization_lora_inference(self, quant_type): + self._test_quantization_lora_inference(TorchAoConfigMixin.TORCHAO_QUANT_TYPES[quant_type]) + + @pytest.mark.parametrize("quant_type", ["int8wo"], ids=["int8wo"]) + def test_torchao_quantization_serialization(self, quant_type, tmp_path): + """Override to use safe_serialization=False for TorchAO (safetensors not supported).""" + config_kwargs = TorchAoConfigMixin.TORCHAO_QUANT_TYPES[quant_type] + model = self._create_quantized_model(config_kwargs) + + model.save_pretrained(str(tmp_path), safe_serialization=False) + + model_loaded = self.model_class.from_pretrained(str(tmp_path), device_map=str(torch_device)) + + inputs = self.get_dummy_inputs() + output = model_loaded(**inputs, return_dict=False)[0] + assert not torch.isnan(output).any(), "Loaded model output contains NaN" + + def test_torchao_modules_to_not_convert(self): + """Test that modules_to_not_convert parameter works correctly.""" + modules_to_exclude = getattr(self, "modules_to_not_convert_for_test", None) + if modules_to_exclude is None: + pytest.skip("modules_to_not_convert_for_test not defined for this model") + + self._test_quantization_modules_to_not_convert( + TorchAoConfigMixin.TORCHAO_QUANT_TYPES["int8wo"], modules_to_exclude + ) + + def test_torchao_device_map(self): + """Test that device_map='auto' works correctly with quantization.""" + self._test_quantization_device_map(TorchAoConfigMixin.TORCHAO_QUANT_TYPES["int8wo"]) + + def test_torchao_dequantize(self): + """Test that dequantize() works correctly.""" + self._test_dequantize(TorchAoConfigMixin.TORCHAO_QUANT_TYPES["int8wo"]) + + def test_torchao_training(self): + """Test that quantized models can be used for training with adapters.""" + self._test_quantization_training(TorchAoConfigMixin.TORCHAO_QUANT_TYPES["int8wo"]) + + +@is_quantization +@is_gguf +@require_accelerate +@require_accelerator +@require_gguf_version_greater_or_equal("0.10.0") +class GGUFConfigMixin: + """ + Base mixin providing GGUF quantization config and model creation. + + Expected from config mixin: + - model_class: The model class to test + + Required properties (must be implemented by subclasses): + - gguf_filename: URL or path to the GGUF file + """ + + @property + def gguf_filename(self): + """URL or path to the GGUF file. Must be implemented by subclasses.""" + raise NotImplementedError("Subclasses must implement the `gguf_filename` property.") + + def _create_quantized_model(self, config_kwargs=None, **extra_kwargs): + if config_kwargs is None: + config_kwargs = {"compute_dtype": torch.bfloat16} + + config = GGUFQuantizationConfig(**config_kwargs) + kwargs = { + "quantization_config": config, + "torch_dtype": config_kwargs.get("compute_dtype", torch.bfloat16), + "device_map": str(torch_device), + } + kwargs.update(extra_kwargs) + return self.model_class.from_single_file(self.gguf_filename, **kwargs) + + def _verify_if_layer_quantized(self, name, module, config_kwargs=None): + from diffusers.quantizers.gguf.utils import GGUFParameter + + assert isinstance(module.weight, GGUFParameter), f"{name} weight is not GGUFParameter" + assert hasattr(module.weight, "quant_type"), f"{name} weight missing quant_type" + assert module.weight.dtype == torch.uint8, f"{name} weight dtype should be uint8" + + +@is_gguf +@require_accelerate +@require_accelerator +@require_gguf_version_greater_or_equal("0.10.0") +class GGUFTesterMixin(GGUFConfigMixin, QuantizationTesterMixin): + """ + Mixin class for testing GGUF quantization on models. + + Expected from config mixin: + - model_class: The model class to test + + Required properties (must be implemented by subclasses): + - gguf_filename: URL or path to the GGUF file + + Expected methods from config mixin: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest mark: gguf + Use `pytest -m "not gguf"` to skip these tests + """ + + def test_gguf_quantization_inference(self): + self._test_quantization_inference({"compute_dtype": torch.bfloat16}) + + def test_gguf_keep_modules_in_fp32(self): + if not hasattr(self.model_class, "_keep_in_fp32_modules"): + pytest.skip(f"{self.model_class.__name__} does not have _keep_in_fp32_modules") + + _keep_in_fp32_modules = self.model_class._keep_in_fp32_modules + self.model_class._keep_in_fp32_modules = ["proj_out"] + + try: + model = self._create_quantized_model() + for name, module in model.named_modules(): + if isinstance(module, torch.nn.Linear) and name in model._keep_in_fp32_modules: + assert module.weight.dtype == torch.float32, f"Module {name} should be FP32" + finally: + self.model_class._keep_in_fp32_modules = _keep_in_fp32_modules + + def test_gguf_quantization_dtype_assignment(self): + self._test_quantization_dtype_assignment({"compute_dtype": torch.bfloat16}) + + def test_gguf_quantization_lora_inference(self): + self._test_quantization_lora_inference({"compute_dtype": torch.bfloat16}) + + def test_gguf_dequantize(self): + """Test that dequantize() works correctly.""" + self._test_dequantize({"compute_dtype": torch.bfloat16}) + + +@is_quantization +@is_modelopt +@require_accelerator +@require_accelerate +@require_modelopt_version_greater_or_equal("0.33.1") +class ModelOptConfigMixin: + """ + Base mixin providing NVIDIA ModelOpt quantization config and model creation. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained + """ + + MODELOPT_CONFIGS = { + "fp8": {"quant_type": "FP8"}, + "int8": {"quant_type": "INT8"}, + "int4": {"quant_type": "INT4"}, + } + + MODELOPT_EXPECTED_MEMORY_REDUCTIONS = { + "fp8": 1.5, + "int8": 1.5, + "int4": 3.0, + } + + def _create_quantized_model(self, config_kwargs, **extra_kwargs): + config = NVIDIAModelOptConfig(**config_kwargs) + kwargs = getattr(self, "pretrained_model_kwargs", {}).copy() + kwargs["quantization_config"] = config + kwargs["device_map"] = str(torch_device) + kwargs.update(extra_kwargs) + return self.model_class.from_pretrained(self.pretrained_model_name_or_path, **kwargs) + + def _verify_if_layer_quantized(self, name, module, config_kwargs): + assert mtq.utils.is_quantized(module), f"Layer {name} does not have weight_quantizer attribute (not quantized)" + + +@is_modelopt +@require_accelerator +@require_accelerate +@require_modelopt_version_greater_or_equal("0.33.1") +class ModelOptTesterMixin(ModelOptConfigMixin, QuantizationTesterMixin): + """ + Mixin class for testing NVIDIA ModelOpt quantization on models. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained (e.g., {"subfolder": "transformer"}) + + Expected methods to be implemented by subclasses: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Optional class attributes: + - MODELOPT_CONFIGS: Dict of config name -> NVIDIAModelOptConfig kwargs to test + + Pytest mark: modelopt + Use `pytest -m "not modelopt"` to skip these tests + """ + + @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) + def test_modelopt_quantization_num_parameters(self, config_name): + self._test_quantization_num_parameters(ModelOptConfigMixin.MODELOPT_CONFIGS[config_name]) + + @pytest.mark.parametrize( + "config_name", + list(ModelOptConfigMixin.MODELOPT_CONFIGS.keys()), + ids=list(ModelOptConfigMixin.MODELOPT_CONFIGS.keys()), + ) + def test_modelopt_quantization_memory_footprint(self, config_name): + expected = ModelOptConfigMixin.MODELOPT_EXPECTED_MEMORY_REDUCTIONS.get(config_name, 1.2) + self._test_quantization_memory_footprint( + ModelOptConfigMixin.MODELOPT_CONFIGS[config_name], expected_memory_reduction=expected + ) + + @pytest.mark.parametrize( + "config_name", + list(ModelOptConfigMixin.MODELOPT_CONFIGS.keys()), + ids=list(ModelOptConfigMixin.MODELOPT_CONFIGS.keys()), + ) + def test_modelopt_quantization_inference(self, config_name): + self._test_quantization_inference(ModelOptConfigMixin.MODELOPT_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) + def test_modelopt_quantization_dtype_assignment(self, config_name): + self._test_quantization_dtype_assignment(ModelOptConfigMixin.MODELOPT_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) + def test_modelopt_quantization_lora_inference(self, config_name): + self._test_quantization_lora_inference(ModelOptConfigMixin.MODELOPT_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) + def test_modelopt_quantization_serialization(self, config_name, tmp_path): + self._test_quantization_serialization(ModelOptConfigMixin.MODELOPT_CONFIGS[config_name], tmp_path) + + @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) + def test_modelopt_quantized_layers(self, config_name): + self._test_quantized_layers(ModelOptConfigMixin.MODELOPT_CONFIGS[config_name]) + + def test_modelopt_modules_to_not_convert(self): + """Test that modules_to_not_convert parameter works correctly.""" + modules_to_exclude = getattr(self, "modules_to_not_convert_for_test", None) + if modules_to_exclude is None: + pytest.skip("modules_to_not_convert_for_test not defined for this model") + + self._test_quantization_modules_to_not_convert(ModelOptConfigMixin.MODELOPT_CONFIGS["fp8"], modules_to_exclude) + + def test_modelopt_device_map(self): + """Test that device_map='auto' works correctly with quantization.""" + self._test_quantization_device_map(ModelOptConfigMixin.MODELOPT_CONFIGS["fp8"]) + + def test_modelopt_dequantize(self): + """Test that dequantize() works correctly.""" + self._test_dequantize(ModelOptConfigMixin.MODELOPT_CONFIGS["fp8"]) + + +@is_quantization +@is_torch_compile +class QuantizationCompileTesterMixin: + """ + Base mixin class providing common test implementations for torch.compile with quantized models. + + Backend-specific compile mixins should: + 1. Inherit from their respective config mixin (e.g., BitsAndBytesConfigMixin) + 2. Inherit from this mixin + 3. Define the config to use for compile tests + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained + + Expected methods in test classes: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + """ + + def setup_method(self): + gc.collect() + backend_empty_cache(torch_device) + torch.compiler.reset() + + def teardown_method(self): + gc.collect() + backend_empty_cache(torch_device) + torch.compiler.reset() + + @torch.no_grad() + def _test_torch_compile(self, config_kwargs): + """ + Test that torch.compile works correctly with a quantized model. + + Args: + config_kwargs: Quantization config parameters + """ + model = self._create_quantized_model(config_kwargs) + model.to(torch_device) + model.eval() + + model = torch.compile(model, fullgraph=True) + + with torch._dynamo.config.patch(error_on_recompile=True): + inputs = self.get_dummy_inputs() + output = model(**inputs, return_dict=False)[0] + assert output is not None, "Model output is None" + assert not torch.isnan(output).any(), "Model output contains NaN" + + @torch.no_grad() + def _test_torch_compile_with_group_offload(self, config_kwargs, use_stream=False): + """ + Test that torch.compile works correctly with a quantized model and group offloading. + + Args: + config_kwargs: Quantization config parameters + use_stream: Whether to use CUDA streams for offloading + """ + torch._dynamo.config.cache_size_limit = 1000 + + model = self._create_quantized_model(config_kwargs) + model.eval() + + if not hasattr(model, "enable_group_offload"): + pytest.skip("Model does not support group offloading") + + group_offload_kwargs = { + "onload_device": torch.device(torch_device), + "offload_device": torch.device("cpu"), + "offload_type": "leaf_level", + "use_stream": use_stream, + } + model.enable_group_offload(**group_offload_kwargs) + model = torch.compile(model) + + inputs = self.get_dummy_inputs() + output = model(**inputs, return_dict=False)[0] + assert output is not None, "Model output is None" + assert not torch.isnan(output).any(), "Model output contains NaN" + + +@is_bitsandbytes +@require_accelerator +@require_bitsandbytes_version_greater("0.43.2") +@require_accelerate +class BitsAndBytesCompileTesterMixin(BitsAndBytesConfigMixin, QuantizationCompileTesterMixin): + """ + Mixin class for testing torch.compile with BitsAndBytes quantized models. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained + + Expected methods to be implemented by subclasses: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest mark: bitsandbytes + Use `pytest -m "not bitsandbytes"` to skip these tests + """ + + @pytest.mark.parametrize("config_name", ["4bit_nf4"], ids=["4bit_nf4"]) + def test_bnb_torch_compile(self, config_name): + self._test_torch_compile(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", ["4bit_nf4"], ids=["4bit_nf4"]) + def test_bnb_torch_compile_with_group_offload(self, config_name): + self._test_torch_compile_with_group_offload(BitsAndBytesConfigMixin.BNB_CONFIGS[config_name]) + + +@is_quanto +@require_quanto +@require_accelerate +@require_accelerator +class QuantoCompileTesterMixin(QuantoConfigMixin, QuantizationCompileTesterMixin): + """ + Mixin class for testing torch.compile with Quanto quantized models. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained + + Expected methods to be implemented by subclasses: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest mark: quanto + Use `pytest -m "not quanto"` to skip these tests + """ + + @pytest.mark.parametrize("weight_type_name", ["int8"], ids=["int8"]) + def test_quanto_torch_compile(self, weight_type_name): + self._test_torch_compile(QuantoConfigMixin.QUANTO_WEIGHT_TYPES[weight_type_name]) + + @pytest.mark.parametrize("weight_type_name", ["int8"], ids=["int8"]) + def test_quanto_torch_compile_with_group_offload(self, weight_type_name): + self._test_torch_compile_with_group_offload(QuantoConfigMixin.QUANTO_WEIGHT_TYPES[weight_type_name]) + + +@is_torchao +@require_accelerator +@require_torchao_version_greater_or_equal("0.7.0") +class TorchAoCompileTesterMixin(TorchAoConfigMixin, QuantizationCompileTesterMixin): + """ + Mixin class for testing torch.compile with TorchAO quantized models. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained + + Expected methods to be implemented by subclasses: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest mark: torchao + Use `pytest -m "not torchao"` to skip these tests + """ + + @pytest.mark.parametrize("quant_type", ["int8wo"], ids=["int8wo"]) + def test_torchao_torch_compile(self, quant_type): + self._test_torch_compile(TorchAoConfigMixin.TORCHAO_QUANT_TYPES[quant_type]) + + @pytest.mark.parametrize("quant_type", ["int8wo"], ids=["int8wo"]) + def test_torchao_torch_compile_with_group_offload(self, quant_type): + self._test_torch_compile_with_group_offload(TorchAoConfigMixin.TORCHAO_QUANT_TYPES[quant_type]) + + +@is_gguf +@require_accelerate +@require_accelerator +@require_gguf_version_greater_or_equal("0.10.0") +class GGUFCompileTesterMixin(GGUFConfigMixin, QuantizationCompileTesterMixin): + """ + Mixin class for testing torch.compile with GGUF quantized models. + + Expected from config mixin: + - model_class: The model class to test + + Required properties (must be implemented by subclasses): + - gguf_filename: URL or path to the GGUF file + + Expected methods from config mixin: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest mark: gguf + Use `pytest -m "not gguf"` to skip these tests + """ + + def test_gguf_torch_compile(self): + self._test_torch_compile({"compute_dtype": torch.bfloat16}) + + def test_gguf_torch_compile_with_group_offload(self): + self._test_torch_compile_with_group_offload({"compute_dtype": torch.bfloat16}) + + +@is_modelopt +@require_accelerator +@require_accelerate +@require_modelopt_version_greater_or_equal("0.33.1") +class ModelOptCompileTesterMixin(ModelOptConfigMixin, QuantizationCompileTesterMixin): + """ + Mixin class for testing torch.compile with NVIDIA ModelOpt quantized models. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained + + Expected methods to be implemented by subclasses: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest mark: modelopt + Use `pytest -m "not modelopt"` to skip these tests + """ + + @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) + def test_modelopt_torch_compile(self, config_name): + self._test_torch_compile(ModelOptConfigMixin.MODELOPT_CONFIGS[config_name]) + + @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) + def test_modelopt_torch_compile_with_group_offload(self, config_name): + self._test_torch_compile_with_group_offload(ModelOptConfigMixin.MODELOPT_CONFIGS[config_name]) diff --git a/tests/models/testing_utils/single_file.py b/tests/models/testing_utils/single_file.py new file mode 100644 index 000000000000..e2b9dadb6140 --- /dev/null +++ b/tests/models/testing_utils/single_file.py @@ -0,0 +1,272 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import gc + +import torch +from huggingface_hub import hf_hub_download, snapshot_download + +from diffusers.loaders.single_file_utils import _extract_repo_id_and_weights_name + +from ...testing_utils import ( + backend_empty_cache, + is_single_file, + nightly, + require_torch_accelerator, + torch_device, +) +from .common import check_device_map_is_respected + + +def download_single_file_checkpoint(pretrained_model_name_or_path, filename, tmpdir): + """Download a single file checkpoint from the Hub to a temporary directory.""" + path = hf_hub_download(pretrained_model_name_or_path, filename=filename, local_dir=tmpdir) + return path + + +def download_diffusers_config(pretrained_model_name_or_path, tmpdir): + """Download diffusers config files (excluding weights) from a repository.""" + path = snapshot_download( + pretrained_model_name_or_path, + ignore_patterns=[ + "**/*.ckpt", + "*.ckpt", + "**/*.bin", + "*.bin", + "**/*.pt", + "*.pt", + "**/*.safetensors", + "*.safetensors", + ], + allow_patterns=["**/*.json", "*.json", "*.txt", "**/*.txt"], + local_dir=tmpdir, + ) + return path + + +@nightly +@require_torch_accelerator +@is_single_file +class SingleFileTesterMixin: + """ + Mixin class for testing single file loading for models. + + Required properties (must be implemented by subclasses): + - ckpt_path: Path or Hub path to the single file checkpoint + + Optional properties: + - torch_dtype: torch dtype to use for testing (default: None) + - alternate_ckpt_paths: List of alternate checkpoint paths for variant testing (default: None) + + Expected from config mixin: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - pretrained_model_kwargs: Additional kwargs for from_pretrained (e.g., subfolder) + + Pytest mark: single_file + Use `pytest -m "not single_file"` to skip these tests + """ + + # ==================== Required Properties ==================== + + @property + def ckpt_path(self) -> str: + """Path or Hub path to the single file checkpoint. Must be implemented by subclasses.""" + raise NotImplementedError("Subclasses must implement the `ckpt_path` property.") + + # ==================== Optional Properties ==================== + + @property + def torch_dtype(self) -> torch.dtype | None: + """torch dtype to use for single file testing.""" + return None + + @property + def alternate_ckpt_paths(self) -> list[str] | None: + """List of alternate checkpoint paths for variant testing.""" + return None + + def setup_method(self): + gc.collect() + backend_empty_cache(torch_device) + + def teardown_method(self): + gc.collect() + backend_empty_cache(torch_device) + + def test_single_file_model_config(self): + pretrained_kwargs = {"device": torch_device, **self.pretrained_model_kwargs} + single_file_kwargs = {"device": torch_device} + + if self.torch_dtype: + pretrained_kwargs["torch_dtype"] = self.torch_dtype + single_file_kwargs["torch_dtype"] = self.torch_dtype + + model = self.model_class.from_pretrained(self.pretrained_model_name_or_path, **pretrained_kwargs) + model_single_file = self.model_class.from_single_file(self.ckpt_path, **single_file_kwargs) + + PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"] + for param_name, param_value in model_single_file.config.items(): + if param_name in PARAMS_TO_IGNORE: + continue + assert model.config[param_name] == param_value, ( + f"{param_name} differs between pretrained loading and single file loading: " + f"pretrained={model.config[param_name]}, single_file={param_value}" + ) + + def test_single_file_model_parameters(self): + pretrained_kwargs = {"device_map": str(torch_device), **self.pretrained_model_kwargs} + single_file_kwargs = {"device": torch_device} + + if self.torch_dtype: + pretrained_kwargs["torch_dtype"] = self.torch_dtype + single_file_kwargs["torch_dtype"] = self.torch_dtype + + # Load pretrained model, get state dict on CPU, then free GPU memory + model = self.model_class.from_pretrained(self.pretrained_model_name_or_path, **pretrained_kwargs) + state_dict = {k: v.cpu() for k, v in model.state_dict().items()} + del model + gc.collect() + backend_empty_cache(torch_device) + + # Load single file model, get state dict on CPU + model_single_file = self.model_class.from_single_file(self.ckpt_path, **single_file_kwargs) + state_dict_single_file = {k: v.cpu() for k, v in model_single_file.state_dict().items()} + del model_single_file + gc.collect() + backend_empty_cache(torch_device) + + assert set(state_dict.keys()) == set(state_dict_single_file.keys()), ( + "Model parameters keys differ between pretrained and single file loading. " + f"Missing in single file: {set(state_dict.keys()) - set(state_dict_single_file.keys())}. " + f"Extra in single file: {set(state_dict_single_file.keys()) - set(state_dict.keys())}" + ) + + for key in state_dict.keys(): + param = state_dict[key] + param_single_file = state_dict_single_file[key] + + assert param.shape == param_single_file.shape, ( + f"Parameter shape mismatch for {key}: " + f"pretrained {param.shape} vs single file {param_single_file.shape}" + ) + + assert torch.equal(param, param_single_file), f"Parameter values differ for {key}" + + def test_single_file_loading_local_files_only(self, tmp_path): + single_file_kwargs = {} + + if self.torch_dtype: + single_file_kwargs["torch_dtype"] = self.torch_dtype + + pretrained_model_name_or_path, weight_name = _extract_repo_id_and_weights_name(self.ckpt_path) + local_ckpt_path = download_single_file_checkpoint(pretrained_model_name_or_path, weight_name, str(tmp_path)) + + model_single_file = self.model_class.from_single_file( + local_ckpt_path, local_files_only=True, **single_file_kwargs + ) + + assert model_single_file is not None, "Failed to load model with local_files_only=True" + + def test_single_file_loading_with_diffusers_config(self): + single_file_kwargs = {} + + if self.torch_dtype: + single_file_kwargs["torch_dtype"] = self.torch_dtype + single_file_kwargs.update(self.pretrained_model_kwargs) + + # Load with config parameter + model_single_file = self.model_class.from_single_file( + self.ckpt_path, config=self.pretrained_model_name_or_path, **single_file_kwargs + ) + + # Load pretrained for comparison + pretrained_kwargs = {**self.pretrained_model_kwargs} + if self.torch_dtype: + pretrained_kwargs["torch_dtype"] = self.torch_dtype + + model = self.model_class.from_pretrained(self.pretrained_model_name_or_path, **pretrained_kwargs) + + # Compare configs + PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"] + for param_name, param_value in model_single_file.config.items(): + if param_name in PARAMS_TO_IGNORE: + continue + assert model.config[param_name] == param_value, ( + f"{param_name} differs: pretrained={model.config[param_name]}, single_file={param_value}" + ) + + def test_single_file_loading_with_diffusers_config_local_files_only(self, tmp_path): + single_file_kwargs = {} + + if self.torch_dtype: + single_file_kwargs["torch_dtype"] = self.torch_dtype + single_file_kwargs.update(self.pretrained_model_kwargs) + + pretrained_model_name_or_path, weight_name = _extract_repo_id_and_weights_name(self.ckpt_path) + local_ckpt_path = download_single_file_checkpoint(pretrained_model_name_or_path, weight_name, str(tmp_path)) + local_diffusers_config = download_diffusers_config(self.pretrained_model_name_or_path, str(tmp_path)) + + model_single_file = self.model_class.from_single_file( + local_ckpt_path, config=local_diffusers_config, local_files_only=True, **single_file_kwargs + ) + + assert model_single_file is not None, "Failed to load model with config and local_files_only=True" + + def test_single_file_loading_dtype(self): + for dtype in [torch.float32, torch.float16]: + if torch_device == "mps" and dtype == torch.bfloat16: + continue + + model_single_file = self.model_class.from_single_file(self.ckpt_path, torch_dtype=dtype) + + assert model_single_file.dtype == dtype, f"Expected dtype {dtype}, got {model_single_file.dtype}" + + # Cleanup + del model_single_file + gc.collect() + backend_empty_cache(torch_device) + + def test_checkpoint_variant_loading(self): + if not self.alternate_ckpt_paths: + return + + for ckpt_path in self.alternate_ckpt_paths: + backend_empty_cache(torch_device) + + single_file_kwargs = {} + if self.torch_dtype: + single_file_kwargs["torch_dtype"] = self.torch_dtype + + model = self.model_class.from_single_file(ckpt_path, **single_file_kwargs) + + assert model is not None, f"Failed to load checkpoint from {ckpt_path}" + + del model + gc.collect() + backend_empty_cache(torch_device) + + def test_single_file_loading_with_device_map(self): + single_file_kwargs = {"device_map": torch_device} + + if self.torch_dtype: + single_file_kwargs["torch_dtype"] = self.torch_dtype + + model = self.model_class.from_single_file(self.ckpt_path, **single_file_kwargs) + + assert model is not None, "Failed to load model with device_map" + assert hasattr(model, "hf_device_map"), "Model should have hf_device_map attribute when loaded with device_map" + assert model.hf_device_map is not None, "hf_device_map should not be None when loaded with device_map" + check_device_map_is_respected(model, model.hf_device_map) diff --git a/tests/models/testing_utils/training.py b/tests/models/testing_utils/training.py new file mode 100644 index 000000000000..44cce6af68e5 --- /dev/null +++ b/tests/models/testing_utils/training.py @@ -0,0 +1,220 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import gc + +import pytest +import torch + +from diffusers.training_utils import EMAModel + +from ...testing_utils import ( + backend_empty_cache, + is_training, + require_torch_accelerator_with_training, + torch_all_close, + torch_device, +) + + +@is_training +@require_torch_accelerator_with_training +class TrainingTesterMixin: + """ + Mixin class for testing training functionality on models. + + Expected from config mixin: + - model_class: The model class to test + - output_shape: Tuple defining the expected output shape + + Expected methods from config mixin: + - get_init_dict(): Returns dict of arguments to initialize the model + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest mark: training + Use `pytest -m "not training"` to skip these tests + """ + + def setup_method(self): + gc.collect() + backend_empty_cache(torch_device) + + def teardown_method(self): + gc.collect() + backend_empty_cache(torch_device) + + def test_training(self): + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + + model = self.model_class(**init_dict) + model.to(torch_device) + model.train() + output = model(**inputs_dict, return_dict=False)[0] + + noise = torch.randn((output.shape[0],) + self.output_shape).to(torch_device) + loss = torch.nn.functional.mse_loss(output, noise) + loss.backward() + + def test_training_with_ema(self): + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + + model = self.model_class(**init_dict) + model.to(torch_device) + model.train() + ema_model = EMAModel(model.parameters()) + + output = model(**inputs_dict, return_dict=False)[0] + + noise = torch.randn((output.shape[0],) + self.output_shape).to(torch_device) + loss = torch.nn.functional.mse_loss(output, noise) + loss.backward() + ema_model.step(model.parameters()) + + def test_gradient_checkpointing(self): + if not self.model_class._supports_gradient_checkpointing: + pytest.skip("Gradient checkpointing is not supported.") + + init_dict = self.get_init_dict() + + # at init model should have gradient checkpointing disabled + model = self.model_class(**init_dict) + assert not model.is_gradient_checkpointing, "Gradient checkpointing should be disabled at init" + + # check enable works + model.enable_gradient_checkpointing() + assert model.is_gradient_checkpointing, "Gradient checkpointing should be enabled" + + # check disable works + model.disable_gradient_checkpointing() + assert not model.is_gradient_checkpointing, "Gradient checkpointing should be disabled" + + def test_gradient_checkpointing_is_applied(self, expected_set=None): + if not self.model_class._supports_gradient_checkpointing: + pytest.skip("Gradient checkpointing is not supported.") + + if expected_set is None: + pytest.skip("expected_set must be provided to verify gradient checkpointing is applied.") + + init_dict = self.get_init_dict() + + model_class_copy = copy.copy(self.model_class) + model = model_class_copy(**init_dict) + model.enable_gradient_checkpointing() + + modules_with_gc_enabled = {} + for submodule in model.modules(): + if hasattr(submodule, "gradient_checkpointing"): + assert submodule.gradient_checkpointing, f"{submodule.__class__.__name__} should have GC enabled" + modules_with_gc_enabled[submodule.__class__.__name__] = True + + assert set(modules_with_gc_enabled.keys()) == expected_set, ( + f"Modules with GC enabled {set(modules_with_gc_enabled.keys())} do not match expected set {expected_set}" + ) + assert all(modules_with_gc_enabled.values()), "All modules should have GC enabled" + + def test_gradient_checkpointing_equivalence(self, loss_tolerance=1e-5, param_grad_tol=5e-5, skip=None): + if not self.model_class._supports_gradient_checkpointing: + pytest.skip("Gradient checkpointing is not supported.") + + if skip is None: + skip = set() + + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + inputs_dict_copy = copy.deepcopy(inputs_dict) + + torch.manual_seed(0) + model = self.model_class(**init_dict) + model.to(torch_device) + + assert not model.is_gradient_checkpointing and model.training + + out = model(**inputs_dict, return_dict=False)[0] + + # run the backwards pass on the model + model.zero_grad() + + labels = torch.randn_like(out) + loss = (out - labels).mean() + loss.backward() + + # re-instantiate the model now enabling gradient checkpointing + torch.manual_seed(0) + model_2 = self.model_class(**init_dict) + # clone model + model_2.load_state_dict(model.state_dict()) + model_2.to(torch_device) + model_2.enable_gradient_checkpointing() + + assert model_2.is_gradient_checkpointing and model_2.training + + out_2 = model_2(**inputs_dict_copy, return_dict=False)[0] + + # run the backwards pass on the model + model_2.zero_grad() + loss_2 = (out_2 - labels).mean() + loss_2.backward() + + # compare the output and parameters gradients + assert (loss - loss_2).abs() < loss_tolerance, ( + f"Loss difference {(loss - loss_2).abs()} exceeds tolerance {loss_tolerance}" + ) + + named_params = dict(model.named_parameters()) + named_params_2 = dict(model_2.named_parameters()) + + for name, param in named_params.items(): + if "post_quant_conv" in name: + continue + if name in skip: + continue + if param.grad is None: + continue + + assert torch_all_close(param.grad.data, named_params_2[name].grad.data, atol=param_grad_tol), ( + f"Gradient mismatch for {name}" + ) + + def test_mixed_precision_training(self): + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + + model = self.model_class(**init_dict) + model.to(torch_device) + model.train() + + # Test with float16 + if torch.device(torch_device).type != "cpu": + with torch.amp.autocast(device_type=torch.device(torch_device).type, dtype=torch.float16): + output = model(**inputs_dict, return_dict=False)[0] + + noise = torch.randn((output.shape[0],) + self.output_shape).to(torch_device) + loss = torch.nn.functional.mse_loss(output, noise) + + loss.backward() + + # Test with bfloat16 + if torch.device(torch_device).type != "cpu": + model.zero_grad() + with torch.amp.autocast(device_type=torch.device(torch_device).type, dtype=torch.bfloat16): + output = model(**inputs_dict, return_dict=False)[0] + + noise = torch.randn((output.shape[0],) + self.output_shape).to(torch_device) + loss = torch.nn.functional.mse_loss(output, noise) + + loss.backward() diff --git a/tests/models/transformers/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py index 3ab02f797b5b..2d39dadfcad1 100644 --- a/tests/models/transformers/test_models_transformer_flux.py +++ b/tests/models/transformers/test_models_transformer_flux.py @@ -13,23 +13,52 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest +from typing import Any +import pytest import torch from diffusers import FluxTransformer2DModel -from diffusers.models.attention_processor import FluxIPAdapterJointAttnProcessor2_0 from diffusers.models.embeddings import ImageProjection - -from ...testing_utils import enable_full_determinism, is_peft_available, torch_device -from ..test_modeling_common import LoraHotSwappingForModelTesterMixin, ModelTesterMixin, TorchCompileTesterMixin +from diffusers.models.transformers.transformer_flux import FluxIPAdapterAttnProcessor +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + BitsAndBytesCompileTesterMixin, + BitsAndBytesTesterMixin, + ContextParallelTesterMixin, + FasterCacheTesterMixin, + FirstBlockCacheTesterMixin, + GGUFCompileTesterMixin, + GGUFTesterMixin, + IPAdapterTesterMixin, + LoraHotSwappingForModelTesterMixin, + LoraTesterMixin, + MemoryTesterMixin, + ModelOptCompileTesterMixin, + ModelOptTesterMixin, + ModelTesterMixin, + PyramidAttentionBroadcastTesterMixin, + QuantoCompileTesterMixin, + QuantoTesterMixin, + SingleFileTesterMixin, + TorchAoCompileTesterMixin, + TorchAoTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, +) enable_full_determinism() -def create_flux_ip_adapter_state_dict(model): - # "ip_adapter" (cross-attention weights) +# TODO: This standalone function maintains backward compatibility with pipeline tests +# (tests/pipelines/test_pipelines_common.py) and will be refactored. +def create_flux_ip_adapter_state_dict(model) -> dict[str, dict[str, Any]]: + """Create a dummy IP Adapter state dict for Flux transformer testing.""" ip_cross_attn_state_dict = {} key_id = 0 @@ -39,7 +68,7 @@ def create_flux_ip_adapter_state_dict(model): joint_attention_dim = model.config["joint_attention_dim"] hidden_size = model.config["num_attention_heads"] * model.config["attention_head_dim"] - sd = FluxIPAdapterJointAttnProcessor2_0( + sd = FluxIPAdapterAttnProcessor( hidden_size=hidden_size, cross_attention_dim=joint_attention_dim, scale=1.0 ).state_dict() ip_cross_attn_state_dict.update( @@ -50,11 +79,8 @@ def create_flux_ip_adapter_state_dict(model): f"{key_id}.to_v_ip.bias": sd["to_v_ip.0.bias"], } ) - key_id += 1 - # "image_proj" (ImageProjection layer weights) - image_projection = ImageProjection( cross_attention_dim=model.config["joint_attention_dim"], image_embed_dim=( @@ -75,57 +101,45 @@ def create_flux_ip_adapter_state_dict(model): ) del sd - ip_state_dict = {} - ip_state_dict.update({"image_proj": ip_image_projection_state_dict, "ip_adapter": ip_cross_attn_state_dict}) - return ip_state_dict + return {"image_proj": ip_image_projection_state_dict, "ip_adapter": ip_cross_attn_state_dict} -class FluxTransformerTests(ModelTesterMixin, unittest.TestCase): - model_class = FluxTransformer2DModel - main_input_name = "hidden_states" - # We override the items here because the transformer under consideration is small. - model_split_percents = [0.7, 0.6, 0.6] +class FluxTransformerTesterConfig(BaseModelTesterConfig): + @property + def model_class(self): + return FluxTransformer2DModel - # Skip setting testing with default: AttnProcessor - uses_custom_attn_processor = True + @property + def pretrained_model_name_or_path(self): + return "hf-internal-testing/tiny-flux-pipe" @property - def dummy_input(self): - return self.prepare_dummy_input() + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} @property - def input_shape(self): + def output_shape(self) -> tuple[int, int]: return (16, 4) @property - def output_shape(self): + def input_shape(self) -> tuple[int, int]: return (16, 4) - def prepare_dummy_input(self, height=4, width=4): - batch_size = 1 - num_latent_channels = 4 - num_image_channels = 3 - sequence_length = 48 - embedding_dim = 32 + @property + def model_split_percents(self) -> list: + return [0.9] - hidden_states = torch.randn((batch_size, height * width, num_latent_channels)).to(torch_device) - encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to(torch_device) - pooled_prompt_embeds = torch.randn((batch_size, embedding_dim)).to(torch_device) - text_ids = torch.randn((sequence_length, num_image_channels)).to(torch_device) - image_ids = torch.randn((height * width, num_image_channels)).to(torch_device) - timestep = torch.tensor([1.0]).to(torch_device).expand(batch_size) + @property + def main_input_name(self) -> str: + return "hidden_states" - return { - "hidden_states": hidden_states, - "encoder_hidden_states": encoder_hidden_states, - "img_ids": image_ids, - "txt_ids": text_ids, - "pooled_projections": pooled_prompt_embeds, - "timestep": timestep, - } + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) - def prepare_init_args_and_inputs_for_common(self): - init_dict = { + def get_init_dict(self) -> dict[str, int | list[int]]: + """Return Flux model initialization arguments.""" + return { "patch_size": 1, "in_channels": 4, "num_layers": 1, @@ -137,11 +151,40 @@ def prepare_init_args_and_inputs_for_common(self): "axes_dims_rope": [4, 4, 8], } - inputs_dict = self.dummy_input - return init_dict, inputs_dict + def get_dummy_inputs(self) -> dict[str, torch.Tensor]: + batch_size = 1 + height = width = 4 + num_latent_channels = 4 + num_image_channels = 3 + sequence_length = 48 + embedding_dim = 32 + return { + "hidden_states": randn_tensor( + (batch_size, height * width, num_latent_channels), generator=self.generator, device=torch_device + ), + "encoder_hidden_states": randn_tensor( + (batch_size, sequence_length, embedding_dim), generator=self.generator, device=torch_device + ), + "pooled_projections": randn_tensor( + (batch_size, embedding_dim), generator=self.generator, device=torch_device + ), + "img_ids": randn_tensor( + (height * width, num_image_channels), generator=self.generator, device=torch_device + ), + "txt_ids": randn_tensor( + (sequence_length, num_image_channels), generator=self.generator, device=torch_device + ), + "timestep": torch.tensor([1.0]).to(torch_device).expand(batch_size), + } + + +class TestFluxTransformer(FluxTransformerTesterConfig, ModelTesterMixin): def test_deprecated_inputs_img_txt_ids_3d(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + """Test that deprecated 3D img_ids and txt_ids still work.""" + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict) model.to(torch_device) model.eval() @@ -162,63 +205,228 @@ def test_deprecated_inputs_img_txt_ids_3d(self): with torch.no_grad(): output_2 = model(**inputs_dict).to_tuple()[0] - self.assertEqual(output_1.shape, output_2.shape) - self.assertTrue( - torch.allclose(output_1, output_2, atol=1e-5), - msg="output with deprecated inputs (img_ids and txt_ids as 3d torch tensors) are not equal as them as 2d inputs", + assert output_1.shape == output_2.shape + assert torch.allclose(output_1, output_2, atol=1e-5), ( + "output with deprecated inputs (img_ids and txt_ids as 3d torch tensors) " + "are not equal as them as 2d inputs" ) - def test_gradient_checkpointing_is_applied(self): - expected_set = {"FluxTransformer2DModel"} - super().test_gradient_checkpointing_is_applied(expected_set=expected_set) - - # The test exists for cases like - # https://github.com/huggingface/diffusers/issues/11874 - @unittest.skipIf(not is_peft_available(), "Only with PEFT") - def test_lora_exclude_modules(self): - from peft import LoraConfig, get_peft_model_state_dict, inject_adapter_in_model, set_peft_model_state_dict - - lora_rank = 4 - target_module = "single_transformer_blocks.0.proj_out" - adapter_name = "foo" - init_dict, _ = self.prepare_init_args_and_inputs_for_common() - model = self.model_class(**init_dict).to(torch_device) - - state_dict = model.state_dict() - target_mod_shape = state_dict[f"{target_module}.weight"].shape - lora_state_dict = { - f"{target_module}.lora_A.weight": torch.ones(lora_rank, target_mod_shape[1]) * 22, - f"{target_module}.lora_B.weight": torch.ones(target_mod_shape[0], lora_rank) * 33, + +class TestFluxTransformerMemory(FluxTransformerTesterConfig, MemoryTesterMixin): + """Memory optimization tests for Flux Transformer.""" + + +class TestFluxTransformerTraining(FluxTransformerTesterConfig, TrainingTesterMixin): + """Training tests for Flux Transformer.""" + + +class TestFluxTransformerAttention(FluxTransformerTesterConfig, AttentionTesterMixin): + """Attention processor tests for Flux Transformer.""" + + +class TestFluxTransformerContextParallel(FluxTransformerTesterConfig, ContextParallelTesterMixin): + """Context Parallel inference tests for Flux Transformer""" + + +class TestFluxTransformerIPAdapter(FluxTransformerTesterConfig, IPAdapterTesterMixin): + """IP Adapter tests for Flux Transformer.""" + + @property + def ip_adapter_processor_cls(self): + return FluxIPAdapterAttnProcessor + + def modify_inputs_for_ip_adapter(self, model, inputs_dict): + torch.manual_seed(0) + # Create dummy image embeds for IP adapter + cross_attention_dim = getattr(model.config, "joint_attention_dim", 32) + image_embeds = torch.randn(1, 1, cross_attention_dim).to(torch_device) + + inputs_dict.update({"joint_attention_kwargs": {"ip_adapter_image_embeds": image_embeds}}) + + return inputs_dict + + def create_ip_adapter_state_dict(self, model: Any) -> dict[str, dict[str, Any]]: + return create_flux_ip_adapter_state_dict(model) + + +class TestFluxTransformerLoRA(FluxTransformerTesterConfig, LoraTesterMixin): + """LoRA adapter tests for Flux Transformer.""" + + +class TestFluxTransformerLoRAHotSwap(FluxTransformerTesterConfig, LoraHotSwappingForModelTesterMixin): + """LoRA hot-swapping tests for Flux Transformer.""" + + @property + def different_shapes_for_compilation(self): + return [(4, 4), (4, 8), (8, 8)] + + def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]: + """Override to support dynamic height/width for LoRA hotswap tests.""" + batch_size = 1 + num_latent_channels = 4 + num_image_channels = 3 + sequence_length = 24 + embedding_dim = 32 + + return { + "hidden_states": randn_tensor((batch_size, height * width, num_latent_channels), device=torch_device), + "encoder_hidden_states": randn_tensor((batch_size, sequence_length, embedding_dim), device=torch_device), + "pooled_projections": randn_tensor((batch_size, embedding_dim), device=torch_device), + "img_ids": randn_tensor((height * width, num_image_channels), device=torch_device), + "txt_ids": randn_tensor((sequence_length, num_image_channels), device=torch_device), + "timestep": torch.tensor([1.0]).to(torch_device).expand(batch_size), + } + + +class TestFluxTransformerCompile(FluxTransformerTesterConfig, TorchCompileTesterMixin): + @property + def different_shapes_for_compilation(self): + return [(4, 4), (4, 8), (8, 8)] + + def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]: + """Override to support dynamic height/width for compilation tests.""" + batch_size = 1 + num_latent_channels = 4 + num_image_channels = 3 + sequence_length = 24 + embedding_dim = 32 + + return { + "hidden_states": randn_tensor((batch_size, height * width, num_latent_channels), device=torch_device), + "encoder_hidden_states": randn_tensor((batch_size, sequence_length, embedding_dim), device=torch_device), + "pooled_projections": randn_tensor((batch_size, embedding_dim), device=torch_device), + "img_ids": randn_tensor((height * width, num_image_channels), device=torch_device), + "txt_ids": randn_tensor((sequence_length, num_image_channels), device=torch_device), + "timestep": torch.tensor([1.0]).to(torch_device).expand(batch_size), } - # Passing exclude_modules should no longer be necessary (or even passing target_modules, for that matter). - config = LoraConfig( - r=lora_rank, target_modules=["single_transformer_blocks.0.proj_out"], exclude_modules=["proj_out"] - ) - inject_adapter_in_model(config, model, adapter_name=adapter_name, state_dict=lora_state_dict) - set_peft_model_state_dict(model, lora_state_dict, adapter_name) - retrieved_lora_state_dict = get_peft_model_state_dict(model, adapter_name=adapter_name) - assert len(retrieved_lora_state_dict) == len(lora_state_dict) - assert (retrieved_lora_state_dict["single_transformer_blocks.0.proj_out.lora_A.weight"] == 22).all() - assert (retrieved_lora_state_dict["single_transformer_blocks.0.proj_out.lora_B.weight"] == 33).all() -class FluxTransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase): - model_class = FluxTransformer2DModel - different_shapes_for_compilation = [(4, 4), (4, 8), (8, 8)] +class TestFluxSingleFile(FluxTransformerTesterConfig, SingleFileTesterMixin): + @property + def ckpt_path(self): + return "https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/flux1-dev.safetensors" + + @property + def alternate_ckpt_paths(self): + return ["https://huggingface.co/Comfy-Org/flux1-dev/blob/main/flux1-dev-fp8.safetensors"] + + @property + def pretrained_model_name_or_path(self): + return "black-forest-labs/FLUX.1-dev" + + +class TestFluxTransformerBitsAndBytes(FluxTransformerTesterConfig, BitsAndBytesTesterMixin): + """BitsAndBytes quantization tests for Flux Transformer.""" + + +class TestFluxTransformerQuanto(FluxTransformerTesterConfig, QuantoTesterMixin): + """Quanto quantization tests for Flux Transformer.""" + + @property + def pretrained_model_name_or_path(self): + return "hf-internal-testing/tiny-flux-transformer" + + @property + def pretrained_model_kwargs(self): + return {} + + +class TestFluxTransformerTorchAo(FluxTransformerTesterConfig, TorchAoTesterMixin): + """TorchAO quantization tests for Flux Transformer.""" + + +class TestFluxTransformerGGUF(FluxTransformerTesterConfig, GGUFTesterMixin): + @property + def gguf_filename(self): + return "https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-dev-Q2_K.gguf" + + @property + def torch_dtype(self): + return torch.bfloat16 + + def get_dummy_inputs(self): + """Override to provide inputs matching the real FLUX model dimensions.""" + return { + "hidden_states": randn_tensor( + (1, 4096, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "encoder_hidden_states": randn_tensor( + (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "pooled_projections": randn_tensor( + (1, 768), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "timestep": torch.tensor([1]).to(torch_device, self.torch_dtype), + "img_ids": randn_tensor((4096, 3), generator=self.generator, device=torch_device, dtype=self.torch_dtype), + "txt_ids": randn_tensor((512, 3), generator=self.generator, device=torch_device, dtype=self.torch_dtype), + "guidance": torch.tensor([3.5]).to(torch_device, self.torch_dtype), + } + + +class TestFluxTransformerQuantoCompile(FluxTransformerTesterConfig, QuantoCompileTesterMixin): + """Quanto + compile tests for Flux Transformer.""" + + +class TestFluxTransformerTorchAoCompile(FluxTransformerTesterConfig, TorchAoCompileTesterMixin): + """TorchAO + compile tests for Flux Transformer.""" + + +class TestFluxTransformerGGUFCompile(FluxTransformerTesterConfig, GGUFCompileTesterMixin): + @property + def gguf_filename(self): + return "https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-dev-Q2_K.gguf" + + @property + def torch_dtype(self): + return torch.bfloat16 + + def get_dummy_inputs(self): + """Override to provide inputs matching the real FLUX model dimensions.""" + return { + "hidden_states": randn_tensor( + (1, 4096, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "encoder_hidden_states": randn_tensor( + (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "pooled_projections": randn_tensor( + (1, 768), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "timestep": torch.tensor([1]).to(torch_device, self.torch_dtype), + "img_ids": randn_tensor((4096, 3), generator=self.generator, device=torch_device, dtype=self.torch_dtype), + "txt_ids": randn_tensor((512, 3), generator=self.generator, device=torch_device, dtype=self.torch_dtype), + "guidance": torch.tensor([3.5]).to(torch_device, self.torch_dtype), + } + + +class TestFluxTransformerModelOpt(FluxTransformerTesterConfig, ModelOptTesterMixin): + """ModelOpt quantization tests for Flux Transformer.""" + + +class TestFluxTransformerModelOptCompile(FluxTransformerTesterConfig, ModelOptCompileTesterMixin): + """ModelOpt + compile tests for Flux Transformer.""" + + +@pytest.mark.skip(reason="torch.compile is not supported by BitsAndBytes") +class TestFluxTransformerBitsAndBytesCompile(FluxTransformerTesterConfig, BitsAndBytesCompileTesterMixin): + """BitsAndBytes + compile tests for Flux Transformer.""" + - def prepare_init_args_and_inputs_for_common(self): - return FluxTransformerTests().prepare_init_args_and_inputs_for_common() +class TestFluxTransformerPABCache(FluxTransformerTesterConfig, PyramidAttentionBroadcastTesterMixin): + """PyramidAttentionBroadcast cache tests for Flux Transformer.""" - def prepare_dummy_input(self, height, width): - return FluxTransformerTests().prepare_dummy_input(height=height, width=width) +class TestFluxTransformerFBCCache(FluxTransformerTesterConfig, FirstBlockCacheTesterMixin): + """FirstBlockCache tests for Flux Transformer.""" -class FluxTransformerLoRAHotSwapTests(LoraHotSwappingForModelTesterMixin, unittest.TestCase): - model_class = FluxTransformer2DModel - different_shapes_for_compilation = [(4, 4), (4, 8), (8, 8)] - def prepare_init_args_and_inputs_for_common(self): - return FluxTransformerTests().prepare_init_args_and_inputs_for_common() +class TestFluxTransformerFasterCache(FluxTransformerTesterConfig, FasterCacheTesterMixin): + """FasterCache tests for Flux Transformer.""" - def prepare_dummy_input(self, height, width): - return FluxTransformerTests().prepare_dummy_input(height=height, width=width) + # Flux is guidance distilled, so we can test at model level without CFG batch handling + FASTER_CACHE_CONFIG = { + "spatial_attention_block_skip_range": 2, + "spatial_attention_timestep_skip_range": (-1, 901), + "tensor_format": "BCHW", + "is_guidance_distilled": True, + } diff --git a/tests/models/transformers/test_models_transformer_helios.py b/tests/models/transformers/test_models_transformer_helios.py new file mode 100644 index 000000000000..c365c258e596 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_helios.py @@ -0,0 +1,168 @@ +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import pytest +import torch + +from diffusers import HeliosTransformer3DModel +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, +) + + +enable_full_determinism() + + +class HeliosTransformer3DTesterConfig(BaseModelTesterConfig): + @property + def model_class(self): + return HeliosTransformer3DModel + + @property + def pretrained_model_name_or_path(self): + return "hf-internal-testing/tiny-helios-base-transformer" + + @property + def output_shape(self) -> tuple[int, ...]: + return (4, 2, 16, 16) + + @property + def input_shape(self) -> tuple[int, ...]: + return (4, 2, 16, 16) + + @property + def main_input_name(self) -> str: + return "hidden_states" + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict[str, int | list[int] | tuple | str | bool]: + return { + "patch_size": (1, 2, 2), + "num_attention_heads": 2, + "attention_head_dim": 12, + "in_channels": 4, + "out_channels": 4, + "text_dim": 16, + "freq_dim": 256, + "ffn_dim": 32, + "num_layers": 2, + "cross_attn_norm": True, + "qk_norm": "rms_norm_across_heads", + "rope_dim": (4, 4, 4), + "has_multi_term_memory_patch": True, + "guidance_cross_attn": True, + "zero_history_timestep": True, + "is_amplify_history": False, + } + + def get_dummy_inputs(self) -> dict[str, torch.Tensor]: + batch_size = 1 + num_channels = 4 + num_frames = 2 + height = 16 + width = 16 + text_encoder_embedding_dim = 16 + sequence_length = 12 + + hidden_states = randn_tensor( + (batch_size, num_channels, num_frames, height, width), + generator=self.generator, + device=torch_device, + ) + timestep = torch.randint(0, 1000, size=(batch_size,), generator=self.generator).to(torch_device) + encoder_hidden_states = randn_tensor( + (batch_size, sequence_length, text_encoder_embedding_dim), + generator=self.generator, + device=torch_device, + ) + indices_hidden_states = torch.ones((batch_size, num_frames)).to(torch_device) + indices_latents_history_short = torch.ones((batch_size, num_frames - 1)).to(torch_device) + indices_latents_history_mid = torch.ones((batch_size, num_frames - 1)).to(torch_device) + indices_latents_history_long = torch.ones((batch_size, (num_frames - 1) * 4)).to(torch_device) + latents_history_short = randn_tensor( + (batch_size, num_channels, num_frames - 1, height, width), + generator=self.generator, + device=torch_device, + ) + latents_history_mid = randn_tensor( + (batch_size, num_channels, num_frames - 1, height, width), + generator=self.generator, + device=torch_device, + ) + latents_history_long = randn_tensor( + (batch_size, num_channels, (num_frames - 1) * 4, height, width), + generator=self.generator, + device=torch_device, + ) + + return { + "hidden_states": hidden_states, + "timestep": timestep, + "encoder_hidden_states": encoder_hidden_states, + "indices_hidden_states": indices_hidden_states, + "indices_latents_history_short": indices_latents_history_short, + "indices_latents_history_mid": indices_latents_history_mid, + "indices_latents_history_long": indices_latents_history_long, + "latents_history_short": latents_history_short, + "latents_history_mid": latents_history_mid, + "latents_history_long": latents_history_long, + } + + +class TestHeliosTransformer3D(HeliosTransformer3DTesterConfig, ModelTesterMixin): + """Core model tests for Helios Transformer 3D.""" + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) + def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype): + # Skip: fp16/bf16 require very high atol to pass, providing little signal. + # Dtype preservation is already tested by test_from_save_pretrained_dtype and test_keep_in_fp32_modules. + pytest.skip("Tolerance requirements too high for meaningful test") + + +class TestHeliosTransformer3DMemory(HeliosTransformer3DTesterConfig, MemoryTesterMixin): + """Memory optimization tests for Helios Transformer 3D.""" + + +class TestHeliosTransformer3DTraining(HeliosTransformer3DTesterConfig, TrainingTesterMixin): + """Training tests for Helios Transformer 3D.""" + + def test_gradient_checkpointing_is_applied(self): + expected_set = {"HeliosTransformer3DModel"} + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + + +class TestHeliosTransformer3DAttention(HeliosTransformer3DTesterConfig, AttentionTesterMixin): + """Attention processor tests for Helios Transformer 3D.""" + + +class TestHeliosTransformer3DCompile(HeliosTransformer3DTesterConfig, TorchCompileTesterMixin): + """Torch compile tests for Helios Transformer 3D.""" + + @pytest.mark.xfail( + reason="Helios DiT does not compile when deterministic algorithms are used due to https://github.com/pytorch/pytorch/issues/170079" + ) + def test_torch_compile_recompilation_and_graph_break(self): + super().test_torch_compile_recompilation_and_graph_break() diff --git a/tests/models/transformers/test_models_transformer_ltx2.py b/tests/models/transformers/test_models_transformer_ltx2.py new file mode 100644 index 000000000000..af9ef0623891 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_ltx2.py @@ -0,0 +1,222 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import torch + +from diffusers import LTX2VideoTransformer3DModel + +from ...testing_utils import enable_full_determinism, torch_device +from ..test_modeling_common import ModelTesterMixin, TorchCompileTesterMixin + + +enable_full_determinism() + + +class LTX2TransformerTests(ModelTesterMixin, unittest.TestCase): + model_class = LTX2VideoTransformer3DModel + main_input_name = "hidden_states" + uses_custom_attn_processor = True + + @property + def dummy_input(self): + # Common + batch_size = 2 + + # Video + num_frames = 2 + num_channels = 4 + height = 16 + width = 16 + + # Audio + audio_num_frames = 9 + audio_num_channels = 2 + num_mel_bins = 2 + + # Text + embedding_dim = 16 + sequence_length = 16 + + hidden_states = torch.randn((batch_size, num_frames * height * width, num_channels)).to(torch_device) + audio_hidden_states = torch.randn((batch_size, audio_num_frames, audio_num_channels * num_mel_bins)).to( + torch_device + ) + encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to(torch_device) + audio_encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to(torch_device) + encoder_attention_mask = torch.ones((batch_size, sequence_length)).bool().to(torch_device) + timestep = torch.rand((batch_size,)).to(torch_device) * 1000 + + return { + "hidden_states": hidden_states, + "audio_hidden_states": audio_hidden_states, + "encoder_hidden_states": encoder_hidden_states, + "audio_encoder_hidden_states": audio_encoder_hidden_states, + "timestep": timestep, + "encoder_attention_mask": encoder_attention_mask, + "num_frames": num_frames, + "height": height, + "width": width, + "audio_num_frames": audio_num_frames, + "fps": 25.0, + } + + @property + def input_shape(self): + return (512, 4) + + @property + def output_shape(self): + return (512, 4) + + def prepare_init_args_and_inputs_for_common(self): + init_dict = { + "in_channels": 4, + "out_channels": 4, + "patch_size": 1, + "patch_size_t": 1, + "num_attention_heads": 2, + "attention_head_dim": 8, + "cross_attention_dim": 16, + "audio_in_channels": 4, + "audio_out_channels": 4, + "audio_num_attention_heads": 2, + "audio_attention_head_dim": 4, + "audio_cross_attention_dim": 8, + "num_layers": 2, + "qk_norm": "rms_norm_across_heads", + "caption_channels": 16, + "rope_double_precision": False, + } + inputs_dict = self.dummy_input + return init_dict, inputs_dict + + def test_gradient_checkpointing_is_applied(self): + expected_set = {"LTX2VideoTransformer3DModel"} + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + + # def test_ltx2_consistency(self, seed=0, dtype=torch.float32): + # torch.manual_seed(seed) + # init_dict, _ = self.prepare_init_args_and_inputs_for_common() + + # # Calculate dummy inputs in a custom manner to ensure compatibility with original code + # batch_size = 2 + # num_frames = 9 + # latent_frames = 2 + # text_embedding_dim = 16 + # text_seq_len = 16 + # fps = 25.0 + # sampling_rate = 16000.0 + # hop_length = 160.0 + + # sigma = torch.rand((1,), generator=torch.manual_seed(seed), dtype=dtype, device="cpu") * 1000 + # timestep = (sigma * torch.ones((batch_size,), dtype=dtype, device="cpu")).to(device=torch_device) + + # num_channels = 4 + # latent_height = 4 + # latent_width = 4 + # hidden_states = torch.randn( + # (batch_size, num_channels, latent_frames, latent_height, latent_width), + # generator=torch.manual_seed(seed), + # dtype=dtype, + # device="cpu", + # ) + # # Patchify video latents (with patch_size (1, 1, 1)) + # hidden_states = hidden_states.reshape(batch_size, -1, latent_frames, 1, latent_height, 1, latent_width, 1) + # hidden_states = hidden_states.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3) + # encoder_hidden_states = torch.randn( + # (batch_size, text_seq_len, text_embedding_dim), + # generator=torch.manual_seed(seed), + # dtype=dtype, + # device="cpu", + # ) + + # audio_num_channels = 2 + # num_mel_bins = 2 + # latent_length = int((sampling_rate / hop_length / 4) * (num_frames / fps)) + # audio_hidden_states = torch.randn( + # (batch_size, audio_num_channels, latent_length, num_mel_bins), + # generator=torch.manual_seed(seed), + # dtype=dtype, + # device="cpu", + # ) + # # Patchify audio latents + # audio_hidden_states = audio_hidden_states.transpose(1, 2).flatten(2, 3) + # audio_encoder_hidden_states = torch.randn( + # (batch_size, text_seq_len, text_embedding_dim), + # generator=torch.manual_seed(seed), + # dtype=dtype, + # device="cpu", + # ) + + # inputs_dict = { + # "hidden_states": hidden_states.to(device=torch_device), + # "audio_hidden_states": audio_hidden_states.to(device=torch_device), + # "encoder_hidden_states": encoder_hidden_states.to(device=torch_device), + # "audio_encoder_hidden_states": audio_encoder_hidden_states.to(device=torch_device), + # "timestep": timestep, + # "num_frames": latent_frames, + # "height": latent_height, + # "width": latent_width, + # "audio_num_frames": num_frames, + # "fps": 25.0, + # } + + # model = self.model_class.from_pretrained( + # "diffusers-internal-dev/dummy-ltx2", + # subfolder="transformer", + # device_map="cpu", + # ) + # # torch.manual_seed(seed) + # # model = self.model_class(**init_dict) + # model.to(torch_device) + # model.eval() + + # with attention_backend("native"): + # with torch.no_grad(): + # output = model(**inputs_dict) + + # video_output, audio_output = output.to_tuple() + + # self.assertIsNotNone(video_output) + # self.assertIsNotNone(audio_output) + + # # input & output have to have the same shape + # video_expected_shape = (batch_size, latent_frames * latent_height * latent_width, num_channels) + # self.assertEqual(video_output.shape, video_expected_shape, "Video input and output shapes do not match") + # audio_expected_shape = (batch_size, latent_length, audio_num_channels * num_mel_bins) + # self.assertEqual(audio_output.shape, audio_expected_shape, "Audio input and output shapes do not match") + + # # Check against expected slice + # # fmt: off + # video_expected_slice = torch.tensor([0.4783, 1.6954, -1.2092, 0.1762, 0.7801, 1.2025, -1.4525, -0.2721, 0.3354, 1.9144, -1.5546, 0.0831, 0.4391, 1.7012, -1.7373, -0.2676]) + # audio_expected_slice = torch.tensor([-0.4236, 0.4750, 0.3901, -0.4339, -0.2782, 0.4357, 0.4526, -0.3927, -0.0980, 0.4870, 0.3964, -0.3169, -0.3974, 0.4408, 0.3809, -0.4692]) + # # fmt: on + + # video_output_flat = video_output.cpu().flatten().float() + # video_generated_slice = torch.cat([video_output_flat[:8], video_output_flat[-8:]]) + # self.assertTrue(torch.allclose(video_generated_slice, video_expected_slice, atol=1e-4)) + + # audio_output_flat = audio_output.cpu().flatten().float() + # audio_generated_slice = torch.cat([audio_output_flat[:8], audio_output_flat[-8:]]) + # self.assertTrue(torch.allclose(audio_generated_slice, audio_expected_slice, atol=1e-4)) + + +class LTX2TransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase): + model_class = LTX2VideoTransformer3DModel + + def prepare_init_args_and_inputs_for_common(self): + return LTX2TransformerTests().prepare_init_args_and_inputs_for_common() diff --git a/tests/models/transformers/test_models_transformer_qwenimage.py b/tests/models/transformers/test_models_transformer_qwenimage.py index b24fa90503ef..e6b19377b14f 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage.py +++ b/tests/models/transformers/test_models_transformer_qwenimage.py @@ -15,10 +15,10 @@ import unittest -import pytest import torch from diffusers import QwenImageTransformer2DModel +from diffusers.models.transformers.transformer_qwenimage import compute_text_seq_len_from_mask from ...testing_utils import enable_full_determinism, torch_device from ..test_modeling_common import ModelTesterMixin, TorchCompileTesterMixin @@ -68,7 +68,6 @@ def prepare_dummy_input(self, height=4, width=4): "encoder_hidden_states_mask": encoder_hidden_states_mask, "timestep": timestep, "img_shapes": img_shapes, - "txt_seq_lens": encoder_hidden_states_mask.sum(dim=1).tolist(), } def prepare_init_args_and_inputs_for_common(self): @@ -91,6 +90,180 @@ def test_gradient_checkpointing_is_applied(self): expected_set = {"QwenImageTransformer2DModel"} super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + def test_infers_text_seq_len_from_mask(self): + """Test that compute_text_seq_len_from_mask correctly infers sequence lengths and returns tensors.""" + init_dict, inputs = self.prepare_init_args_and_inputs_for_common() + model = self.model_class(**init_dict).to(torch_device) + + # Test 1: Contiguous mask with padding at the end (only first 2 tokens valid) + encoder_hidden_states_mask = inputs["encoder_hidden_states_mask"].clone() + encoder_hidden_states_mask[:, 2:] = 0 # Only first 2 tokens are valid + + rope_text_seq_len, per_sample_len, normalized_mask = compute_text_seq_len_from_mask( + inputs["encoder_hidden_states"], encoder_hidden_states_mask + ) + + # Verify rope_text_seq_len is returned as an int (for torch.compile compatibility) + self.assertIsInstance(rope_text_seq_len, int) + + # Verify per_sample_len is computed correctly (max valid position + 1 = 2) + self.assertIsInstance(per_sample_len, torch.Tensor) + self.assertEqual(int(per_sample_len.max().item()), 2) + + # Verify mask is normalized to bool dtype + self.assertTrue(normalized_mask.dtype == torch.bool) + self.assertEqual(normalized_mask.sum().item(), 2) # Only 2 True values + + # Verify rope_text_seq_len is at least the sequence length + self.assertGreaterEqual(rope_text_seq_len, inputs["encoder_hidden_states"].shape[1]) + + # Test 2: Verify model runs successfully with inferred values + inputs["encoder_hidden_states_mask"] = normalized_mask + with torch.no_grad(): + output = model(**inputs) + self.assertEqual(output.sample.shape[1], inputs["hidden_states"].shape[1]) + + # Test 3: Different mask pattern (padding at beginning) + encoder_hidden_states_mask2 = inputs["encoder_hidden_states_mask"].clone() + encoder_hidden_states_mask2[:, :3] = 0 # First 3 tokens are padding + encoder_hidden_states_mask2[:, 3:] = 1 # Last 4 tokens are valid + + rope_text_seq_len2, per_sample_len2, normalized_mask2 = compute_text_seq_len_from_mask( + inputs["encoder_hidden_states"], encoder_hidden_states_mask2 + ) + + # Max valid position is 6 (last token), so per_sample_len should be 7 + self.assertEqual(int(per_sample_len2.max().item()), 7) + self.assertEqual(normalized_mask2.sum().item(), 4) # 4 True values + + # Test 4: No mask provided (None case) + rope_text_seq_len_none, per_sample_len_none, normalized_mask_none = compute_text_seq_len_from_mask( + inputs["encoder_hidden_states"], None + ) + self.assertEqual(rope_text_seq_len_none, inputs["encoder_hidden_states"].shape[1]) + self.assertIsInstance(rope_text_seq_len_none, int) + self.assertIsNone(per_sample_len_none) + self.assertIsNone(normalized_mask_none) + + def test_non_contiguous_attention_mask(self): + """Test that non-contiguous masks work correctly (e.g., [1, 0, 1, 0, 1, 0, 0])""" + init_dict, inputs = self.prepare_init_args_and_inputs_for_common() + model = self.model_class(**init_dict).to(torch_device) + + # Create a non-contiguous mask pattern: valid, padding, valid, padding, etc. + encoder_hidden_states_mask = inputs["encoder_hidden_states_mask"].clone() + # Pattern: [True, False, True, False, True, False, False] + encoder_hidden_states_mask[:, 1] = 0 + encoder_hidden_states_mask[:, 3] = 0 + encoder_hidden_states_mask[:, 5:] = 0 + + inferred_rope_len, per_sample_len, normalized_mask = compute_text_seq_len_from_mask( + inputs["encoder_hidden_states"], encoder_hidden_states_mask + ) + self.assertEqual(int(per_sample_len.max().item()), 5) + self.assertEqual(inferred_rope_len, inputs["encoder_hidden_states"].shape[1]) + self.assertIsInstance(inferred_rope_len, int) + self.assertTrue(normalized_mask.dtype == torch.bool) + + inputs["encoder_hidden_states_mask"] = normalized_mask + + with torch.no_grad(): + output = model(**inputs) + + self.assertEqual(output.sample.shape[1], inputs["hidden_states"].shape[1]) + + def test_txt_seq_lens_deprecation(self): + """Test that passing txt_seq_lens raises a deprecation warning.""" + init_dict, inputs = self.prepare_init_args_and_inputs_for_common() + model = self.model_class(**init_dict).to(torch_device) + + # Prepare inputs with txt_seq_lens (deprecated parameter) + txt_seq_lens = [inputs["encoder_hidden_states"].shape[1]] + + # Remove encoder_hidden_states_mask to use the deprecated path + inputs_with_deprecated = inputs.copy() + inputs_with_deprecated.pop("encoder_hidden_states_mask") + inputs_with_deprecated["txt_seq_lens"] = txt_seq_lens + + # Test that deprecation warning is raised + with self.assertWarns(FutureWarning) as warning_context: + with torch.no_grad(): + output = model(**inputs_with_deprecated) + + # Verify the warning message mentions the deprecation + warning_message = str(warning_context.warning) + self.assertIn("txt_seq_lens", warning_message) + self.assertIn("deprecated", warning_message) + self.assertIn("encoder_hidden_states_mask", warning_message) + + # Verify the model still works correctly despite the deprecation + self.assertEqual(output.sample.shape[1], inputs["hidden_states"].shape[1]) + + def test_layered_model_with_mask(self): + """Test QwenImageTransformer2DModel with use_layer3d_rope=True (layered model).""" + # Create layered model config + init_dict = { + "patch_size": 2, + "in_channels": 16, + "out_channels": 4, + "num_layers": 2, + "attention_head_dim": 16, + "num_attention_heads": 3, + "joint_attention_dim": 16, + "axes_dims_rope": (8, 4, 4), # Must match attention_head_dim (8+4+4=16) + "use_layer3d_rope": True, # Enable layered RoPE + "use_additional_t_cond": True, # Enable additional time conditioning + } + + model = self.model_class(**init_dict).to(torch_device) + + # Verify the model uses QwenEmbedLayer3DRope + from diffusers.models.transformers.transformer_qwenimage import QwenEmbedLayer3DRope + + self.assertIsInstance(model.pos_embed, QwenEmbedLayer3DRope) + + # Test single generation with layered structure + batch_size = 1 + text_seq_len = 7 + img_h, img_w = 4, 4 + layers = 4 + + # For layered model: (layers + 1) because we have N layers + 1 combined image + hidden_states = torch.randn(batch_size, (layers + 1) * img_h * img_w, 16).to(torch_device) + encoder_hidden_states = torch.randn(batch_size, text_seq_len, 16).to(torch_device) + + # Create mask with some padding + encoder_hidden_states_mask = torch.ones(batch_size, text_seq_len).to(torch_device) + encoder_hidden_states_mask[0, 5:] = 0 # Only 5 valid tokens + + timestep = torch.tensor([1.0]).to(torch_device) + + # additional_t_cond for use_additional_t_cond=True (0 or 1 index for embedding) + addition_t_cond = torch.tensor([0], dtype=torch.long).to(torch_device) + + # Layer structure: 4 layers + 1 condition image + img_shapes = [ + [ + (1, img_h, img_w), # layer 0 + (1, img_h, img_w), # layer 1 + (1, img_h, img_w), # layer 2 + (1, img_h, img_w), # layer 3 + (1, img_h, img_w), # condition image (last one gets special treatment) + ] + ] + + with torch.no_grad(): + output = model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states_mask=encoder_hidden_states_mask, + timestep=timestep, + img_shapes=img_shapes, + additional_t_cond=addition_t_cond, + ) + + self.assertEqual(output.sample.shape[1], hidden_states.shape[1]) + class QwenImageTransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase): model_class = QwenImageTransformer2DModel @@ -101,6 +274,76 @@ def prepare_init_args_and_inputs_for_common(self): def prepare_dummy_input(self, height, width): return QwenImageTransformerTests().prepare_dummy_input(height=height, width=width) - @pytest.mark.xfail(condition=True, reason="RoPE needs to be revisited.", strict=True) def test_torch_compile_recompilation_and_graph_break(self): super().test_torch_compile_recompilation_and_graph_break() + + def test_torch_compile_with_and_without_mask(self): + """Test that torch.compile works with both None mask and padding mask.""" + init_dict, inputs = self.prepare_init_args_and_inputs_for_common() + model = self.model_class(**init_dict).to(torch_device) + model.eval() + model.compile(mode="default", fullgraph=True) + + # Test 1: Run with None mask (no padding, all tokens are valid) + inputs_no_mask = inputs.copy() + inputs_no_mask["encoder_hidden_states_mask"] = None + + # First run to allow compilation + with torch.no_grad(): + output_no_mask = model(**inputs_no_mask) + + # Second run to verify no recompilation + with ( + torch._inductor.utils.fresh_inductor_cache(), + torch._dynamo.config.patch(error_on_recompile=True), + torch.no_grad(), + ): + output_no_mask_2 = model(**inputs_no_mask) + + self.assertEqual(output_no_mask.sample.shape[1], inputs["hidden_states"].shape[1]) + self.assertEqual(output_no_mask_2.sample.shape[1], inputs["hidden_states"].shape[1]) + + # Test 2: Run with all-ones mask (should behave like None) + inputs_all_ones = inputs.copy() + # Keep the all-ones mask + self.assertTrue(inputs_all_ones["encoder_hidden_states_mask"].all().item()) + + # First run to allow compilation + with torch.no_grad(): + output_all_ones = model(**inputs_all_ones) + + # Second run to verify no recompilation + with ( + torch._inductor.utils.fresh_inductor_cache(), + torch._dynamo.config.patch(error_on_recompile=True), + torch.no_grad(), + ): + output_all_ones_2 = model(**inputs_all_ones) + + self.assertEqual(output_all_ones.sample.shape[1], inputs["hidden_states"].shape[1]) + self.assertEqual(output_all_ones_2.sample.shape[1], inputs["hidden_states"].shape[1]) + + # Test 3: Run with actual padding mask (has zeros) + inputs_with_padding = inputs.copy() + mask_with_padding = inputs["encoder_hidden_states_mask"].clone() + mask_with_padding[:, 4:] = 0 # Last 3 tokens are padding + + inputs_with_padding["encoder_hidden_states_mask"] = mask_with_padding + + # First run to allow compilation + with torch.no_grad(): + output_with_padding = model(**inputs_with_padding) + + # Second run to verify no recompilation + with ( + torch._inductor.utils.fresh_inductor_cache(), + torch._dynamo.config.patch(error_on_recompile=True), + torch.no_grad(), + ): + output_with_padding_2 = model(**inputs_with_padding) + + self.assertEqual(output_with_padding.sample.shape[1], inputs["hidden_states"].shape[1]) + self.assertEqual(output_with_padding_2.sample.shape[1], inputs["hidden_states"].shape[1]) + + # Verify that outputs are different (mask should affect results) + self.assertFalse(torch.allclose(output_no_mask.sample, output_with_padding.sample, atol=1e-3)) diff --git a/tests/models/transformers/test_models_transformer_wan.py b/tests/models/transformers/test_models_transformer_wan.py index 9f248f990c8a..26b0ac946434 100644 --- a/tests/models/transformers/test_models_transformer_wan.py +++ b/tests/models/transformers/test_models_transformer_wan.py @@ -12,57 +12,57 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest - +import pytest import torch from diffusers import WanTransformer3DModel - -from ...testing_utils import ( - enable_full_determinism, - torch_device, +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + BitsAndBytesTesterMixin, + GGUFCompileTesterMixin, + GGUFTesterMixin, + MemoryTesterMixin, + ModelTesterMixin, + TorchAoTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, ) -from ..test_modeling_common import ModelTesterMixin, TorchCompileTesterMixin enable_full_determinism() -class WanTransformer3DTests(ModelTesterMixin, unittest.TestCase): - model_class = WanTransformer3DModel - main_input_name = "hidden_states" - uses_custom_attn_processor = True +class WanTransformer3DTesterConfig(BaseModelTesterConfig): + @property + def model_class(self): + return WanTransformer3DModel @property - def dummy_input(self): - batch_size = 1 - num_channels = 4 - num_frames = 2 - height = 16 - width = 16 - text_encoder_embedding_dim = 16 - sequence_length = 12 + def pretrained_model_name_or_path(self): + return "hf-internal-testing/tiny-wan22-transformer" - hidden_states = torch.randn((batch_size, num_channels, num_frames, height, width)).to(torch_device) - timestep = torch.randint(0, 1000, size=(batch_size,)).to(torch_device) - encoder_hidden_states = torch.randn((batch_size, sequence_length, text_encoder_embedding_dim)).to(torch_device) + @property + def output_shape(self) -> tuple[int, ...]: + return (4, 2, 16, 16) - return { - "hidden_states": hidden_states, - "encoder_hidden_states": encoder_hidden_states, - "timestep": timestep, - } + @property + def input_shape(self) -> tuple[int, ...]: + return (4, 2, 16, 16) @property - def input_shape(self): - return (4, 1, 16, 16) + def main_input_name(self) -> str: + return "hidden_states" @property - def output_shape(self): - return (4, 1, 16, 16) + def generator(self): + return torch.Generator("cpu").manual_seed(0) - def prepare_init_args_and_inputs_for_common(self): - init_dict = { + def get_init_dict(self) -> dict[str, int | list[int] | tuple | str | bool]: + return { "patch_size": (1, 2, 2), "num_attention_heads": 2, "attention_head_dim": 12, @@ -76,16 +76,160 @@ def prepare_init_args_and_inputs_for_common(self): "qk_norm": "rms_norm_across_heads", "rope_max_seq_len": 32, } - inputs_dict = self.dummy_input - return init_dict, inputs_dict + + def get_dummy_inputs(self) -> dict[str, torch.Tensor]: + batch_size = 1 + num_channels = 4 + num_frames = 2 + height = 16 + width = 16 + text_encoder_embedding_dim = 16 + sequence_length = 12 + + return { + "hidden_states": randn_tensor( + (batch_size, num_channels, num_frames, height, width), + generator=self.generator, + device=torch_device, + ), + "encoder_hidden_states": randn_tensor( + (batch_size, sequence_length, text_encoder_embedding_dim), + generator=self.generator, + device=torch_device, + ), + "timestep": torch.randint(0, 1000, size=(batch_size,), generator=self.generator).to(torch_device), + } + + +class TestWanTransformer3D(WanTransformer3DTesterConfig, ModelTesterMixin): + """Core model tests for Wan Transformer 3D.""" + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) + def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype): + # Skip: fp16/bf16 require very high atol to pass, providing little signal. + # Dtype preservation is already tested by test_from_save_pretrained_dtype and test_keep_in_fp32_modules. + pytest.skip("Tolerance requirements too high for meaningful test") + + +class TestWanTransformer3DMemory(WanTransformer3DTesterConfig, MemoryTesterMixin): + """Memory optimization tests for Wan Transformer 3D.""" + + +class TestWanTransformer3DTraining(WanTransformer3DTesterConfig, TrainingTesterMixin): + """Training tests for Wan Transformer 3D.""" def test_gradient_checkpointing_is_applied(self): expected_set = {"WanTransformer3DModel"} super().test_gradient_checkpointing_is_applied(expected_set=expected_set) -class WanTransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase): - model_class = WanTransformer3DModel +class TestWanTransformer3DAttention(WanTransformer3DTesterConfig, AttentionTesterMixin): + """Attention processor tests for Wan Transformer 3D.""" + + +class TestWanTransformer3DCompile(WanTransformer3DTesterConfig, TorchCompileTesterMixin): + """Torch compile tests for Wan Transformer 3D.""" - def prepare_init_args_and_inputs_for_common(self): - return WanTransformer3DTests().prepare_init_args_and_inputs_for_common() + +class TestWanTransformer3DBitsAndBytes(WanTransformer3DTesterConfig, BitsAndBytesTesterMixin): + """BitsAndBytes quantization tests for Wan Transformer 3D.""" + + @property + def torch_dtype(self): + return torch.float16 + + def get_dummy_inputs(self): + """Override to provide inputs matching the tiny Wan model dimensions.""" + return { + "hidden_states": randn_tensor( + (1, 36, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "encoder_hidden_states": randn_tensor( + (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype), + } + + +class TestWanTransformer3DTorchAo(WanTransformer3DTesterConfig, TorchAoTesterMixin): + """TorchAO quantization tests for Wan Transformer 3D.""" + + @property + def torch_dtype(self): + return torch.bfloat16 + + def get_dummy_inputs(self): + """Override to provide inputs matching the tiny Wan model dimensions.""" + return { + "hidden_states": randn_tensor( + (1, 36, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "encoder_hidden_states": randn_tensor( + (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype), + } + + +class TestWanTransformer3DGGUF(WanTransformer3DTesterConfig, GGUFTesterMixin): + """GGUF quantization tests for Wan Transformer 3D.""" + + @property + def gguf_filename(self): + return "https://huggingface.co/QuantStack/Wan2.2-I2V-A14B-GGUF/blob/main/LowNoise/Wan2.2-I2V-A14B-LowNoise-Q2_K.gguf" + + @property + def torch_dtype(self): + return torch.bfloat16 + + def _create_quantized_model(self, config_kwargs=None, **extra_kwargs): + return super()._create_quantized_model( + config_kwargs, config="Wan-AI/Wan2.2-I2V-A14B-Diffusers", subfolder="transformer", **extra_kwargs + ) + + def get_dummy_inputs(self): + """Override to provide inputs matching the real Wan I2V model dimensions. + + Wan 2.2 I2V: in_channels=36, text_dim=4096 + """ + return { + "hidden_states": randn_tensor( + (1, 36, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "encoder_hidden_states": randn_tensor( + (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype), + } + + +class TestWanTransformer3DGGUFCompile(WanTransformer3DTesterConfig, GGUFCompileTesterMixin): + """GGUF + compile tests for Wan Transformer 3D.""" + + @property + def gguf_filename(self): + return "https://huggingface.co/QuantStack/Wan2.2-I2V-A14B-GGUF/blob/main/LowNoise/Wan2.2-I2V-A14B-LowNoise-Q2_K.gguf" + + @property + def torch_dtype(self): + return torch.bfloat16 + + def _create_quantized_model(self, config_kwargs=None, **extra_kwargs): + return super()._create_quantized_model( + config_kwargs, config="Wan-AI/Wan2.2-I2V-A14B-Diffusers", subfolder="transformer", **extra_kwargs + ) + + def get_dummy_inputs(self): + """Override to provide inputs matching the real Wan I2V model dimensions. + + Wan 2.2 I2V: in_channels=36, text_dim=4096 + """ + return { + "hidden_states": randn_tensor( + (1, 36, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "encoder_hidden_states": randn_tensor( + (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype), + } diff --git a/tests/models/transformers/test_models_transformer_wan_animate.py b/tests/models/transformers/test_models_transformer_wan_animate.py index 5d571b8c2e7d..ac0ef0698c63 100644 --- a/tests/models/transformers/test_models_transformer_wan_animate.py +++ b/tests/models/transformers/test_models_transformer_wan_animate.py @@ -12,76 +12,62 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest - +import pytest import torch from diffusers import WanAnimateTransformer3DModel +from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import ( - enable_full_determinism, - torch_device, +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + BitsAndBytesTesterMixin, + GGUFCompileTesterMixin, + GGUFTesterMixin, + MemoryTesterMixin, + ModelTesterMixin, + TorchAoTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, ) -from ..test_modeling_common import ModelTesterMixin, TorchCompileTesterMixin enable_full_determinism() -class WanAnimateTransformer3DTests(ModelTesterMixin, unittest.TestCase): - model_class = WanAnimateTransformer3DModel - main_input_name = "hidden_states" - uses_custom_attn_processor = True - +class WanAnimateTransformer3DTesterConfig(BaseModelTesterConfig): @property - def dummy_input(self): - batch_size = 1 - num_channels = 4 - num_frames = 20 # To make the shapes work out; for complicated reasons we want 21 to divide num_frames + 1 - height = 16 - width = 16 - text_encoder_embedding_dim = 16 - sequence_length = 12 - - clip_seq_len = 12 - clip_dim = 16 + def model_class(self): + return WanAnimateTransformer3DModel - inference_segment_length = 77 # The inference segment length in the full Wan2.2-Animate-14B model - face_height = 16 # Should be square and match `motion_encoder_size` below - face_width = 16 + @property + def pretrained_model_name_or_path(self): + return "hf-internal-testing/tiny-wan-animate-transformer" - hidden_states = torch.randn((batch_size, 2 * num_channels + 4, num_frames + 1, height, width)).to(torch_device) - timestep = torch.randint(0, 1000, size=(batch_size,)).to(torch_device) - encoder_hidden_states = torch.randn((batch_size, sequence_length, text_encoder_embedding_dim)).to(torch_device) - clip_ref_features = torch.randn((batch_size, clip_seq_len, clip_dim)).to(torch_device) - pose_latents = torch.randn((batch_size, num_channels, num_frames, height, width)).to(torch_device) - face_pixel_values = torch.randn((batch_size, 3, inference_segment_length, face_height, face_width)).to( - torch_device - ) + @property + def output_shape(self) -> tuple[int, ...]: + # Output has fewer channels than input (4 vs 12) + return (4, 21, 16, 16) - return { - "hidden_states": hidden_states, - "timestep": timestep, - "encoder_hidden_states": encoder_hidden_states, - "encoder_hidden_states_image": clip_ref_features, - "pose_hidden_states": pose_latents, - "face_pixel_values": face_pixel_values, - } + @property + def input_shape(self) -> tuple[int, ...]: + return (12, 21, 16, 16) @property - def input_shape(self): - return (12, 1, 16, 16) + def main_input_name(self) -> str: + return "hidden_states" @property - def output_shape(self): - return (4, 1, 16, 16) + def generator(self): + return torch.Generator("cpu").manual_seed(0) - def prepare_init_args_and_inputs_for_common(self): + def get_init_dict(self) -> dict[str, int | list[int] | tuple | str | bool | float | dict]: # Use custom channel sizes since the default Wan Animate channel sizes will cause the motion encoder to # contain the vast majority of the parameters in the test model channel_sizes = {"4": 16, "8": 16, "16": 16} - init_dict = { + return { "patch_size": (1, 2, 2), "num_attention_heads": 2, "attention_head_dim": 12, @@ -105,22 +91,219 @@ def prepare_init_args_and_inputs_for_common(self): "face_encoder_num_heads": 2, "inject_face_latents_blocks": 2, } - inputs_dict = self.dummy_input - return init_dict, inputs_dict - def test_gradient_checkpointing_is_applied(self): - expected_set = {"WanAnimateTransformer3DModel"} - super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + def get_dummy_inputs(self) -> dict[str, torch.Tensor]: + batch_size = 1 + num_channels = 4 + num_frames = 20 # To make the shapes work out; for complicated reasons we want 21 to divide num_frames + 1 + height = 16 + width = 16 + text_encoder_embedding_dim = 16 + sequence_length = 12 + + clip_seq_len = 12 + clip_dim = 16 + + inference_segment_length = 77 # The inference segment length in the full Wan2.2-Animate-14B model + face_height = 16 # Should be square and match `motion_encoder_size` + face_width = 16 + + return { + "hidden_states": randn_tensor( + (batch_size, 2 * num_channels + 4, num_frames + 1, height, width), + generator=self.generator, + device=torch_device, + ), + "timestep": torch.randint(0, 1000, size=(batch_size,), generator=self.generator).to(torch_device), + "encoder_hidden_states": randn_tensor( + (batch_size, sequence_length, text_encoder_embedding_dim), + generator=self.generator, + device=torch_device, + ), + "encoder_hidden_states_image": randn_tensor( + (batch_size, clip_seq_len, clip_dim), + generator=self.generator, + device=torch_device, + ), + "pose_hidden_states": randn_tensor( + (batch_size, num_channels, num_frames, height, width), + generator=self.generator, + device=torch_device, + ), + "face_pixel_values": randn_tensor( + (batch_size, 3, inference_segment_length, face_height, face_width), + generator=self.generator, + device=torch_device, + ), + } + + +class TestWanAnimateTransformer3D(WanAnimateTransformer3DTesterConfig, ModelTesterMixin): + """Core model tests for Wan Animate Transformer 3D.""" - # Override test_output because the transformer output is expected to have less channels than the main transformer - # input. def test_output(self): + # Override test_output because the transformer output is expected to have less channels + # than the main transformer input. expected_output_shape = (1, 4, 21, 16, 16) super().test_output(expected_output_shape=expected_output_shape) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) + def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype): + # Skip: fp16/bf16 require very high atol (~1e-2) to pass, providing little signal. + # Dtype preservation is already tested by test_from_save_pretrained_dtype and test_keep_in_fp32_modules. + pytest.skip("Tolerance requirements too high for meaningful test") + + +class TestWanAnimateTransformer3DMemory(WanAnimateTransformer3DTesterConfig, MemoryTesterMixin): + """Memory optimization tests for Wan Animate Transformer 3D.""" -class WanAnimateTransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase): - model_class = WanAnimateTransformer3DModel - def prepare_init_args_and_inputs_for_common(self): - return WanAnimateTransformer3DTests().prepare_init_args_and_inputs_for_common() +class TestWanAnimateTransformer3DTraining(WanAnimateTransformer3DTesterConfig, TrainingTesterMixin): + """Training tests for Wan Animate Transformer 3D.""" + + def test_gradient_checkpointing_is_applied(self): + expected_set = {"WanAnimateTransformer3DModel"} + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + + +class TestWanAnimateTransformer3DAttention(WanAnimateTransformer3DTesterConfig, AttentionTesterMixin): + """Attention processor tests for Wan Animate Transformer 3D.""" + + +class TestWanAnimateTransformer3DCompile(WanAnimateTransformer3DTesterConfig, TorchCompileTesterMixin): + """Torch compile tests for Wan Animate Transformer 3D.""" + + def test_torch_compile_recompilation_and_graph_break(self): + # Skip: F.pad with mode="replicate" in WanAnimateFaceEncoder triggers importlib.import_module + # internally, which dynamo doesn't support tracing through. + pytest.skip("F.pad with replicate mode triggers unsupported import in torch.compile") + + +class TestWanAnimateTransformer3DBitsAndBytes(WanAnimateTransformer3DTesterConfig, BitsAndBytesTesterMixin): + """BitsAndBytes quantization tests for Wan Animate Transformer 3D.""" + + @property + def torch_dtype(self): + return torch.float16 + + def get_dummy_inputs(self): + """Override to provide inputs matching the tiny Wan Animate model dimensions.""" + return { + "hidden_states": randn_tensor( + (1, 36, 21, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "encoder_hidden_states": randn_tensor( + (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "encoder_hidden_states_image": randn_tensor( + (1, 257, 1280), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "pose_hidden_states": randn_tensor( + (1, 16, 20, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "face_pixel_values": randn_tensor( + (1, 3, 77, 512, 512), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype), + } + + +class TestWanAnimateTransformer3DTorchAo(WanAnimateTransformer3DTesterConfig, TorchAoTesterMixin): + """TorchAO quantization tests for Wan Animate Transformer 3D.""" + + @property + def torch_dtype(self): + return torch.bfloat16 + + def get_dummy_inputs(self): + """Override to provide inputs matching the tiny Wan Animate model dimensions.""" + return { + "hidden_states": randn_tensor( + (1, 36, 21, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "encoder_hidden_states": randn_tensor( + (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "encoder_hidden_states_image": randn_tensor( + (1, 257, 1280), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "pose_hidden_states": randn_tensor( + (1, 16, 20, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "face_pixel_values": randn_tensor( + (1, 3, 77, 512, 512), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype), + } + + +class TestWanAnimateTransformer3DGGUF(WanAnimateTransformer3DTesterConfig, GGUFTesterMixin): + """GGUF quantization tests for Wan Animate Transformer 3D.""" + + @property + def gguf_filename(self): + return "https://huggingface.co/QuantStack/Wan2.2-Animate-14B-GGUF/blob/main/Wan2.2-Animate-14B-Q2_K.gguf" + + @property + def torch_dtype(self): + return torch.bfloat16 + + def get_dummy_inputs(self): + """Override to provide inputs matching the real Wan Animate model dimensions. + + Wan 2.2 Animate: in_channels=36 (2*16+4), text_dim=4096, image_dim=1280 + """ + return { + "hidden_states": randn_tensor( + (1, 36, 21, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "encoder_hidden_states": randn_tensor( + (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "encoder_hidden_states_image": randn_tensor( + (1, 257, 1280), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "pose_hidden_states": randn_tensor( + (1, 16, 20, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "face_pixel_values": randn_tensor( + (1, 3, 77, 512, 512), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype), + } + + +class TestWanAnimateTransformer3DGGUFCompile(WanAnimateTransformer3DTesterConfig, GGUFCompileTesterMixin): + """GGUF + compile tests for Wan Animate Transformer 3D.""" + + @property + def gguf_filename(self): + return "https://huggingface.co/QuantStack/Wan2.2-Animate-14B-GGUF/blob/main/Wan2.2-Animate-14B-Q2_K.gguf" + + @property + def torch_dtype(self): + return torch.bfloat16 + + def get_dummy_inputs(self): + """Override to provide inputs matching the real Wan Animate model dimensions. + + Wan 2.2 Animate: in_channels=36 (2*16+4), text_dim=4096, image_dim=1280 + """ + return { + "hidden_states": randn_tensor( + (1, 36, 21, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "encoder_hidden_states": randn_tensor( + (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "encoder_hidden_states_image": randn_tensor( + (1, 257, 1280), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "pose_hidden_states": randn_tensor( + (1, 16, 20, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "face_pixel_values": randn_tensor( + (1, 3, 77, 512, 512), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype), + } diff --git a/tests/models/transformers/test_models_transformer_wan_vace.py b/tests/models/transformers/test_models_transformer_wan_vace.py new file mode 100644 index 000000000000..5ab51bbb9003 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_wan_vace.py @@ -0,0 +1,256 @@ +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers import WanVACETransformer3DModel +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + BitsAndBytesTesterMixin, + GGUFCompileTesterMixin, + GGUFTesterMixin, + MemoryTesterMixin, + ModelTesterMixin, + TorchAoTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, +) + + +enable_full_determinism() + + +class WanVACETransformer3DTesterConfig(BaseModelTesterConfig): + @property + def model_class(self): + return WanVACETransformer3DModel + + @property + def pretrained_model_name_or_path(self): + return "hf-internal-testing/tiny-wan-vace-transformer" + + @property + def output_shape(self) -> tuple[int, ...]: + return (16, 2, 16, 16) + + @property + def input_shape(self) -> tuple[int, ...]: + return (16, 2, 16, 16) + + @property + def main_input_name(self) -> str: + return "hidden_states" + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict[str, int | list[int] | tuple | str | bool | None]: + return { + "patch_size": (1, 2, 2), + "num_attention_heads": 2, + "attention_head_dim": 12, + "in_channels": 16, + "out_channels": 16, + "text_dim": 32, + "freq_dim": 256, + "ffn_dim": 32, + "num_layers": 4, + "cross_attn_norm": True, + "qk_norm": "rms_norm_across_heads", + "rope_max_seq_len": 32, + "vace_layers": [0, 2], + "vace_in_channels": 48, # 3 * in_channels = 3 * 16 = 48 + } + + def get_dummy_inputs(self) -> dict[str, torch.Tensor]: + batch_size = 1 + num_channels = 16 + num_frames = 2 + height = 16 + width = 16 + text_encoder_embedding_dim = 32 + sequence_length = 12 + + # VACE requires control_hidden_states with vace_in_channels (3 * in_channels) + vace_in_channels = 48 + + return { + "hidden_states": randn_tensor( + (batch_size, num_channels, num_frames, height, width), + generator=self.generator, + device=torch_device, + ), + "encoder_hidden_states": randn_tensor( + (batch_size, sequence_length, text_encoder_embedding_dim), + generator=self.generator, + device=torch_device, + ), + "control_hidden_states": randn_tensor( + (batch_size, vace_in_channels, num_frames, height, width), + generator=self.generator, + device=torch_device, + ), + "timestep": torch.randint(0, 1000, size=(batch_size,), generator=self.generator).to(torch_device), + } + + +class TestWanVACETransformer3D(WanVACETransformer3DTesterConfig, ModelTesterMixin): + """Core model tests for Wan VACE Transformer 3D.""" + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) + def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype): + # Skip: fp16/bf16 require very high atol to pass, providing little signal. + # Dtype preservation is already tested by test_from_save_pretrained_dtype and test_keep_in_fp32_modules. + pytest.skip("Tolerance requirements too high for meaningful test") + + def test_model_parallelism(self, tmp_path): + # Skip: Device mismatch between cuda:0 and cuda:1 in VACE control flow + pytest.skip("Model parallelism not yet supported for WanVACE") + + +class TestWanVACETransformer3DMemory(WanVACETransformer3DTesterConfig, MemoryTesterMixin): + """Memory optimization tests for Wan VACE Transformer 3D.""" + + +class TestWanVACETransformer3DTraining(WanVACETransformer3DTesterConfig, TrainingTesterMixin): + """Training tests for Wan VACE Transformer 3D.""" + + def test_gradient_checkpointing_is_applied(self): + expected_set = {"WanVACETransformer3DModel"} + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + + +class TestWanVACETransformer3DAttention(WanVACETransformer3DTesterConfig, AttentionTesterMixin): + """Attention processor tests for Wan VACE Transformer 3D.""" + + +class TestWanVACETransformer3DCompile(WanVACETransformer3DTesterConfig, TorchCompileTesterMixin): + """Torch compile tests for Wan VACE Transformer 3D.""" + + def test_torch_compile_repeated_blocks(self): + # WanVACE has two block types (WanTransformerBlock and WanVACETransformerBlock), + # so we need recompile_limit=2 instead of the default 1. + super().test_torch_compile_repeated_blocks(recompile_limit=2) + + +class TestWanVACETransformer3DBitsAndBytes(WanVACETransformer3DTesterConfig, BitsAndBytesTesterMixin): + """BitsAndBytes quantization tests for Wan VACE Transformer 3D.""" + + @property + def torch_dtype(self): + return torch.float16 + + def get_dummy_inputs(self): + """Override to provide inputs matching the tiny Wan VACE model dimensions.""" + return { + "hidden_states": randn_tensor( + (1, 16, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "encoder_hidden_states": randn_tensor( + (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "control_hidden_states": randn_tensor( + (1, 96, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype), + } + + +class TestWanVACETransformer3DTorchAo(WanVACETransformer3DTesterConfig, TorchAoTesterMixin): + """TorchAO quantization tests for Wan VACE Transformer 3D.""" + + @property + def torch_dtype(self): + return torch.bfloat16 + + def get_dummy_inputs(self): + """Override to provide inputs matching the tiny Wan VACE model dimensions.""" + return { + "hidden_states": randn_tensor( + (1, 16, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "encoder_hidden_states": randn_tensor( + (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "control_hidden_states": randn_tensor( + (1, 96, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype), + } + + +class TestWanVACETransformer3DGGUF(WanVACETransformer3DTesterConfig, GGUFTesterMixin): + """GGUF quantization tests for Wan VACE Transformer 3D.""" + + @property + def gguf_filename(self): + return "https://huggingface.co/QuantStack/Wan2.1_14B_VACE-GGUF/blob/main/Wan2.1_14B_VACE-Q3_K_S.gguf" + + @property + def torch_dtype(self): + return torch.bfloat16 + + def get_dummy_inputs(self): + """Override to provide inputs matching the real Wan VACE model dimensions. + + Wan 2.1 VACE: in_channels=16, text_dim=4096, vace_in_channels=96 + """ + return { + "hidden_states": randn_tensor( + (1, 16, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "encoder_hidden_states": randn_tensor( + (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "control_hidden_states": randn_tensor( + (1, 96, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype), + } + + +class TestWanVACETransformer3DGGUFCompile(WanVACETransformer3DTesterConfig, GGUFCompileTesterMixin): + """GGUF + compile tests for Wan VACE Transformer 3D.""" + + @property + def gguf_filename(self): + return "https://huggingface.co/QuantStack/Wan2.1_14B_VACE-GGUF/blob/main/Wan2.1_14B_VACE-Q3_K_S.gguf" + + @property + def torch_dtype(self): + return torch.bfloat16 + + def get_dummy_inputs(self): + """Override to provide inputs matching the real Wan VACE model dimensions. + + Wan 2.1 VACE: in_channels=16, text_dim=4096, vace_in_channels=96 + """ + return { + "hidden_states": randn_tensor( + (1, 16, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "encoder_hidden_states": randn_tensor( + (1, 512, 4096), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "control_hidden_states": randn_tensor( + (1, 96, 2, 64, 64), generator=self.generator, device=torch_device, dtype=self.torch_dtype + ), + "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype), + } diff --git a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py index 854b5218c617..9a6b4b9b6fb4 100644 --- a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py +++ b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py @@ -33,6 +33,19 @@ from ..test_modular_pipelines_common import ModularPipelineTesterMixin +FLUX_TEXT2IMAGE_WORKFLOWS = { + "text2image": [ + ("text_encoder", "FluxTextEncoderStep"), + ("denoise.input", "FluxTextInputStep"), + ("denoise.before_denoise.prepare_latents", "FluxPrepareLatentsStep"), + ("denoise.before_denoise.set_timesteps", "FluxSetTimestepsStep"), + ("denoise.before_denoise.prepare_rope_inputs", "FluxRoPEInputsStep"), + ("denoise.denoise", "FluxDenoiseStep"), + ("decode", "FluxDecodeStep"), + ] +} + + class TestFluxModularPipelineFast(ModularPipelineTesterMixin): pipeline_class = FluxModularPipeline pipeline_blocks_class = FluxAutoBlocks @@ -40,6 +53,7 @@ class TestFluxModularPipelineFast(ModularPipelineTesterMixin): params = frozenset(["prompt", "height", "width", "guidance_scale"]) batch_params = frozenset(["prompt"]) + expected_workflow_blocks = FLUX_TEXT2IMAGE_WORKFLOWS def get_dummy_inputs(self, seed=0): generator = self.get_generator(seed) @@ -59,6 +73,23 @@ def test_float16_inference(self): super().test_float16_inference(9e-2) +FLUX_IMAGE2IMAGE_WORKFLOWS = { + "image2image": [ + ("text_encoder", "FluxTextEncoderStep"), + ("vae_encoder.preprocess", "FluxProcessImagesInputStep"), + ("vae_encoder.encode", "FluxVaeEncoderStep"), + ("denoise.input.text_inputs", "FluxTextInputStep"), + ("denoise.input.additional_inputs", "FluxAdditionalInputsStep"), + ("denoise.before_denoise.prepare_latents", "FluxPrepareLatentsStep"), + ("denoise.before_denoise.set_timesteps", "FluxImg2ImgSetTimestepsStep"), + ("denoise.before_denoise.prepare_img2img_latents", "FluxImg2ImgPrepareLatentsStep"), + ("denoise.before_denoise.prepare_rope_inputs", "FluxRoPEInputsStep"), + ("denoise.denoise", "FluxDenoiseStep"), + ("decode", "FluxDecodeStep"), + ] +} + + class TestFluxImg2ImgModularPipelineFast(ModularPipelineTesterMixin): pipeline_class = FluxModularPipeline pipeline_blocks_class = FluxAutoBlocks @@ -66,6 +97,7 @@ class TestFluxImg2ImgModularPipelineFast(ModularPipelineTesterMixin): params = frozenset(["prompt", "height", "width", "guidance_scale", "image"]) batch_params = frozenset(["prompt", "image"]) + expected_workflow_blocks = FLUX_IMAGE2IMAGE_WORKFLOWS def get_pipeline(self, components_manager=None, torch_dtype=torch.float32): pipeline = super().get_pipeline(components_manager, torch_dtype) @@ -125,6 +157,32 @@ def test_float16_inference(self): super().test_float16_inference(8e-2) +FLUX_KONTEXT_WORKFLOWS = { + "text2image": [ + ("text_encoder", "FluxTextEncoderStep"), + ("denoise.input", "FluxTextInputStep"), + ("denoise.before_denoise.prepare_latents", "FluxPrepareLatentsStep"), + ("denoise.before_denoise.set_timesteps", "FluxSetTimestepsStep"), + ("denoise.before_denoise.prepare_rope_inputs", "FluxRoPEInputsStep"), + ("denoise.denoise", "FluxKontextDenoiseStep"), + ("decode", "FluxDecodeStep"), + ], + "image_conditioned": [ + ("text_encoder", "FluxTextEncoderStep"), + ("vae_encoder.preprocess", "FluxKontextProcessImagesInputStep"), + ("vae_encoder.encode", "FluxVaeEncoderStep"), + ("denoise.input.set_resolution", "FluxKontextSetResolutionStep"), + ("denoise.input.text_inputs", "FluxTextInputStep"), + ("denoise.input.additional_inputs", "FluxKontextAdditionalInputsStep"), + ("denoise.before_denoise.prepare_latents", "FluxPrepareLatentsStep"), + ("denoise.before_denoise.set_timesteps", "FluxSetTimestepsStep"), + ("denoise.before_denoise.prepare_rope_inputs", "FluxKontextRoPEInputsStep"), + ("denoise.denoise", "FluxKontextDenoiseStep"), + ("decode", "FluxDecodeStep"), + ], +} + + class TestFluxKontextModularPipelineFast(ModularPipelineTesterMixin): pipeline_class = FluxKontextModularPipeline pipeline_blocks_class = FluxKontextAutoBlocks @@ -132,6 +190,7 @@ class TestFluxKontextModularPipelineFast(ModularPipelineTesterMixin): params = frozenset(["prompt", "height", "width", "guidance_scale", "image"]) batch_params = frozenset(["prompt", "image"]) + expected_workflow_blocks = FLUX_KONTEXT_WORKFLOWS def get_dummy_inputs(self, seed=0): generator = self.get_generator(seed) diff --git a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py index 8fd529e97e71..3045af636841 100644 --- a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py +++ b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py @@ -28,6 +28,21 @@ from ..test_modular_pipelines_common import ModularPipelineTesterMixin +FLUX2_TEXT2IMAGE_WORKFLOWS = { + "text2image": [ + ("text_encoder", "Flux2TextEncoderStep"), + ("denoise.input", "Flux2TextInputStep"), + ("denoise.prepare_latents", "Flux2PrepareLatentsStep"), + ("denoise.set_timesteps", "Flux2SetTimestepsStep"), + ("denoise.prepare_guidance", "Flux2PrepareGuidanceStep"), + ("denoise.prepare_rope_inputs", "Flux2RoPEInputsStep"), + ("denoise.denoise", "Flux2DenoiseStep"), + ("denoise.after_denoise", "Flux2UnpackLatentsStep"), + ("decode", "Flux2DecodeStep"), + ], +} + + class TestFlux2ModularPipelineFast(ModularPipelineTesterMixin): pipeline_class = Flux2ModularPipeline pipeline_blocks_class = Flux2AutoBlocks @@ -35,6 +50,7 @@ class TestFlux2ModularPipelineFast(ModularPipelineTesterMixin): params = frozenset(["prompt", "height", "width", "guidance_scale"]) batch_params = frozenset(["prompt"]) + expected_workflow_blocks = FLUX2_TEXT2IMAGE_WORKFLOWS def get_dummy_inputs(self, seed=0): generator = self.get_generator(seed) @@ -56,6 +72,24 @@ def test_float16_inference(self): super().test_float16_inference(9e-2) +FLUX2_IMAGE_CONDITIONED_WORKFLOWS = { + "image_conditioned": [ + ("text_encoder", "Flux2TextEncoderStep"), + ("vae_encoder.preprocess", "Flux2ProcessImagesInputStep"), + ("vae_encoder.encode", "Flux2VaeEncoderStep"), + ("denoise.input", "Flux2TextInputStep"), + ("denoise.prepare_image_latents", "Flux2PrepareImageLatentsStep"), + ("denoise.prepare_latents", "Flux2PrepareLatentsStep"), + ("denoise.set_timesteps", "Flux2SetTimestepsStep"), + ("denoise.prepare_guidance", "Flux2PrepareGuidanceStep"), + ("denoise.prepare_rope_inputs", "Flux2RoPEInputsStep"), + ("denoise.denoise", "Flux2DenoiseStep"), + ("denoise.after_denoise", "Flux2UnpackLatentsStep"), + ("decode", "Flux2DecodeStep"), + ], +} + + class TestFlux2ImageConditionedModularPipelineFast(ModularPipelineTesterMixin): pipeline_class = Flux2ModularPipeline pipeline_blocks_class = Flux2AutoBlocks @@ -63,6 +97,7 @@ class TestFlux2ImageConditionedModularPipelineFast(ModularPipelineTesterMixin): params = frozenset(["prompt", "height", "width", "guidance_scale", "image"]) batch_params = frozenset(["prompt", "image"]) + expected_workflow_blocks = FLUX2_IMAGE_CONDITIONED_WORKFLOWS def get_dummy_inputs(self, seed=0): generator = self.get_generator(seed) diff --git a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py new file mode 100644 index 000000000000..ad295a961357 --- /dev/null +++ b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py @@ -0,0 +1,124 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import random + +import numpy as np +import PIL +import pytest + +from diffusers.modular_pipelines import ( + Flux2KleinAutoBlocks, + Flux2KleinModularPipeline, +) + +from ...testing_utils import floats_tensor, torch_device +from ..test_modular_pipelines_common import ModularPipelineTesterMixin + + +FLUX2_KLEIN_WORKFLOWS = { + "text2image": [ + ("text_encoder", "Flux2KleinTextEncoderStep"), + ("denoise.input", "Flux2TextInputStep"), + ("denoise.prepare_latents", "Flux2PrepareLatentsStep"), + ("denoise.set_timesteps", "Flux2SetTimestepsStep"), + ("denoise.prepare_rope_inputs", "Flux2RoPEInputsStep"), + ("denoise.denoise", "Flux2KleinDenoiseStep"), + ("denoise.after_denoise", "Flux2UnpackLatentsStep"), + ("decode", "Flux2DecodeStep"), + ], +} + + +class TestFlux2KleinModularPipelineFast(ModularPipelineTesterMixin): + pipeline_class = Flux2KleinModularPipeline + pipeline_blocks_class = Flux2KleinAutoBlocks + pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-klein-modular" + + params = frozenset(["prompt", "height", "width"]) + batch_params = frozenset(["prompt"]) + expected_workflow_blocks = FLUX2_KLEIN_WORKFLOWS + + def get_dummy_inputs(self, seed=0): + generator = self.get_generator(seed) + inputs = { + "prompt": "A painting of a squirrel eating a burger", + # TODO (Dhruv): Update text encoder config so that vocab_size matches tokenizer + "max_sequence_length": 8, # bit of a hack to workaround vocab size mismatch + "text_encoder_out_layers": (1,), + "generator": generator, + "num_inference_steps": 2, + "height": 32, + "width": 32, + "output_type": "pt", + } + return inputs + + def test_float16_inference(self): + super().test_float16_inference(9e-2) + + +FLUX2_KLEIN_IMAGE_CONDITIONED_WORKFLOWS = { + "image_conditioned": [ + ("text_encoder", "Flux2KleinTextEncoderStep"), + ("vae_encoder.preprocess", "Flux2ProcessImagesInputStep"), + ("vae_encoder.encode", "Flux2VaeEncoderStep"), + ("denoise.input", "Flux2TextInputStep"), + ("denoise.prepare_image_latents", "Flux2PrepareImageLatentsStep"), + ("denoise.prepare_latents", "Flux2PrepareLatentsStep"), + ("denoise.set_timesteps", "Flux2SetTimestepsStep"), + ("denoise.prepare_rope_inputs", "Flux2RoPEInputsStep"), + ("denoise.denoise", "Flux2KleinDenoiseStep"), + ("denoise.after_denoise", "Flux2UnpackLatentsStep"), + ("decode", "Flux2DecodeStep"), + ], +} + + +class TestFlux2KleinImageConditionedModularPipelineFast(ModularPipelineTesterMixin): + pipeline_class = Flux2KleinModularPipeline + pipeline_blocks_class = Flux2KleinAutoBlocks + pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-klein-modular" + + params = frozenset(["prompt", "height", "width", "image"]) + batch_params = frozenset(["prompt", "image"]) + expected_workflow_blocks = FLUX2_KLEIN_IMAGE_CONDITIONED_WORKFLOWS + + def get_dummy_inputs(self, seed=0): + generator = self.get_generator(seed) + inputs = { + "prompt": "A painting of a squirrel eating a burger", + # TODO (Dhruv): Update text encoder config so that vocab_size matches tokenizer + "max_sequence_length": 8, # bit of a hack to workaround vocab size mismatch + "text_encoder_out_layers": (1,), + "generator": generator, + "num_inference_steps": 2, + "height": 32, + "width": 32, + "output_type": "pt", + } + image = floats_tensor((1, 3, 64, 64), rng=random.Random(seed)).to(torch_device) + image = image.cpu().permute(0, 2, 3, 1)[0] + init_image = PIL.Image.fromarray(np.uint8(image * 255)).convert("RGB") + inputs["image"] = init_image + + return inputs + + def test_float16_inference(self): + super().test_float16_inference(9e-2) + + @pytest.mark.skip(reason="batched inference is currently not supported") + def test_inference_batch_single_identical(self, batch_size=2, expected_max_diff=0.0001): + return diff --git a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein_base.py b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein_base.py new file mode 100644 index 000000000000..b3aa79040317 --- /dev/null +++ b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein_base.py @@ -0,0 +1,124 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import random + +import numpy as np +import PIL +import pytest + +from diffusers.modular_pipelines import ( + Flux2KleinBaseAutoBlocks, + Flux2KleinBaseModularPipeline, +) + +from ...testing_utils import floats_tensor, torch_device +from ..test_modular_pipelines_common import ModularPipelineTesterMixin + + +FLUX2_KLEIN_BASE_WORKFLOWS = { + "text2image": [ + ("text_encoder", "Flux2KleinBaseTextEncoderStep"), + ("denoise.input", "Flux2KleinBaseTextInputStep"), + ("denoise.prepare_latents", "Flux2PrepareLatentsStep"), + ("denoise.set_timesteps", "Flux2SetTimestepsStep"), + ("denoise.prepare_rope_inputs", "Flux2KleinBaseRoPEInputsStep"), + ("denoise.denoise", "Flux2KleinBaseDenoiseStep"), + ("denoise.after_denoise", "Flux2UnpackLatentsStep"), + ("decode", "Flux2DecodeStep"), + ], +} + + +class TestFlux2KleinBaseModularPipelineFast(ModularPipelineTesterMixin): + pipeline_class = Flux2KleinBaseModularPipeline + pipeline_blocks_class = Flux2KleinBaseAutoBlocks + pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-klein-base-modular" + + params = frozenset(["prompt", "height", "width"]) + batch_params = frozenset(["prompt"]) + expected_workflow_blocks = FLUX2_KLEIN_BASE_WORKFLOWS + + def get_dummy_inputs(self, seed=0): + generator = self.get_generator(seed) + inputs = { + "prompt": "A painting of a squirrel eating a burger", + # TODO (Dhruv): Update text encoder config so that vocab_size matches tokenizer + "max_sequence_length": 8, # bit of a hack to workaround vocab size mismatch + "text_encoder_out_layers": (1,), + "generator": generator, + "num_inference_steps": 2, + "height": 32, + "width": 32, + "output_type": "pt", + } + return inputs + + def test_float16_inference(self): + super().test_float16_inference(9e-2) + + +FLUX2_KLEIN_BASE_IMAGE_CONDITIONED_WORKFLOWS = { + "image_conditioned": [ + ("text_encoder", "Flux2KleinBaseTextEncoderStep"), + ("vae_encoder.preprocess", "Flux2ProcessImagesInputStep"), + ("vae_encoder.encode", "Flux2VaeEncoderStep"), + ("denoise.input", "Flux2KleinBaseTextInputStep"), + ("denoise.prepare_latents", "Flux2PrepareLatentsStep"), + ("denoise.prepare_image_latents", "Flux2PrepareImageLatentsStep"), + ("denoise.set_timesteps", "Flux2SetTimestepsStep"), + ("denoise.prepare_rope_inputs", "Flux2KleinBaseRoPEInputsStep"), + ("denoise.denoise", "Flux2KleinBaseDenoiseStep"), + ("denoise.after_denoise", "Flux2UnpackLatentsStep"), + ("decode", "Flux2DecodeStep"), + ], +} + + +class TestFlux2KleinBaseImageConditionedModularPipelineFast(ModularPipelineTesterMixin): + pipeline_class = Flux2KleinBaseModularPipeline + pipeline_blocks_class = Flux2KleinBaseAutoBlocks + pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-klein-base-modular" + + params = frozenset(["prompt", "height", "width", "image"]) + batch_params = frozenset(["prompt", "image"]) + expected_workflow_blocks = FLUX2_KLEIN_BASE_IMAGE_CONDITIONED_WORKFLOWS + + def get_dummy_inputs(self, seed=0): + generator = self.get_generator(seed) + inputs = { + "prompt": "A painting of a squirrel eating a burger", + # TODO (Dhruv): Update text encoder config so that vocab_size matches tokenizer + "max_sequence_length": 8, # bit of a hack to workaround vocab size mismatch + "text_encoder_out_layers": (1,), + "generator": generator, + "num_inference_steps": 2, + "height": 32, + "width": 32, + "output_type": "pt", + } + image = floats_tensor((1, 3, 64, 64), rng=random.Random(seed)).to(torch_device) + image = image.cpu().permute(0, 2, 3, 1)[0] + init_image = PIL.Image.fromarray(np.uint8(image * 255)).convert("RGB") + inputs["image"] = init_image + + return inputs + + def test_float16_inference(self): + super().test_float16_inference(9e-2) + + @pytest.mark.skip(reason="batched inference is currently not supported") + def test_inference_batch_single_identical(self, batch_size=2, expected_max_diff=0.0001): + return diff --git a/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py b/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py index f4bd27b7ea47..92573c202e49 100644 --- a/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py +++ b/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py @@ -30,6 +30,103 @@ from ..test_modular_pipelines_common import ModularGuiderTesterMixin, ModularPipelineTesterMixin +QWEN_IMAGE_TEXT2IMAGE_WORKFLOWS = { + "text2image": [ + ("text_encoder", "QwenImageTextEncoderStep"), + ("denoise.input", "QwenImageTextInputsStep"), + ("denoise.prepare_latents", "QwenImagePrepareLatentsStep"), + ("denoise.set_timesteps", "QwenImageSetTimestepsStep"), + ("denoise.prepare_rope_inputs", "QwenImageRoPEInputsStep"), + ("denoise.denoise", "QwenImageDenoiseStep"), + ("denoise.after_denoise", "QwenImageAfterDenoiseStep"), + ("decode.decode", "QwenImageDecoderStep"), + ("decode.postprocess", "QwenImageProcessImagesOutputStep"), + ], + "image2image": [ + ("text_encoder", "QwenImageTextEncoderStep"), + ("vae_encoder.preprocess", "QwenImageProcessImagesInputStep"), + ("vae_encoder.encode", "QwenImageVaeEncoderStep"), + ("denoise.input.text_inputs", "QwenImageTextInputsStep"), + ("denoise.input.additional_inputs", "QwenImageAdditionalInputsStep"), + ("denoise.prepare_latents", "QwenImagePrepareLatentsStep"), + ("denoise.set_timesteps", "QwenImageSetTimestepsWithStrengthStep"), + ("denoise.prepare_img2img_latents", "QwenImagePrepareLatentsWithStrengthStep"), + ("denoise.prepare_rope_inputs", "QwenImageRoPEInputsStep"), + ("denoise.denoise", "QwenImageDenoiseStep"), + ("denoise.after_denoise", "QwenImageAfterDenoiseStep"), + ("decode.decode", "QwenImageDecoderStep"), + ("decode.postprocess", "QwenImageProcessImagesOutputStep"), + ], + "inpainting": [ + ("text_encoder", "QwenImageTextEncoderStep"), + ("vae_encoder.preprocess", "QwenImageInpaintProcessImagesInputStep"), + ("vae_encoder.encode", "QwenImageVaeEncoderStep"), + ("denoise.input.text_inputs", "QwenImageTextInputsStep"), + ("denoise.input.additional_inputs", "QwenImageAdditionalInputsStep"), + ("denoise.prepare_latents", "QwenImagePrepareLatentsStep"), + ("denoise.set_timesteps", "QwenImageSetTimestepsWithStrengthStep"), + ("denoise.prepare_inpaint_latents.add_noise_to_latents", "QwenImagePrepareLatentsWithStrengthStep"), + ("denoise.prepare_inpaint_latents.create_mask_latents", "QwenImageCreateMaskLatentsStep"), + ("denoise.prepare_rope_inputs", "QwenImageRoPEInputsStep"), + ("denoise.denoise", "QwenImageInpaintDenoiseStep"), + ("denoise.after_denoise", "QwenImageAfterDenoiseStep"), + ("decode.decode", "QwenImageDecoderStep"), + ("decode.postprocess", "QwenImageInpaintProcessImagesOutputStep"), + ], + "controlnet_text2image": [ + ("text_encoder", "QwenImageTextEncoderStep"), + ("controlnet_vae_encoder", "QwenImageControlNetVaeEncoderStep"), + ("denoise.input", "QwenImageTextInputsStep"), + ("denoise.controlnet_input", "QwenImageControlNetInputsStep"), + ("denoise.prepare_latents", "QwenImagePrepareLatentsStep"), + ("denoise.set_timesteps", "QwenImageSetTimestepsStep"), + ("denoise.prepare_rope_inputs", "QwenImageRoPEInputsStep"), + ("denoise.controlnet_before_denoise", "QwenImageControlNetBeforeDenoiserStep"), + ("denoise.controlnet_denoise", "QwenImageControlNetDenoiseStep"), + ("denoise.after_denoise", "QwenImageAfterDenoiseStep"), + ("decode.decode", "QwenImageDecoderStep"), + ("decode.postprocess", "QwenImageProcessImagesOutputStep"), + ], + "controlnet_image2image": [ + ("text_encoder", "QwenImageTextEncoderStep"), + ("vae_encoder.preprocess", "QwenImageProcessImagesInputStep"), + ("vae_encoder.encode", "QwenImageVaeEncoderStep"), + ("controlnet_vae_encoder", "QwenImageControlNetVaeEncoderStep"), + ("denoise.input.text_inputs", "QwenImageTextInputsStep"), + ("denoise.input.additional_inputs", "QwenImageAdditionalInputsStep"), + ("denoise.controlnet_input", "QwenImageControlNetInputsStep"), + ("denoise.prepare_latents", "QwenImagePrepareLatentsStep"), + ("denoise.set_timesteps", "QwenImageSetTimestepsWithStrengthStep"), + ("denoise.prepare_img2img_latents", "QwenImagePrepareLatentsWithStrengthStep"), + ("denoise.prepare_rope_inputs", "QwenImageRoPEInputsStep"), + ("denoise.controlnet_before_denoise", "QwenImageControlNetBeforeDenoiserStep"), + ("denoise.controlnet_denoise", "QwenImageControlNetDenoiseStep"), + ("denoise.after_denoise", "QwenImageAfterDenoiseStep"), + ("decode.decode", "QwenImageDecoderStep"), + ("decode.postprocess", "QwenImageProcessImagesOutputStep"), + ], + "controlnet_inpainting": [ + ("text_encoder", "QwenImageTextEncoderStep"), + ("vae_encoder.preprocess", "QwenImageInpaintProcessImagesInputStep"), + ("vae_encoder.encode", "QwenImageVaeEncoderStep"), + ("controlnet_vae_encoder", "QwenImageControlNetVaeEncoderStep"), + ("denoise.input.text_inputs", "QwenImageTextInputsStep"), + ("denoise.input.additional_inputs", "QwenImageAdditionalInputsStep"), + ("denoise.controlnet_input", "QwenImageControlNetInputsStep"), + ("denoise.prepare_latents", "QwenImagePrepareLatentsStep"), + ("denoise.set_timesteps", "QwenImageSetTimestepsWithStrengthStep"), + ("denoise.prepare_inpaint_latents.add_noise_to_latents", "QwenImagePrepareLatentsWithStrengthStep"), + ("denoise.prepare_inpaint_latents.create_mask_latents", "QwenImageCreateMaskLatentsStep"), + ("denoise.prepare_rope_inputs", "QwenImageRoPEInputsStep"), + ("denoise.controlnet_before_denoise", "QwenImageControlNetBeforeDenoiserStep"), + ("denoise.controlnet_denoise", "QwenImageInpaintControlNetDenoiseStep"), + ("denoise.after_denoise", "QwenImageAfterDenoiseStep"), + ("decode.decode", "QwenImageDecoderStep"), + ("decode.postprocess", "QwenImageInpaintProcessImagesOutputStep"), + ], +} + + class TestQwenImageModularPipelineFast(ModularPipelineTesterMixin, ModularGuiderTesterMixin): pipeline_class = QwenImageModularPipeline pipeline_blocks_class = QwenImageAutoBlocks @@ -37,6 +134,7 @@ class TestQwenImageModularPipelineFast(ModularPipelineTesterMixin, ModularGuider params = frozenset(["prompt", "height", "width", "negative_prompt", "attention_kwargs", "image", "mask_image"]) batch_params = frozenset(["prompt", "negative_prompt", "image", "mask_image"]) + expected_workflow_blocks = QWEN_IMAGE_TEXT2IMAGE_WORKFLOWS def get_dummy_inputs(self): generator = self.get_generator() @@ -56,6 +154,44 @@ def test_inference_batch_single_identical(self): super().test_inference_batch_single_identical(expected_max_diff=5e-4) +QWEN_IMAGE_EDIT_WORKFLOWS = { + "image_conditioned": [ + ("text_encoder.resize", "QwenImageEditResizeStep"), + ("text_encoder.encode", "QwenImageEditTextEncoderStep"), + ("vae_encoder.resize", "QwenImageEditResizeStep"), + ("vae_encoder.preprocess", "QwenImageEditProcessImagesInputStep"), + ("vae_encoder.encode", "QwenImageVaeEncoderStep"), + ("denoise.input.text_inputs", "QwenImageTextInputsStep"), + ("denoise.input.additional_inputs", "QwenImageAdditionalInputsStep"), + ("denoise.prepare_latents", "QwenImagePrepareLatentsStep"), + ("denoise.set_timesteps", "QwenImageSetTimestepsStep"), + ("denoise.prepare_rope_inputs", "QwenImageEditRoPEInputsStep"), + ("denoise.denoise", "QwenImageEditDenoiseStep"), + ("denoise.after_denoise", "QwenImageAfterDenoiseStep"), + ("decode.decode", "QwenImageDecoderStep"), + ("decode.postprocess", "QwenImageProcessImagesOutputStep"), + ], + "image_conditioned_inpainting": [ + ("text_encoder.resize", "QwenImageEditResizeStep"), + ("text_encoder.encode", "QwenImageEditTextEncoderStep"), + ("vae_encoder.resize", "QwenImageEditResizeStep"), + ("vae_encoder.preprocess", "QwenImageEditInpaintProcessImagesInputStep"), + ("vae_encoder.encode", "QwenImageVaeEncoderStep"), + ("denoise.input.text_inputs", "QwenImageTextInputsStep"), + ("denoise.input.additional_inputs", "QwenImageAdditionalInputsStep"), + ("denoise.prepare_latents", "QwenImagePrepareLatentsStep"), + ("denoise.set_timesteps", "QwenImageSetTimestepsWithStrengthStep"), + ("denoise.prepare_inpaint_latents.add_noise_to_latents", "QwenImagePrepareLatentsWithStrengthStep"), + ("denoise.prepare_inpaint_latents.create_mask_latents", "QwenImageCreateMaskLatentsStep"), + ("denoise.prepare_rope_inputs", "QwenImageEditRoPEInputsStep"), + ("denoise.denoise", "QwenImageEditInpaintDenoiseStep"), + ("denoise.after_denoise", "QwenImageAfterDenoiseStep"), + ("decode.decode", "QwenImageDecoderStep"), + ("decode.postprocess", "QwenImageInpaintProcessImagesOutputStep"), + ], +} + + class TestQwenImageEditModularPipelineFast(ModularPipelineTesterMixin, ModularGuiderTesterMixin): pipeline_class = QwenImageEditModularPipeline pipeline_blocks_class = QwenImageEditAutoBlocks @@ -63,6 +199,7 @@ class TestQwenImageEditModularPipelineFast(ModularPipelineTesterMixin, ModularGu params = frozenset(["prompt", "height", "width", "negative_prompt", "attention_kwargs", "image", "mask_image"]) batch_params = frozenset(["prompt", "negative_prompt", "image", "mask_image"]) + expected_workflow_blocks = QWEN_IMAGE_EDIT_WORKFLOWS def get_dummy_inputs(self): generator = self.get_generator() diff --git a/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py b/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py index 7b55933e4caf..f640f0ec83f2 100644 --- a/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py +++ b/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py @@ -267,6 +267,60 @@ def test_controlnet_cfg(self): assert max_diff > 1e-2, "Output with CFG must be different from normal inference" +TEXT2IMAGE_WORKFLOWS = { + "text2image": [ + ("text_encoder", "StableDiffusionXLTextEncoderStep"), + ("denoise.input", "StableDiffusionXLInputStep"), + ("denoise.before_denoise.set_timesteps", "StableDiffusionXLSetTimestepsStep"), + ("denoise.before_denoise.prepare_latents", "StableDiffusionXLPrepareLatentsStep"), + ("denoise.before_denoise.prepare_add_cond", "StableDiffusionXLPrepareAdditionalConditioningStep"), + ("denoise.denoise", "StableDiffusionXLDenoiseStep"), + ("decode", "StableDiffusionXLDecodeStep"), + ], + "controlnet_text2image": [ + ("text_encoder", "StableDiffusionXLTextEncoderStep"), + ("denoise.input", "StableDiffusionXLInputStep"), + ("denoise.before_denoise.set_timesteps", "StableDiffusionXLSetTimestepsStep"), + ("denoise.before_denoise.prepare_latents", "StableDiffusionXLPrepareLatentsStep"), + ("denoise.before_denoise.prepare_add_cond", "StableDiffusionXLPrepareAdditionalConditioningStep"), + ("denoise.controlnet_input", "StableDiffusionXLControlNetInputStep"), + ("denoise.denoise", "StableDiffusionXLControlNetDenoiseStep"), + ("decode", "StableDiffusionXLDecodeStep"), + ], + "controlnet_union_text2image": [ + ("text_encoder", "StableDiffusionXLTextEncoderStep"), + ("denoise.input", "StableDiffusionXLInputStep"), + ("denoise.before_denoise.set_timesteps", "StableDiffusionXLSetTimestepsStep"), + ("denoise.before_denoise.prepare_latents", "StableDiffusionXLPrepareLatentsStep"), + ("denoise.before_denoise.prepare_add_cond", "StableDiffusionXLPrepareAdditionalConditioningStep"), + ("denoise.controlnet_input", "StableDiffusionXLControlNetUnionInputStep"), + ("denoise.denoise", "StableDiffusionXLControlNetDenoiseStep"), + ("decode", "StableDiffusionXLDecodeStep"), + ], + "ip_adapter_text2image": [ + ("text_encoder", "StableDiffusionXLTextEncoderStep"), + ("ip_adapter", "StableDiffusionXLIPAdapterStep"), + ("denoise.input", "StableDiffusionXLInputStep"), + ("denoise.before_denoise.set_timesteps", "StableDiffusionXLSetTimestepsStep"), + ("denoise.before_denoise.prepare_latents", "StableDiffusionXLPrepareLatentsStep"), + ("denoise.before_denoise.prepare_add_cond", "StableDiffusionXLPrepareAdditionalConditioningStep"), + ("denoise.denoise", "StableDiffusionXLDenoiseStep"), + ("decode", "StableDiffusionXLDecodeStep"), + ], + "ip_adapter_controlnet_text2image": [ + ("text_encoder", "StableDiffusionXLTextEncoderStep"), + ("ip_adapter", "StableDiffusionXLIPAdapterStep"), + ("denoise.input", "StableDiffusionXLInputStep"), + ("denoise.before_denoise.set_timesteps", "StableDiffusionXLSetTimestepsStep"), + ("denoise.before_denoise.prepare_latents", "StableDiffusionXLPrepareLatentsStep"), + ("denoise.before_denoise.prepare_add_cond", "StableDiffusionXLPrepareAdditionalConditioningStep"), + ("denoise.controlnet_input", "StableDiffusionXLControlNetInputStep"), + ("denoise.denoise", "StableDiffusionXLControlNetDenoiseStep"), + ("decode", "StableDiffusionXLDecodeStep"), + ], +} + + class TestSDXLModularPipelineFast( SDXLModularTesterMixin, SDXLModularIPAdapterTesterMixin, @@ -291,6 +345,8 @@ class TestSDXLModularPipelineFast( batch_params = frozenset(["prompt", "negative_prompt"]) expected_image_output_shape = (1, 3, 64, 64) + expected_workflow_blocks = TEXT2IMAGE_WORKFLOWS + def get_dummy_inputs(self, seed=0): generator = self.get_generator(seed) inputs = { @@ -314,6 +370,65 @@ def test_inference_batch_single_identical(self): super().test_inference_batch_single_identical(expected_max_diff=3e-3) +IMAGE2IMAGE_WORKFLOWS = { + "image2image": [ + ("text_encoder", "StableDiffusionXLTextEncoderStep"), + ("vae_encoder", "StableDiffusionXLVaeEncoderStep"), + ("denoise.input", "StableDiffusionXLInputStep"), + ("denoise.before_denoise.set_timesteps", "StableDiffusionXLImg2ImgSetTimestepsStep"), + ("denoise.before_denoise.prepare_latents", "StableDiffusionXLImg2ImgPrepareLatentsStep"), + ("denoise.before_denoise.prepare_add_cond", "StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep"), + ("denoise.denoise", "StableDiffusionXLDenoiseStep"), + ("decode", "StableDiffusionXLDecodeStep"), + ], + "controlnet_image2image": [ + ("text_encoder", "StableDiffusionXLTextEncoderStep"), + ("vae_encoder", "StableDiffusionXLVaeEncoderStep"), + ("denoise.input", "StableDiffusionXLInputStep"), + ("denoise.before_denoise.set_timesteps", "StableDiffusionXLImg2ImgSetTimestepsStep"), + ("denoise.before_denoise.prepare_latents", "StableDiffusionXLImg2ImgPrepareLatentsStep"), + ("denoise.before_denoise.prepare_add_cond", "StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep"), + ("denoise.controlnet_input", "StableDiffusionXLControlNetInputStep"), + ("denoise.denoise", "StableDiffusionXLControlNetDenoiseStep"), + ("decode", "StableDiffusionXLDecodeStep"), + ], + "controlnet_union_image2image": [ + ("text_encoder", "StableDiffusionXLTextEncoderStep"), + ("vae_encoder", "StableDiffusionXLVaeEncoderStep"), + ("denoise.input", "StableDiffusionXLInputStep"), + ("denoise.before_denoise.set_timesteps", "StableDiffusionXLImg2ImgSetTimestepsStep"), + ("denoise.before_denoise.prepare_latents", "StableDiffusionXLImg2ImgPrepareLatentsStep"), + ("denoise.before_denoise.prepare_add_cond", "StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep"), + ("denoise.controlnet_input", "StableDiffusionXLControlNetUnionInputStep"), + ("denoise.denoise", "StableDiffusionXLControlNetDenoiseStep"), + ("decode", "StableDiffusionXLDecodeStep"), + ], + "ip_adapter_image2image": [ + ("text_encoder", "StableDiffusionXLTextEncoderStep"), + ("ip_adapter", "StableDiffusionXLIPAdapterStep"), + ("vae_encoder", "StableDiffusionXLVaeEncoderStep"), + ("denoise.input", "StableDiffusionXLInputStep"), + ("denoise.before_denoise.set_timesteps", "StableDiffusionXLImg2ImgSetTimestepsStep"), + ("denoise.before_denoise.prepare_latents", "StableDiffusionXLImg2ImgPrepareLatentsStep"), + ("denoise.before_denoise.prepare_add_cond", "StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep"), + ("denoise.denoise", "StableDiffusionXLDenoiseStep"), + ("decode", "StableDiffusionXLDecodeStep"), + ], + "ip_adapter_controlnet_image2image": [ + ("text_encoder", "StableDiffusionXLTextEncoderStep"), + ("ip_adapter", "StableDiffusionXLIPAdapterStep"), + ("vae_encoder", "StableDiffusionXLVaeEncoderStep"), + ("denoise.input", "StableDiffusionXLInputStep"), + ("denoise.before_denoise.set_timesteps", "StableDiffusionXLImg2ImgSetTimestepsStep"), + ("denoise.before_denoise.prepare_latents", "StableDiffusionXLImg2ImgPrepareLatentsStep"), + ("denoise.before_denoise.prepare_add_cond", "StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep"), + ("denoise.controlnet_input", "StableDiffusionXLControlNetInputStep"), + ("denoise.denoise", "StableDiffusionXLControlNetDenoiseStep"), + ("decode", "StableDiffusionXLDecodeStep"), + ], +} + + class TestSDXLImg2ImgModularPipelineFast( SDXLModularTesterMixin, SDXLModularIPAdapterTesterMixin, @@ -338,6 +453,7 @@ class TestSDXLImg2ImgModularPipelineFast( ) batch_params = frozenset(["prompt", "negative_prompt", "image"]) expected_image_output_shape = (1, 3, 64, 64) + expected_workflow_blocks = IMAGE2IMAGE_WORKFLOWS def get_dummy_inputs(self, seed=0): generator = self.get_generator(seed) @@ -367,6 +483,65 @@ def test_inference_batch_single_identical(self): super().test_inference_batch_single_identical(expected_max_diff=3e-3) +INPAINTING_WORKFLOWS = { + "inpainting": [ + ("text_encoder", "StableDiffusionXLTextEncoderStep"), + ("vae_encoder", "StableDiffusionXLInpaintVaeEncoderStep"), + ("input", "StableDiffusionXLInputStep"), + ("set_timesteps", "StableDiffusionXLImg2ImgSetTimestepsStep"), + ("prepare_latents", "StableDiffusionXLInpaintPrepareLatentsStep"), + ("prepare_add_cond", "StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep"), + ("denoise", "StableDiffusionXLInpaintDenoiseStep"), + ("decode", "StableDiffusionXLInpaintDecodeStep"), + ], + "controlnet_inpainting": [ + ("text_encoder", "StableDiffusionXLTextEncoderStep"), + ("vae_encoder", "StableDiffusionXLInpaintVaeEncoderStep"), + ("input", "StableDiffusionXLInputStep"), + ("set_timesteps", "StableDiffusionXLImg2ImgSetTimestepsStep"), + ("prepare_latents", "StableDiffusionXLInpaintPrepareLatentsStep"), + ("prepare_add_cond", "StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep"), + ("controlnet_input", "StableDiffusionXLControlNetInputStep"), + ("denoise", "StableDiffusionXLInpaintControlNetDenoiseStep"), + ("decode", "StableDiffusionXLInpaintDecodeStep"), + ], + "controlnet_union_inpainting": [ + ("text_encoder", "StableDiffusionXLTextEncoderStep"), + ("vae_encoder", "StableDiffusionXLInpaintVaeEncoderStep"), + ("input", "StableDiffusionXLInputStep"), + ("set_timesteps", "StableDiffusionXLImg2ImgSetTimestepsStep"), + ("prepare_latents", "StableDiffusionXLInpaintPrepareLatentsStep"), + ("prepare_add_cond", "StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep"), + ("controlnet_input", "StableDiffusionXLControlNetUnionInputStep"), + ("denoise", "StableDiffusionXLInpaintControlNetDenoiseStep"), + ("decode", "StableDiffusionXLInpaintDecodeStep"), + ], + "ip_adapter_inpainting": [ + ("text_encoder", "StableDiffusionXLTextEncoderStep"), + ("ip_adapter", "StableDiffusionXLIPAdapterStep"), + ("vae_encoder", "StableDiffusionXLInpaintVaeEncoderStep"), + ("input", "StableDiffusionXLInputStep"), + ("set_timesteps", "StableDiffusionXLImg2ImgSetTimestepsStep"), + ("prepare_latents", "StableDiffusionXLInpaintPrepareLatentsStep"), + ("prepare_add_cond", "StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep"), + ("denoise", "StableDiffusionXLInpaintDenoiseStep"), + ("decode", "StableDiffusionXLInpaintDecodeStep"), + ], + "ip_adapter_controlnet_inpainting": [ + ("text_encoder", "StableDiffusionXLTextEncoderStep"), + ("ip_adapter", "StableDiffusionXLIPAdapterStep"), + ("vae_encoder", "StableDiffusionXLInpaintVaeEncoderStep"), + ("input", "StableDiffusionXLInputStep"), + ("set_timesteps", "StableDiffusionXLImg2ImgSetTimestepsStep"), + ("prepare_latents", "StableDiffusionXLInpaintPrepareLatentsStep"), + ("prepare_add_cond", "StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep"), + ("controlnet_input", "StableDiffusionXLControlNetInputStep"), + ("denoise", "StableDiffusionXLInpaintControlNetDenoiseStep"), + ("decode", "StableDiffusionXLInpaintDecodeStep"), + ], +} + + class SDXLInpaintingModularPipelineFastTests( SDXLModularTesterMixin, SDXLModularIPAdapterTesterMixin, @@ -392,6 +567,7 @@ class SDXLInpaintingModularPipelineFastTests( ) batch_params = frozenset(["prompt", "negative_prompt", "image", "mask_image"]) expected_image_output_shape = (1, 3, 64, 64) + expected_workflow_blocks = INPAINTING_WORKFLOWS def get_dummy_inputs(self, device, seed=0): generator = self.get_generator(seed) diff --git a/tests/modular_pipelines/test_modular_pipelines_common.py b/tests/modular_pipelines/test_modular_pipelines_common.py index 661fcc253795..c1a402a2fd8f 100644 --- a/tests/modular_pipelines/test_modular_pipelines_common.py +++ b/tests/modular_pipelines/test_modular_pipelines_common.py @@ -1,16 +1,36 @@ import gc +import json +import os import tempfile -from typing import Callable, Union +from typing import Callable import pytest import torch import diffusers -from diffusers import ComponentsManager, ModularPipeline, ModularPipelineBlocks +from diffusers import AutoModel, ComponentsManager, ModularPipeline, ModularPipelineBlocks from diffusers.guiders import ClassifierFreeGuidance +from diffusers.modular_pipelines import ( + ConditionalPipelineBlocks, + LoopSequentialPipelineBlocks, + SequentialPipelineBlocks, +) +from diffusers.modular_pipelines.modular_pipeline_utils import ( + ComponentSpec, + ConfigSpec, + InputParam, + OutputParam, + generate_modular_model_card_content, +) from diffusers.utils import logging -from ..testing_utils import backend_empty_cache, numpy_cosine_similarity_distance, require_accelerator, torch_device +from ..testing_utils import ( + CaptureLogger, + backend_empty_cache, + numpy_cosine_similarity_distance, + require_accelerator, + torch_device, +) class ModularPipelineTesterMixin: @@ -30,13 +50,16 @@ class ModularPipelineTesterMixin: optional_params = frozenset(["num_inference_steps", "num_images_per_prompt", "latents", "output_type"]) # this is modular specific: generator needs to be a intermediate input because it's mutable intermediate_params = frozenset(["generator"]) + # Output type for the pipeline (e.g., "images" for image pipelines, "videos" for video pipelines) + # Subclasses can override this to change the expected output type + output_name = "images" def get_generator(self, seed=0): generator = torch.Generator("cpu").manual_seed(seed) return generator @property - def pipeline_class(self) -> Union[Callable, ModularPipeline]: + def pipeline_class(self) -> Callable | ModularPipeline: raise NotImplementedError( "You need to set the attribute `pipeline_class = ClassNameOfPipeline` in the child test class. " "See existing pipeline tests for reference." @@ -49,7 +72,7 @@ def pretrained_model_name_or_path(self) -> str: ) @property - def pipeline_blocks_class(self) -> Union[Callable, ModularPipelineBlocks]: + def pipeline_blocks_class(self) -> Callable | ModularPipelineBlocks: raise NotImplementedError( "You need to set the attribute `pipeline_blocks_class = ClassNameOfPipelineBlocks` in the child test class. " "See existing pipeline tests for reference." @@ -90,6 +113,14 @@ def batch_params(self) -> frozenset: "See existing pipeline tests for reference." ) + @property + def expected_workflow_blocks(self) -> dict: + raise NotImplementedError( + "You need to set the attribute `expected_workflow_blocks` in the child test class. " + "`expected_workflow_blocks` is a dictionary that maps workflow names to list of block names. " + "See existing pipeline tests for reference." + ) + def setup_method(self): # clean up the VRAM before each test torch.compiler.reset() @@ -156,7 +187,7 @@ def test_inference_batch_consistent(self, batch_sizes=[2], batch_generator=True) logger.setLevel(level=diffusers.logging.WARNING) for batch_size, batched_input in zip(batch_sizes, batched_inputs): - output = pipe(**batched_input, output="images") + output = pipe(**batched_input, output=self.output_name) assert len(output) == batch_size, "Output is different from expected batch size" def test_inference_batch_single_identical( @@ -190,12 +221,16 @@ def test_inference_batch_single_identical( if "batch_size" in inputs: batched_inputs["batch_size"] = batch_size - output = pipe(**inputs, output="images") - output_batch = pipe(**batched_inputs, output="images") + output = pipe(**inputs, output=self.output_name) + output_batch = pipe(**batched_inputs, output=self.output_name) assert output_batch.shape[0] == batch_size - max_diff = torch.abs(output_batch[0] - output[0]).max() + # For batch comparison, we only need to compare the first item + if output_batch.shape[0] == batch_size and output.shape[0] == 1: + output_batch = output_batch[0:1] + + max_diff = torch.abs(output_batch - output).max() assert max_diff < expected_max_diff, "Batch inference results different from single inference results" @require_accelerator @@ -210,19 +245,32 @@ def test_float16_inference(self, expected_max_diff=5e-2): # Reset generator in case it is used inside dummy inputs if "generator" in inputs: inputs["generator"] = self.get_generator(0) - output = pipe(**inputs, output="images") + + output = pipe(**inputs, output=self.output_name) fp16_inputs = self.get_dummy_inputs() # Reset generator in case it is used inside dummy inputs if "generator" in fp16_inputs: fp16_inputs["generator"] = self.get_generator(0) - output_fp16 = pipe_fp16(**fp16_inputs, output="images") - output = output.cpu() - output_fp16 = output_fp16.cpu() + output_fp16 = pipe_fp16(**fp16_inputs, output=self.output_name) + + output_tensor = output.float().cpu() + output_fp16_tensor = output_fp16.float().cpu() + + # Check for NaNs in outputs (can happen with tiny models in FP16) + if torch.isnan(output_tensor).any() or torch.isnan(output_fp16_tensor).any(): + pytest.skip("FP16 inference produces NaN values - this is a known issue with tiny models") + + max_diff = numpy_cosine_similarity_distance( + output_tensor.flatten().numpy(), output_fp16_tensor.flatten().numpy() + ) + + # Check if cosine similarity is NaN (which can happen if vectors are zero or very small) + if torch.isnan(torch.tensor(max_diff)): + pytest.skip("Cosine similarity is NaN - outputs may be too small for reliable comparison") - max_diff = numpy_cosine_similarity_distance(output.flatten(), output_fp16.flatten()) - assert max_diff < expected_max_diff, "FP16 inference is different from FP32 inference" + assert max_diff < expected_max_diff, f"FP16 inference is different from FP32 inference (max_diff: {max_diff})" @require_accelerator def test_to_device(self): @@ -244,14 +292,16 @@ def test_to_device(self): def test_inference_is_not_nan_cpu(self): pipe = self.get_pipeline().to("cpu") - output = pipe(**self.get_dummy_inputs(), output="images") + inputs = self.get_dummy_inputs() + output = pipe(**inputs, output=self.output_name) assert torch.isnan(output).sum() == 0, "CPU Inference returns NaN" @require_accelerator def test_inference_is_not_nan(self): pipe = self.get_pipeline().to(torch_device) - output = pipe(**self.get_dummy_inputs(), output="images") + inputs = self.get_dummy_inputs() + output = pipe(**inputs, output=self.output_name) assert torch.isnan(output).sum() == 0, "Accelerator Inference returns NaN" def test_num_images_per_prompt(self): @@ -271,7 +321,7 @@ def test_num_images_per_prompt(self): if key in self.batch_params: inputs[key] = batch_size * [inputs[key]] - images = pipe(**inputs, num_images_per_prompt=num_images_per_prompt, output="images") + images = pipe(**inputs, num_images_per_prompt=num_images_per_prompt, output=self.output_name) assert images.shape[0] == batch_size * num_images_per_prompt @@ -286,8 +336,7 @@ def test_components_auto_cpu_offload_inference_consistent(self): image_slices = [] for pipe in [base_pipe, offload_pipe]: inputs = self.get_dummy_inputs() - image = pipe(**inputs, output="images") - + image = pipe(**inputs, output=self.output_name) image_slices.append(image[0, -3:, -3:, -1].flatten()) assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3 @@ -308,12 +357,66 @@ def test_save_from_pretrained(self): image_slices = [] for pipe in pipes: inputs = self.get_dummy_inputs() - image = pipe(**inputs, output="images") - + image = pipe(**inputs, output=self.output_name) image_slices.append(image[0, -3:, -3:, -1].flatten()) assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3 + def test_modular_index_consistency(self): + pipe = self.get_pipeline() + components_spec = pipe._component_specs + components = sorted(components_spec.keys()) + + with tempfile.TemporaryDirectory() as tmpdir: + pipe.save_pretrained(tmpdir) + index_file = os.path.join(tmpdir, "modular_model_index.json") + assert os.path.exists(index_file) + + with open(index_file) as f: + index_contents = json.load(f) + + compulsory_keys = {"_blocks_class_name", "_class_name", "_diffusers_version"} + for k in compulsory_keys: + assert k in index_contents + + to_check_attrs = {"pretrained_model_name_or_path", "revision", "subfolder"} + for component in components: + spec = components_spec[component] + for attr in to_check_attrs: + if getattr(spec, "pretrained_model_name_or_path", None) is not None: + for attr in to_check_attrs: + assert component in index_contents, f"{component} should be present in index but isn't." + attr_value_from_index = index_contents[component][2][attr] + assert getattr(spec, attr) == attr_value_from_index + + def test_workflow_map(self): + blocks = self.pipeline_blocks_class() + if blocks._workflow_map is None: + pytest.skip("Skipping test as _workflow_map is not set") + + assert hasattr(self, "expected_workflow_blocks") and self.expected_workflow_blocks, ( + "expected_workflow_blocks must be defined in the test class" + ) + + for workflow_name, expected_blocks in self.expected_workflow_blocks.items(): + workflow_blocks = blocks.get_workflow(workflow_name) + actual_blocks = list(workflow_blocks.sub_blocks.items()) + + # Check that the number of blocks matches + assert len(actual_blocks) == len(expected_blocks), ( + f"Workflow '{workflow_name}' has {len(actual_blocks)} blocks, expected {len(expected_blocks)}" + ) + + # Check that each block name and type matches + for i, ((actual_name, actual_block), (expected_name, expected_class_name)) in enumerate( + zip(actual_blocks, expected_blocks) + ): + assert actual_name == expected_name + assert actual_block.__class__.__name__ == expected_class_name, ( + f"Workflow '{workflow_name}': block '{actual_name}' has type " + f"{actual_block.__class__.__name__}, expected {expected_class_name}" + ) + class ModularGuiderTesterMixin: def test_guider_cfg(self, expected_max_diff=1e-2): @@ -324,14 +427,540 @@ def test_guider_cfg(self, expected_max_diff=1e-2): pipe.update_components(guider=guider) inputs = self.get_dummy_inputs() - out_no_cfg = pipe(**inputs, output="images") + out_no_cfg = pipe(**inputs, output=self.output_name) # forward pass with CFG applied guider = ClassifierFreeGuidance(guidance_scale=7.5) pipe.update_components(guider=guider) inputs = self.get_dummy_inputs() - out_cfg = pipe(**inputs, output="images") + out_cfg = pipe(**inputs, output=self.output_name) assert out_cfg.shape == out_no_cfg.shape max_diff = torch.abs(out_cfg - out_no_cfg).max() assert max_diff > expected_max_diff, "Output with CFG must be different from normal inference" + + +class TestCustomBlockRequirements: + def get_dummy_block_pipe(self): + class DummyBlockOne: + # keep two arbitrary deps so that we can test warnings. + _requirements = {"xyz": ">=0.8.0", "abc": ">=10.0.0"} + + class DummyBlockTwo: + # keep two dependencies that will be available during testing. + _requirements = {"transformers": ">=4.44.0", "diffusers": ">=0.2.0"} + + pipe = SequentialPipelineBlocks.from_blocks_dict( + {"dummy_block_one": DummyBlockOne, "dummy_block_two": DummyBlockTwo} + ) + return pipe + + def get_dummy_conditional_block_pipe(self): + class DummyBlockOne: + _requirements = {"xyz": ">=0.8.0", "abc": ">=10.0.0"} + + class DummyBlockTwo: + _requirements = {"transformers": ">=4.44.0", "diffusers": ">=0.2.0"} + + class DummyConditionalBlocks(ConditionalPipelineBlocks): + block_classes = [DummyBlockOne, DummyBlockTwo] + block_names = ["block_one", "block_two"] + block_trigger_inputs = [] + + def select_block(self, **kwargs): + return "block_one" + + return DummyConditionalBlocks() + + def get_dummy_loop_block_pipe(self): + class DummyBlockOne: + _requirements = {"xyz": ">=0.8.0", "abc": ">=10.0.0"} + + class DummyBlockTwo: + _requirements = {"transformers": ">=4.44.0", "diffusers": ">=0.2.0"} + + return LoopSequentialPipelineBlocks.from_blocks_dict({"block_one": DummyBlockOne, "block_two": DummyBlockTwo}) + + def test_sequential_block_requirements_save_load(self, tmp_path): + pipe = self.get_dummy_block_pipe() + pipe.save_pretrained(tmp_path) + + config_path = tmp_path / "modular_config.json" + + with open(config_path, "r") as f: + config = json.load(f) + + assert "requirements" in config + requirements = config["requirements"] + + expected_requirements = { + "xyz": ">=0.8.0", + "abc": ">=10.0.0", + "transformers": ">=4.44.0", + "diffusers": ">=0.2.0", + } + assert expected_requirements == requirements + + def test_sequential_block_requirements_warnings(self, tmp_path): + pipe = self.get_dummy_block_pipe() + + logger = logging.get_logger("diffusers.modular_pipelines.modular_pipeline_utils") + logger.setLevel(30) + + with CaptureLogger(logger) as cap_logger: + pipe.save_pretrained(tmp_path) + + template = "{req} was specified in the requirements but wasn't found in the current environment" + msg_xyz = template.format(req="xyz") + msg_abc = template.format(req="abc") + assert msg_xyz in str(cap_logger.out) + assert msg_abc in str(cap_logger.out) + + def test_conditional_block_requirements_save_load(self, tmp_path): + pipe = self.get_dummy_conditional_block_pipe() + pipe.save_pretrained(tmp_path) + + config_path = tmp_path / "modular_config.json" + with open(config_path, "r") as f: + config = json.load(f) + + assert "requirements" in config + expected_requirements = { + "xyz": ">=0.8.0", + "abc": ">=10.0.0", + "transformers": ">=4.44.0", + "diffusers": ">=0.2.0", + } + assert expected_requirements == config["requirements"] + + def test_loop_block_requirements_save_load(self, tmp_path): + pipe = self.get_dummy_loop_block_pipe() + pipe.save_pretrained(tmp_path) + + config_path = tmp_path / "modular_config.json" + with open(config_path, "r") as f: + config = json.load(f) + + assert "requirements" in config + expected_requirements = { + "xyz": ">=0.8.0", + "abc": ">=10.0.0", + "transformers": ">=4.44.0", + "diffusers": ">=0.2.0", + } + assert expected_requirements == config["requirements"] + + +class TestModularModelCardContent: + def create_mock_block(self, name="TestBlock", description="Test block description"): + class MockBlock: + def __init__(self, name, description): + self.__class__.__name__ = name + self.description = description + self.sub_blocks = {} + + return MockBlock(name, description) + + def create_mock_blocks( + self, + class_name="TestBlocks", + description="Test pipeline description", + num_blocks=2, + components=None, + configs=None, + inputs=None, + outputs=None, + trigger_inputs=None, + model_name=None, + ): + class MockBlocks: + def __init__(self): + self.__class__.__name__ = class_name + self.description = description + self.sub_blocks = {} + self.expected_components = components or [] + self.expected_configs = configs or [] + self.inputs = inputs or [] + self.outputs = outputs or [] + self.trigger_inputs = trigger_inputs + self.model_name = model_name + + blocks = MockBlocks() + + # Add mock sub-blocks + for i in range(num_blocks): + block_name = f"block_{i}" + blocks.sub_blocks[block_name] = self.create_mock_block(f"Block{i}", f"Description for block {i}") + + return blocks + + def test_basic_model_card_content_structure(self): + """Test that all expected keys are present in the output.""" + blocks = self.create_mock_blocks() + content = generate_modular_model_card_content(blocks) + + expected_keys = [ + "pipeline_name", + "model_description", + "blocks_description", + "components_description", + "configs_section", + "io_specification_section", + "trigger_inputs_section", + "tags", + ] + + for key in expected_keys: + assert key in content, f"Expected key '{key}' not found in model card content" + + assert isinstance(content["tags"], list), "Tags should be a list" + + def test_pipeline_name_generation(self): + """Test that pipeline name is correctly generated from blocks class name.""" + blocks = self.create_mock_blocks(class_name="StableDiffusionBlocks") + content = generate_modular_model_card_content(blocks) + + assert content["pipeline_name"] == "StableDiffusion Pipeline" + + def test_tags_generation_text_to_image(self): + """Test that text-to-image tags are correctly generated.""" + blocks = self.create_mock_blocks(trigger_inputs=None) + content = generate_modular_model_card_content(blocks) + + assert "modular-diffusers" in content["tags"] + assert "diffusers" in content["tags"] + assert "text-to-image" in content["tags"] + + def test_tags_generation_with_trigger_inputs(self): + """Test that tags are correctly generated based on trigger inputs.""" + # Test inpainting + blocks = self.create_mock_blocks(trigger_inputs=["mask", "prompt"]) + content = generate_modular_model_card_content(blocks) + assert "inpainting" in content["tags"] + + # Test image-to-image + blocks = self.create_mock_blocks(trigger_inputs=["image", "prompt"]) + content = generate_modular_model_card_content(blocks) + assert "image-to-image" in content["tags"] + + # Test controlnet + blocks = self.create_mock_blocks(trigger_inputs=["control_image", "prompt"]) + content = generate_modular_model_card_content(blocks) + assert "controlnet" in content["tags"] + + def test_tags_with_model_name(self): + """Test that model name is included in tags when present.""" + blocks = self.create_mock_blocks(model_name="stable-diffusion-xl") + content = generate_modular_model_card_content(blocks) + + assert "stable-diffusion-xl" in content["tags"] + + def test_components_description_formatting(self): + """Test that components are correctly formatted.""" + components = [ + ComponentSpec(name="vae", description="VAE component"), + ComponentSpec(name="text_encoder", description="Text encoder component"), + ] + blocks = self.create_mock_blocks(components=components) + content = generate_modular_model_card_content(blocks) + + assert "vae" in content["components_description"] + assert "text_encoder" in content["components_description"] + # Should be enumerated + assert "1." in content["components_description"] + + def test_components_description_empty(self): + """Test handling of pipelines without components.""" + blocks = self.create_mock_blocks(components=None) + content = generate_modular_model_card_content(blocks) + + assert "No specific components required" in content["components_description"] + + def test_configs_section_with_configs(self): + """Test that configs section is generated when configs are present.""" + configs = [ + ConfigSpec(name="num_train_timesteps", default=1000, description="Number of training timesteps"), + ] + blocks = self.create_mock_blocks(configs=configs) + content = generate_modular_model_card_content(blocks) + + assert "## Configuration Parameters" in content["configs_section"] + + def test_configs_section_empty(self): + """Test that configs section is empty when no configs are present.""" + blocks = self.create_mock_blocks(configs=None) + content = generate_modular_model_card_content(blocks) + + assert content["configs_section"] == "" + + def test_inputs_description_required_and_optional(self): + """Test that required and optional inputs are correctly formatted.""" + inputs = [ + InputParam(name="prompt", type_hint=str, required=True, description="The input prompt"), + InputParam(name="num_steps", type_hint=int, required=False, default=50, description="Number of steps"), + ] + blocks = self.create_mock_blocks(inputs=inputs) + content = generate_modular_model_card_content(blocks) + + io_section = content["io_specification_section"] + assert "**Inputs:**" in io_section + assert "prompt" in io_section + assert "num_steps" in io_section + assert "*optional*" in io_section + assert "defaults to `50`" in io_section + + def test_inputs_description_empty(self): + """Test handling of pipelines without specific inputs.""" + blocks = self.create_mock_blocks(inputs=[]) + content = generate_modular_model_card_content(blocks) + + assert "No specific inputs defined" in content["io_specification_section"] + + def test_outputs_description_formatting(self): + """Test that outputs are correctly formatted.""" + outputs = [ + OutputParam(name="images", type_hint=torch.Tensor, description="Generated images"), + ] + blocks = self.create_mock_blocks(outputs=outputs) + content = generate_modular_model_card_content(blocks) + + io_section = content["io_specification_section"] + assert "images" in io_section + assert "Generated images" in io_section + + def test_outputs_description_empty(self): + """Test handling of pipelines without specific outputs.""" + blocks = self.create_mock_blocks(outputs=[]) + content = generate_modular_model_card_content(blocks) + + assert "Standard pipeline outputs" in content["io_specification_section"] + + def test_trigger_inputs_section_with_triggers(self): + """Test that trigger inputs section is generated when present.""" + blocks = self.create_mock_blocks(trigger_inputs=["mask", "image"]) + content = generate_modular_model_card_content(blocks) + + assert "### Conditional Execution" in content["trigger_inputs_section"] + assert "`mask`" in content["trigger_inputs_section"] + assert "`image`" in content["trigger_inputs_section"] + + def test_trigger_inputs_section_empty(self): + """Test that trigger inputs section is empty when not present.""" + blocks = self.create_mock_blocks(trigger_inputs=None) + content = generate_modular_model_card_content(blocks) + + assert content["trigger_inputs_section"] == "" + + def test_model_description_includes_block_count(self): + """Test that model description includes the number of blocks.""" + blocks = self.create_mock_blocks(num_blocks=5) + content = generate_modular_model_card_content(blocks) + + assert "5-block architecture" in content["model_description"] + + +class TestAutoModelLoadIdTagging: + def test_automodel_tags_load_id(self): + model = AutoModel.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe", subfolder="unet") + + assert hasattr(model, "_diffusers_load_id"), "Model should have _diffusers_load_id attribute" + assert model._diffusers_load_id != "null", "_diffusers_load_id should not be 'null'" + + # Verify load_id contains the expected fields + load_id = model._diffusers_load_id + assert "hf-internal-testing/tiny-stable-diffusion-xl-pipe" in load_id + assert "unet" in load_id + + def test_automodel_update_components(self): + pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") + pipe.load_components(torch_dtype=torch.float32) + + auto_model = AutoModel.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe", subfolder="unet") + + pipe.update_components(unet=auto_model) + + assert pipe.unet is auto_model + + assert "unet" in pipe._component_specs + spec = pipe._component_specs["unet"] + assert spec.pretrained_model_name_or_path == "hf-internal-testing/tiny-stable-diffusion-xl-pipe" + assert spec.subfolder == "unet" + + +class TestLoadComponentsSkipBehavior: + def test_load_components_skips_already_loaded(self): + pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") + pipe.load_components(torch_dtype=torch.float32) + + original_unet = pipe.unet + + pipe.load_components() + + # Verify that the unet is the same object (not reloaded) + assert pipe.unet is original_unet, "load_components should skip already loaded components" + + def test_load_components_selective_loading(self): + pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") + + pipe.load_components(names="unet", torch_dtype=torch.float32) + + # Verify only requested component was loaded. + assert hasattr(pipe, "unet") + assert pipe.unet is not None + assert getattr(pipe, "vae", None) is None + + def test_load_components_selective_loading_incremental(self): + """Loading a subset of components should not affect already-loaded components.""" + pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") + + pipe.load_components(names="unet", torch_dtype=torch.float32) + pipe.load_components(names="text_encoder", torch_dtype=torch.float32) + + assert hasattr(pipe, "unet") + assert pipe.unet is not None + assert hasattr(pipe, "text_encoder") + assert pipe.text_encoder is not None + + def test_load_components_skips_invalid_pretrained_path(self): + pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") + + pipe._component_specs["test_component"] = ComponentSpec( + name="test_component", + type_hint=torch.nn.Module, + pretrained_model_name_or_path=None, + default_creation_method="from_pretrained", + ) + pipe.load_components(torch_dtype=torch.float32) + + # Verify test_component was not loaded + assert not hasattr(pipe, "test_component") or pipe.test_component is None + + +class TestCustomModelSavePretrained: + def test_save_pretrained_updates_index_for_local_model(self, tmp_path): + """When a component without _diffusers_load_id (custom/local model) is saved, + modular_model_index.json should point to the save directory.""" + import json + + pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") + pipe.load_components(torch_dtype=torch.float32) + + pipe.unet._diffusers_load_id = "null" + + save_dir = str(tmp_path / "my-pipeline") + pipe.save_pretrained(save_dir) + + with open(os.path.join(save_dir, "modular_model_index.json")) as f: + index = json.load(f) + + _library, _cls, unet_spec = index["unet"] + assert unet_spec["pretrained_model_name_or_path"] == save_dir + assert unet_spec["subfolder"] == "unet" + + _library, _cls, vae_spec = index["vae"] + assert vae_spec["pretrained_model_name_or_path"] == "hf-internal-testing/tiny-stable-diffusion-xl-pipe" + + def test_save_pretrained_roundtrip_with_local_model(self, tmp_path): + """A pipeline with a custom/local model should be saveable and re-loadable with identical outputs.""" + pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") + pipe.load_components(torch_dtype=torch.float32) + + pipe.unet._diffusers_load_id = "null" + + original_state_dict = pipe.unet.state_dict() + + save_dir = str(tmp_path / "my-pipeline") + pipe.save_pretrained(save_dir) + + loaded_pipe = ModularPipeline.from_pretrained(save_dir) + loaded_pipe.load_components(torch_dtype=torch.float32) + + assert loaded_pipe.unet is not None + assert loaded_pipe.unet.__class__.__name__ == pipe.unet.__class__.__name__ + + loaded_state_dict = loaded_pipe.unet.state_dict() + assert set(original_state_dict.keys()) == set(loaded_state_dict.keys()) + for key in original_state_dict: + assert torch.equal(original_state_dict[key], loaded_state_dict[key]), f"Mismatch in {key}" + + def test_save_pretrained_updates_index_for_model_with_no_load_id(self, tmp_path): + """testing the workflow of update the pipeline with a custom model and save the pipeline, + the modular_model_index.json should point to the save directory.""" + import json + + from diffusers import UNet2DConditionModel + + pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") + pipe.load_components(torch_dtype=torch.float32) + + unet = UNet2DConditionModel.from_pretrained( + "hf-internal-testing/tiny-stable-diffusion-xl-pipe", subfolder="unet" + ) + assert not hasattr(unet, "_diffusers_load_id") + + pipe.update_components(unet=unet) + + save_dir = str(tmp_path / "my-pipeline") + pipe.save_pretrained(save_dir) + + with open(os.path.join(save_dir, "modular_model_index.json")) as f: + index = json.load(f) + + _library, _cls, unet_spec = index["unet"] + assert unet_spec["pretrained_model_name_or_path"] == save_dir + assert unet_spec["subfolder"] == "unet" + + _library, _cls, vae_spec = index["vae"] + assert vae_spec["pretrained_model_name_or_path"] == "hf-internal-testing/tiny-stable-diffusion-xl-pipe" + + def test_save_pretrained_overwrite_modular_index(self, tmp_path): + """With overwrite_modular_index=True, all component references should point to the save directory.""" + import json + + pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") + pipe.load_components(torch_dtype=torch.float32) + + save_dir = str(tmp_path / "my-pipeline") + pipe.save_pretrained(save_dir, overwrite_modular_index=True) + + with open(os.path.join(save_dir, "modular_model_index.json")) as f: + index = json.load(f) + + for component_name in ["unet", "vae", "text_encoder", "text_encoder_2"]: + if component_name not in index: + continue + _library, _cls, spec = index[component_name] + assert spec["pretrained_model_name_or_path"] == save_dir, ( + f"{component_name} should point to save dir but got {spec['pretrained_model_name_or_path']}" + ) + assert spec["subfolder"] == component_name + + loaded_pipe = ModularPipeline.from_pretrained(save_dir) + loaded_pipe.load_components(torch_dtype=torch.float32) + + assert loaded_pipe.unet is not None + assert loaded_pipe.vae is not None + + +class TestModularPipelineInitFallback: + """Test that ModularPipeline.__init__ falls back to default_blocks_name when + _blocks_class_name is a base class (e.g. SequentialPipelineBlocks saved by from_blocks_dict).""" + + def test_init_fallback_when_blocks_class_name_is_base_class(self, tmp_path): + # 1. Load pipeline and get a workflow (returns a base SequentialPipelineBlocks) + pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe") + t2i_blocks = pipe.blocks.get_workflow("text2image") + assert t2i_blocks.__class__.__name__ == "SequentialPipelineBlocks" + + # 2. Use init_pipeline to create a new pipeline from the workflow blocks + t2i_pipe = t2i_blocks.init_pipeline("hf-internal-testing/tiny-stable-diffusion-xl-pipe") + + # 3. Save and reload — the saved config will have _blocks_class_name="SequentialPipelineBlocks" + save_dir = str(tmp_path / "pipeline") + t2i_pipe.save_pretrained(save_dir) + loaded_pipe = ModularPipeline.from_pretrained(save_dir) + + # 4. Verify it fell back to default_blocks_name and has correct blocks + assert loaded_pipe.__class__.__name__ == pipe.__class__.__name__ + assert loaded_pipe._blocks.__class__.__name__ == pipe._blocks.__class__.__name__ + assert len(loaded_pipe._blocks.sub_blocks) == len(pipe._blocks.sub_blocks) diff --git a/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py new file mode 100644 index 000000000000..766ca0c16f86 --- /dev/null +++ b/tests/modular_pipelines/test_modular_pipelines_custom_blocks.py @@ -0,0 +1,422 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os +import tempfile +from collections import deque +from typing import List + +import numpy as np +import torch + +from diffusers import FluxTransformer2DModel +from diffusers.modular_pipelines import ( + ComponentSpec, + InputParam, + ModularPipelineBlocks, + OutputParam, + PipelineState, + WanModularPipeline, +) + +from ..testing_utils import nightly, require_torch, slow + + +class DummyCustomBlockSimple(ModularPipelineBlocks): + def __init__(self, use_dummy_model_component=False): + self.use_dummy_model_component = use_dummy_model_component + super().__init__() + + @property + def expected_components(self): + if self.use_dummy_model_component: + return [ComponentSpec("transformer", FluxTransformer2DModel)] + else: + return [] + + @property + def inputs(self) -> List[InputParam]: + return [InputParam("prompt", type_hint=str, required=True, description="Prompt to use")] + + @property + def intermediate_inputs(self) -> List[InputParam]: + return [] + + @property + def intermediate_outputs(self) -> List[OutputParam]: + return [ + OutputParam( + "output_prompt", + type_hint=str, + description="Modified prompt", + ) + ] + + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + old_prompt = block_state.prompt + block_state.output_prompt = "Modular diffusers + " + old_prompt + self.set_block_state(state, block_state) + + return components, state + + +CODE_STR = """ +from diffusers.modular_pipelines import ( + ComponentSpec, + InputParam, + ModularPipelineBlocks, + OutputParam, + PipelineState, + WanModularPipeline, +) +from typing import List + +class DummyCustomBlockSimple(ModularPipelineBlocks): + def __init__(self, use_dummy_model_component=False): + self.use_dummy_model_component = use_dummy_model_component + super().__init__() + + @property + def expected_components(self): + if self.use_dummy_model_component: + return [ComponentSpec("transformer", FluxTransformer2DModel)] + else: + return [] + + @property + def inputs(self) -> List[InputParam]: + return [InputParam("prompt", type_hint=str, required=True, description="Prompt to use")] + + @property + def intermediate_inputs(self) -> List[InputParam]: + return [] + + @property + def intermediate_outputs(self) -> List[OutputParam]: + return [ + OutputParam( + "output_prompt", + type_hint=str, + description="Modified prompt", + ) + ] + + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + old_prompt = block_state.prompt + block_state.output_prompt = "Modular diffusers + " + old_prompt + self.set_block_state(state, block_state) + + return components, state +""" + + +class TestModularCustomBlocks: + def _test_block_properties(self, block): + assert not block.expected_components + assert not block.intermediate_inputs + + actual_inputs = [inp.name for inp in block.inputs] + actual_intermediate_outputs = [out.name for out in block.intermediate_outputs] + assert actual_inputs == ["prompt"] + assert actual_intermediate_outputs == ["output_prompt"] + + def test_custom_block_properties(self): + custom_block = DummyCustomBlockSimple() + self._test_block_properties(custom_block) + + def test_custom_block_output(self): + custom_block = DummyCustomBlockSimple() + pipe = custom_block.init_pipeline() + prompt = "Diffusers is nice" + output = pipe(prompt=prompt) + + actual_inputs = [inp.name for inp in custom_block.inputs] + actual_intermediate_outputs = [out.name for out in custom_block.intermediate_outputs] + assert sorted(output.values) == sorted(actual_inputs + actual_intermediate_outputs) + + output_prompt = output.values["output_prompt"] + assert output_prompt.startswith("Modular diffusers + ") + + def test_custom_block_saving_loading(self): + custom_block = DummyCustomBlockSimple() + + with tempfile.TemporaryDirectory() as tmpdir: + custom_block.save_pretrained(tmpdir) + assert any("modular_config.json" in k for k in os.listdir(tmpdir)) + + with open(os.path.join(tmpdir, "modular_config.json"), "r") as f: + config = json.load(f) + auto_map = config["auto_map"] + assert auto_map == {"ModularPipelineBlocks": "test_modular_pipelines_custom_blocks.DummyCustomBlockSimple"} + + # For now, the Python script that implements the custom block has to be manually pushed to the Hub. + # This is why, we have to separately save the Python script here. + code_path = os.path.join(tmpdir, "test_modular_pipelines_custom_blocks.py") + with open(code_path, "w") as f: + f.write(CODE_STR) + + loaded_custom_block = ModularPipelineBlocks.from_pretrained(tmpdir, trust_remote_code=True) + + pipe = loaded_custom_block.init_pipeline() + prompt = "Diffusers is nice" + output = pipe(prompt=prompt) + + actual_inputs = [inp.name for inp in loaded_custom_block.inputs] + actual_intermediate_outputs = [out.name for out in loaded_custom_block.intermediate_outputs] + assert sorted(output.values) == sorted(actual_inputs + actual_intermediate_outputs) + + output_prompt = output.values["output_prompt"] + assert output_prompt.startswith("Modular diffusers + ") + + def test_custom_block_supported_components(self): + custom_block = DummyCustomBlockSimple(use_dummy_model_component=True) + pipe = custom_block.init_pipeline("hf-internal-testing/tiny-flux-kontext-pipe") + pipe.load_components() + + assert len(pipe.components) == 1 + assert pipe.component_names[0] == "transformer" + + def test_trust_remote_code_not_propagated_to_external_repo(self): + """When a modular pipeline repo references a component from an external repo that has custom + code (auto_map in config), calling load_components(trust_remote_code=True) should NOT + propagate trust_remote_code to that external component. The external component should fail + to load.""" + + from diffusers import ModularPipeline + + CUSTOM_MODEL_CODE = ( + "import torch\n" + "from diffusers import ModelMixin, ConfigMixin\n" + "from diffusers.configuration_utils import register_to_config\n" + "\n" + "class CustomModel(ModelMixin, ConfigMixin):\n" + " @register_to_config\n" + " def __init__(self, hidden_size=8):\n" + " super().__init__()\n" + " self.linear = torch.nn.Linear(hidden_size, hidden_size)\n" + "\n" + " def forward(self, x):\n" + " return self.linear(x)\n" + ) + + with tempfile.TemporaryDirectory() as external_repo_dir, tempfile.TemporaryDirectory() as pipeline_repo_dir: + # Step 1: Create an external model repo with custom code (requires trust_remote_code) + with open(os.path.join(external_repo_dir, "modeling.py"), "w") as f: + f.write(CUSTOM_MODEL_CODE) + + config = { + "_class_name": "CustomModel", + "_diffusers_version": "0.0.0", + "auto_map": {"AutoModel": "modeling.CustomModel"}, + "hidden_size": 8, + } + with open(os.path.join(external_repo_dir, "config.json"), "w") as f: + json.dump(config, f) + + torch.save({}, os.path.join(external_repo_dir, "diffusion_pytorch_model.bin")) + + # Step 2: Create a custom block that references the external repo. + # Define both the class (for direct use) and its code string (for block.py). + class ExternalRefBlock(ModularPipelineBlocks): + @property + def expected_components(self): + return [ + ComponentSpec( + "custom_model", + AutoModel, + pretrained_model_name_or_path=external_repo_dir, + ) + ] + + @property + def inputs(self) -> List[InputParam]: + return [InputParam("prompt", type_hint=str, required=True)] + + @property + def intermediate_inputs(self) -> List[InputParam]: + return [] + + @property + def intermediate_outputs(self) -> List[OutputParam]: + return [OutputParam("output", type_hint=str)] + + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + block_state.output = "test" + self.set_block_state(state, block_state) + return components, state + + EXTERNAL_REF_BLOCK_CODE_STR = ( + "from typing import List\n" + "from diffusers import AutoModel\n" + "from diffusers.modular_pipelines import (\n" + " ComponentSpec,\n" + " InputParam,\n" + " ModularPipelineBlocks,\n" + " OutputParam,\n" + " PipelineState,\n" + ")\n" + "\n" + "class ExternalRefBlock(ModularPipelineBlocks):\n" + " @property\n" + " def expected_components(self):\n" + " return [\n" + " ComponentSpec(\n" + ' "custom_model",\n' + " AutoModel,\n" + f' pretrained_model_name_or_path="{external_repo_dir}",\n' + " )\n" + " ]\n" + "\n" + " @property\n" + " def inputs(self) -> List[InputParam]:\n" + ' return [InputParam("prompt", type_hint=str, required=True)]\n' + "\n" + " @property\n" + " def intermediate_inputs(self) -> List[InputParam]:\n" + " return []\n" + "\n" + " @property\n" + " def intermediate_outputs(self) -> List[OutputParam]:\n" + ' return [OutputParam("output", type_hint=str)]\n' + "\n" + " def __call__(self, components, state: PipelineState) -> PipelineState:\n" + " block_state = self.get_block_state(state)\n" + ' block_state.output = "test"\n' + " self.set_block_state(state, block_state)\n" + " return components, state\n" + ) + + # Save the block config, write block.py, then load back via from_pretrained + block = ExternalRefBlock() + block.save_pretrained(pipeline_repo_dir) + + # auto_map will reference the module name derived from ExternalRefBlock.__module__, + # which is "test_modular_pipelines_custom_blocks". Write the code file with that name. + code_path = os.path.join(pipeline_repo_dir, "test_modular_pipelines_custom_blocks.py") + with open(code_path, "w") as f: + f.write(EXTERNAL_REF_BLOCK_CODE_STR) + + block = ModularPipelineBlocks.from_pretrained(pipeline_repo_dir, trust_remote_code=True) + pipe = block.init_pipeline() + pipe.save_pretrained(pipeline_repo_dir) + + # Step 3: Load the pipeline from the saved directory. + loaded_pipe = ModularPipeline.from_pretrained(pipeline_repo_dir, trust_remote_code=True) + + assert loaded_pipe._pretrained_model_name_or_path == pipeline_repo_dir + assert loaded_pipe._component_specs["custom_model"].pretrained_model_name_or_path == external_repo_dir + assert getattr(loaded_pipe, "custom_model", None) is None + + # Step 4a: load_components WITHOUT trust_remote_code. + # It should still fail + loaded_pipe.load_components() + assert getattr(loaded_pipe, "custom_model", None) is None + + # Step 4b: load_components with trust_remote_code=True. + # trust_remote_code should be stripped for the external component, so it fails. + # The warning should contain guidance about manually loading with trust_remote_code. + loaded_pipe.load_components(trust_remote_code=True) + assert getattr(loaded_pipe, "custom_model", None) is None + + # Step 4c: Manually load with AutoModel and update_components — this should work. + from diffusers import AutoModel + + custom_model = AutoModel.from_pretrained(external_repo_dir, trust_remote_code=True) + loaded_pipe.update_components(custom_model=custom_model) + assert getattr(loaded_pipe, "custom_model", None) is not None + + def test_custom_block_loads_from_hub(self): + repo_id = "hf-internal-testing/tiny-modular-diffusers-block" + block = ModularPipelineBlocks.from_pretrained(repo_id, trust_remote_code=True) + self._test_block_properties(block) + + pipe = block.init_pipeline() + + prompt = "Diffusers is nice" + output = pipe(prompt=prompt) + output_prompt = output.values["output_prompt"] + assert output_prompt.startswith("Modular diffusers + ") + + +@slow +@nightly +@require_torch +class TestKreaCustomBlocksIntegration: + repo_id = "krea/krea-realtime-video" + + def test_loading_from_hub(self): + blocks = ModularPipelineBlocks.from_pretrained(self.repo_id, trust_remote_code=True) + block_names = sorted(blocks.sub_blocks) + + assert block_names == sorted(["text_encoder", "before_denoise", "denoise", "decode"]) + + pipe = WanModularPipeline(blocks, self.repo_id) + pipe.load_components( + trust_remote_code=True, + device_map="cuda", + torch_dtype={"default": torch.bfloat16, "vae": torch.float16}, + ) + assert len(pipe.components) == 7 + assert sorted(pipe.components) == sorted( + ["text_encoder", "tokenizer", "guider", "scheduler", "vae", "transformer", "video_processor"] + ) + + def test_forward(self): + blocks = ModularPipelineBlocks.from_pretrained(self.repo_id, trust_remote_code=True) + pipe = WanModularPipeline(blocks, self.repo_id) + pipe.load_components( + trust_remote_code=True, + device_map="cuda", + torch_dtype={"default": torch.bfloat16, "vae": torch.float16}, + ) + + num_frames_per_block = 2 + num_blocks = 2 + + state = PipelineState() + state.set("frame_cache_context", deque(maxlen=pipe.config.frame_cache_len)) + + prompt = ["a cat sitting on a boat"] + + for block in pipe.transformer.blocks: + block.self_attn.fuse_projections() + + for block_idx in range(num_blocks): + state = pipe( + state, + prompt=prompt, + num_inference_steps=2, + num_blocks=num_blocks, + num_frames_per_block=num_frames_per_block, + block_idx=block_idx, + generator=torch.manual_seed(42), + ) + current_frames = np.array(state.values["videos"][0]) + current_frames_flat = current_frames.flatten() + actual_slices = np.concatenate([current_frames_flat[:4], current_frames_flat[-4:]]).tolist() + + if block_idx == 0: + assert current_frames.shape == (5, 480, 832, 3) + expected_slices = np.array([211, 229, 238, 208, 195, 180, 188, 193]) + else: + assert current_frames.shape == (8, 480, 832, 3) + expected_slices = np.array([179, 203, 214, 176, 194, 181, 187, 191]) + + assert np.allclose(actual_slices, expected_slices) diff --git a/tests/modular_pipelines/wan/__init__.py b/tests/modular_pipelines/wan/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/modular_pipelines/wan/test_modular_pipeline_wan.py b/tests/modular_pipelines/wan/test_modular_pipeline_wan.py new file mode 100644 index 000000000000..c5ed9613e40f --- /dev/null +++ b/tests/modular_pipelines/wan/test_modular_pipeline_wan.py @@ -0,0 +1,49 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from diffusers.modular_pipelines import WanBlocks, WanModularPipeline + +from ..test_modular_pipelines_common import ModularPipelineTesterMixin + + +class TestWanModularPipelineFast(ModularPipelineTesterMixin): + pipeline_class = WanModularPipeline + pipeline_blocks_class = WanBlocks + pretrained_model_name_or_path = "hf-internal-testing/tiny-wan-modular-pipe" + + params = frozenset(["prompt", "height", "width", "num_frames"]) + batch_params = frozenset(["prompt"]) + optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt", "latents"]) + output_name = "videos" + + def get_dummy_inputs(self, seed=0): + generator = self.get_generator(seed) + inputs = { + "prompt": "A painting of a squirrel eating a burger", + "generator": generator, + "num_inference_steps": 2, + "height": 16, + "width": 16, + "num_frames": 9, + "max_sequence_length": 16, + "output_type": "pt", + } + return inputs + + @pytest.mark.skip(reason="num_videos_per_prompt") + def test_num_images_per_prompt(self): + pass diff --git a/tests/modular_pipelines/z_image/__init__.py b/tests/modular_pipelines/z_image/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py b/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py new file mode 100644 index 000000000000..ab45def3ef30 --- /dev/null +++ b/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py @@ -0,0 +1,69 @@ +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from diffusers.modular_pipelines import ZImageAutoBlocks, ZImageModularPipeline + +from ..test_modular_pipelines_common import ModularPipelineTesterMixin + + +ZIMAGE_WORKFLOWS = { + "text2image": [ + ("text_encoder", "ZImageTextEncoderStep"), + ("denoise.input", "ZImageTextInputStep"), + ("denoise.prepare_latents", "ZImagePrepareLatentsStep"), + ("denoise.set_timesteps", "ZImageSetTimestepsStep"), + ("denoise.denoise", "ZImageDenoiseStep"), + ("decode", "ZImageVaeDecoderStep"), + ], + "image2image": [ + ("text_encoder", "ZImageTextEncoderStep"), + ("vae_encoder", "ZImageVaeImageEncoderStep"), + ("denoise.input", "ZImageTextInputStep"), + ("denoise.additional_inputs", "ZImageAdditionalInputsStep"), + ("denoise.prepare_latents", "ZImagePrepareLatentsStep"), + ("denoise.set_timesteps", "ZImageSetTimestepsStep"), + ("denoise.set_timesteps_with_strength", "ZImageSetTimestepsWithStrengthStep"), + ("denoise.prepare_latents_with_image", "ZImagePrepareLatentswithImageStep"), + ("denoise.denoise", "ZImageDenoiseStep"), + ("decode", "ZImageVaeDecoderStep"), + ], +} + + +class TestZImageModularPipelineFast(ModularPipelineTesterMixin): + pipeline_class = ZImageModularPipeline + pipeline_blocks_class = ZImageAutoBlocks + pretrained_model_name_or_path = "hf-internal-testing/tiny-zimage-modular-pipe" + + params = frozenset(["prompt", "height", "width"]) + batch_params = frozenset(["prompt"]) + expected_workflow_blocks = ZIMAGE_WORKFLOWS + + def get_dummy_inputs(self, seed=0): + generator = self.get_generator(seed) + inputs = { + "prompt": "A painting of a squirrel eating a burger", + "generator": generator, + "num_inference_steps": 2, + "height": 32, + "width": 32, + "max_sequence_length": 16, + "output_type": "pt", + } + return inputs + + def test_inference_batch_single_identical(self): + super().test_inference_batch_single_identical(expected_max_diff=5e-3) diff --git a/tests/others/test_check_copies.py b/tests/others/test_check_copies.py index 4b6fa28eb9ac..134faf5ec6f0 100644 --- a/tests/others/test_check_copies.py +++ b/tests/others/test_check_copies.py @@ -41,7 +41,7 @@ \""" prev_sample: torch.Tensor - pred_original_sample: Optional[torch.Tensor] = None + pred_original_sample: torch.Tensor | None = None """ diff --git a/tests/others/test_outputs.py b/tests/others/test_outputs.py index c8069e6916ed..90b8bfe9464e 100644 --- a/tests/others/test_outputs.py +++ b/tests/others/test_outputs.py @@ -1,7 +1,6 @@ import pickle as pkl import unittest from dataclasses import dataclass -from typing import List, Union import numpy as np import PIL.Image @@ -13,7 +12,7 @@ @dataclass class CustomOutput(BaseOutput): - images: Union[List[PIL.Image.Image], np.ndarray] + images: list[PIL.Image.Image] | np.ndarray class ConfigTester(unittest.TestCase): diff --git a/tests/pipelines/allegro/test_allegro.py b/tests/pipelines/allegro/test_allegro.py index b2e588de0647..c126a94ce10e 100644 --- a/tests/pipelines/allegro/test_allegro.py +++ b/tests/pipelines/allegro/test_allegro.py @@ -158,6 +158,10 @@ def test_save_load_local(self): def test_save_load_optional_components(self): pass + @unittest.skip("Decoding without tiling is not yet implemented") + def test_pipeline_with_accelerator_device_map(self): + pass + def test_inference(self): device = "cpu" diff --git a/tests/pipelines/audioldm2/test_audioldm2.py b/tests/pipelines/audioldm2/test_audioldm2.py index 5ccba1dabbfe..33f41bce8f16 100644 --- a/tests/pipelines/audioldm2/test_audioldm2.py +++ b/tests/pipelines/audioldm2/test_audioldm2.py @@ -282,6 +282,8 @@ def test_audioldm2_prompt_embeds(self): text_inputs = text_inputs["input_ids"].to(torch_device) clap_prompt_embeds = audioldm_pipe.text_encoder.get_text_features(text_inputs) + if hasattr(clap_prompt_embeds, "pooler_output"): + clap_prompt_embeds = clap_prompt_embeds.pooler_output clap_prompt_embeds = clap_prompt_embeds[:, None, :] text_inputs = audioldm_pipe.tokenizer_2( @@ -341,6 +343,8 @@ def test_audioldm2_negative_prompt_embeds(self): text_inputs = text_inputs["input_ids"].to(torch_device) clap_prompt_embeds = audioldm_pipe.text_encoder.get_text_features(text_inputs) + if hasattr(clap_prompt_embeds, "pooler_output"): + clap_prompt_embeds = clap_prompt_embeds.pooler_output clap_prompt_embeds = clap_prompt_embeds[:, None, :] text_inputs = audioldm_pipe.tokenizer_2( diff --git a/tests/pipelines/bria/test_pipeline_bria.py b/tests/pipelines/bria/test_pipeline_bria.py index 844488e76f2e..dac9c428cfc9 100644 --- a/tests/pipelines/bria/test_pipeline_bria.py +++ b/tests/pipelines/bria/test_pipeline_bria.py @@ -19,7 +19,7 @@ import numpy as np import torch from huggingface_hub import hf_hub_download -from transformers import T5EncoderModel, T5TokenizerFast +from transformers import AutoConfig, T5EncoderModel, T5TokenizerFast from diffusers import ( AutoencoderKL, @@ -89,7 +89,8 @@ def get_dummy_components(self): scheduler = FlowMatchEulerDiscreteScheduler() torch.manual_seed(0) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = T5TokenizerFast.from_pretrained("hf-internal-testing/tiny-random-t5") components = { diff --git a/tests/pipelines/bria_fibo_edit/__init__.py b/tests/pipelines/bria_fibo_edit/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/pipelines/bria_fibo_edit/test_pipeline_bria_fibo_edit.py b/tests/pipelines/bria_fibo_edit/test_pipeline_bria_fibo_edit.py new file mode 100644 index 000000000000..5376c4b5e03f --- /dev/null +++ b/tests/pipelines/bria_fibo_edit/test_pipeline_bria_fibo_edit.py @@ -0,0 +1,192 @@ +# Copyright 2024 Bria AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import numpy as np +import torch +from PIL import Image +from transformers import AutoTokenizer +from transformers.models.smollm3.modeling_smollm3 import SmolLM3Config, SmolLM3ForCausalLM + +from diffusers import ( + AutoencoderKLWan, + BriaFiboEditPipeline, + FlowMatchEulerDiscreteScheduler, +) +from diffusers.models.transformers.transformer_bria_fibo import BriaFiboTransformer2DModel +from tests.pipelines.test_pipelines_common import PipelineTesterMixin + +from ...testing_utils import ( + enable_full_determinism, + torch_device, +) + + +enable_full_determinism() + + +class BriaFiboPipelineFastTests(PipelineTesterMixin, unittest.TestCase): + pipeline_class = BriaFiboEditPipeline + params = frozenset(["prompt", "height", "width", "guidance_scale"]) + batch_params = frozenset(["prompt"]) + test_xformers_attention = False + test_layerwise_casting = False + test_group_offloading = False + supports_dduf = False + + def get_dummy_components(self): + torch.manual_seed(0) + transformer = BriaFiboTransformer2DModel( + patch_size=1, + in_channels=16, + num_layers=1, + num_single_layers=1, + attention_head_dim=8, + num_attention_heads=2, + joint_attention_dim=64, + text_encoder_dim=32, + pooled_projection_dim=None, + axes_dims_rope=[0, 4, 4], + ) + + vae = AutoencoderKLWan( + base_dim=80, + decoder_base_dim=128, + dim_mult=[1, 2, 4, 4], + dropout=0.0, + in_channels=12, + latents_mean=[0.0] * 16, + latents_std=[1.0] * 16, + is_residual=True, + num_res_blocks=2, + out_channels=12, + patch_size=2, + scale_factor_spatial=16, + scale_factor_temporal=4, + temperal_downsample=[False, True, True], + z_dim=16, + ) + scheduler = FlowMatchEulerDiscreteScheduler() + text_encoder = SmolLM3ForCausalLM(SmolLM3Config(hidden_size=32)) + tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") + + components = { + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + "transformer": transformer, + "vae": vae, + } + return components + + def get_dummy_inputs(self, device, seed=0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device="cpu").manual_seed(seed) + inputs = { + "prompt": '{"text": "A painting of a squirrel eating a burger","edit_instruction": "A painting of a squirrel eating a burger"}', + "negative_prompt": "bad, ugly", + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 5.0, + "height": 192, + "width": 336, + "output_type": "np", + } + image = Image.new("RGB", (336, 192), (255, 255, 255)) + inputs["image"] = image + return inputs + + @unittest.skip(reason="will not be supported due to dim-fusion") + def test_encode_prompt_works_in_isolation(self): + pass + + @unittest.skip(reason="Batching is not supported yet") + def test_num_images_per_prompt(self): + pass + + @unittest.skip(reason="Batching is not supported yet") + def test_inference_batch_consistent(self): + pass + + @unittest.skip(reason="Batching is not supported yet") + def test_inference_batch_single_identical(self): + pass + + def test_bria_fibo_different_prompts(self): + pipe = self.pipeline_class(**self.get_dummy_components()) + pipe = pipe.to(torch_device) + inputs = self.get_dummy_inputs(torch_device) + output_same_prompt = pipe(**inputs).images[0] + + inputs = self.get_dummy_inputs(torch_device) + inputs["prompt"] = {"edit_instruction": "a different prompt"} + output_different_prompts = pipe(**inputs).images[0] + + max_diff = np.abs(output_same_prompt - output_different_prompts).max() + assert max_diff > 1e-6 + + def test_image_output_shape(self): + pipe = self.pipeline_class(**self.get_dummy_components()) + pipe = pipe.to(torch_device) + inputs = self.get_dummy_inputs(torch_device) + + height_width_pairs = [(32, 32), (64, 64), (32, 64)] + for height, width in height_width_pairs: + expected_height = height + expected_width = width + + inputs.update({"height": height, "width": width}) + image = pipe(**inputs).images[0] + output_height, output_width, _ = image.shape + assert (output_height, output_width) == (expected_height, expected_width) + + def test_bria_fibo_edit_mask(self): + pipe = self.pipeline_class(**self.get_dummy_components()) + pipe = pipe.to(torch_device) + inputs = self.get_dummy_inputs(torch_device) + + mask = Image.fromarray((np.ones((192, 336)) * 255).astype(np.uint8), mode="L") + + inputs.update({"mask": mask}) + output = pipe(**inputs).images[0] + + assert output.shape == (192, 336, 3) + + def test_bria_fibo_edit_mask_image_size_mismatch(self): + pipe = self.pipeline_class(**self.get_dummy_components()) + pipe = pipe.to(torch_device) + inputs = self.get_dummy_inputs(torch_device) + + mask = Image.fromarray((np.ones((64, 64)) * 255).astype(np.uint8), mode="L") + + inputs.update({"mask": mask}) + with self.assertRaisesRegex(ValueError, "Mask and image must have the same size"): + pipe(**inputs) + + def test_bria_fibo_edit_mask_no_image(self): + pipe = self.pipeline_class(**self.get_dummy_components()) + pipe = pipe.to(torch_device) + inputs = self.get_dummy_inputs(torch_device) + + mask = Image.fromarray((np.ones((32, 32)) * 255).astype(np.uint8), mode="L") + + # Remove image from inputs if it's there (it shouldn't be by default from get_dummy_inputs) + inputs.pop("image", None) + inputs.update({"mask": mask}) + + with self.assertRaisesRegex(ValueError, "If mask is provided, image must also be provided"): + pipe(**inputs) diff --git a/tests/pipelines/chroma/test_pipeline_chroma.py b/tests/pipelines/chroma/test_pipeline_chroma.py index 3edd58b75f82..6b856128dff0 100644 --- a/tests/pipelines/chroma/test_pipeline_chroma.py +++ b/tests/pipelines/chroma/test_pipeline_chroma.py @@ -2,7 +2,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKL, ChromaPipeline, ChromaTransformer2DModel, FlowMatchEulerDiscreteScheduler @@ -41,7 +41,8 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): ) torch.manual_seed(0) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") diff --git a/tests/pipelines/chroma/test_pipeline_chroma_img2img.py b/tests/pipelines/chroma/test_pipeline_chroma_img2img.py index 4ed1393037b9..8d991c42c749 100644 --- a/tests/pipelines/chroma/test_pipeline_chroma_img2img.py +++ b/tests/pipelines/chroma/test_pipeline_chroma_img2img.py @@ -3,7 +3,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKL, ChromaImg2ImgPipeline, ChromaTransformer2DModel, FlowMatchEulerDiscreteScheduler @@ -42,7 +42,8 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): ) torch.manual_seed(0) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") diff --git a/tests/pipelines/chronoedit/test_chronoedit.py b/tests/pipelines/chronoedit/test_chronoedit.py index 43e5b3159b1c..0b72f93eed3c 100644 --- a/tests/pipelines/chronoedit/test_chronoedit.py +++ b/tests/pipelines/chronoedit/test_chronoedit.py @@ -17,6 +17,7 @@ import torch from PIL import Image from transformers import ( + AutoConfig, AutoTokenizer, CLIPImageProcessor, CLIPVisionConfig, @@ -71,7 +72,8 @@ def get_dummy_components(self): torch.manual_seed(0) # TODO: impl FlowDPMSolverMultistepScheduler scheduler = FlowMatchEulerDiscreteScheduler(shift=7.0) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") torch.manual_seed(0) diff --git a/tests/pipelines/cogvideo/test_cogvideox.py b/tests/pipelines/cogvideo/test_cogvideox.py index dca1725d8a74..e1797cfaac9c 100644 --- a/tests/pipelines/cogvideo/test_cogvideox.py +++ b/tests/pipelines/cogvideo/test_cogvideox.py @@ -18,7 +18,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKLCogVideoX, CogVideoXPipeline, CogVideoXTransformer3DModel, DDIMScheduler @@ -117,7 +117,8 @@ def get_dummy_components(self, num_layers: int = 1): torch.manual_seed(0) scheduler = DDIMScheduler() - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") components = { @@ -235,6 +236,9 @@ def test_attention_slicing_forward_pass( return components = self.get_dummy_components() + for key in components: + if "text_encoder" in key and hasattr(components[key], "eval"): + components[key].eval() pipe = self.pipeline_class(**components) for component in pipe.components.values(): if hasattr(component, "set_default_attn_processor"): diff --git a/tests/pipelines/cogvideo/test_cogvideox_fun_control.py b/tests/pipelines/cogvideo/test_cogvideox_fun_control.py index 097e8df7b35f..30eab652c011 100644 --- a/tests/pipelines/cogvideo/test_cogvideox_fun_control.py +++ b/tests/pipelines/cogvideo/test_cogvideox_fun_control.py @@ -18,7 +18,7 @@ import numpy as np import torch from PIL import Image -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKLCogVideoX, CogVideoXFunControlPipeline, CogVideoXTransformer3DModel, DDIMScheduler @@ -104,7 +104,8 @@ def get_dummy_components(self): torch.manual_seed(0) scheduler = DDIMScheduler() - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") components = { @@ -228,6 +229,9 @@ def test_attention_slicing_forward_pass( return components = self.get_dummy_components() + for key in components: + if "text_encoder" in key and hasattr(components[key], "eval"): + components[key].eval() pipe = self.pipeline_class(**components) for component in pipe.components.values(): if hasattr(component, "set_default_attn_processor"): diff --git a/tests/pipelines/cogvideo/test_cogvideox_image2video.py b/tests/pipelines/cogvideo/test_cogvideox_image2video.py index 1dd5e2ae1405..c30cf56a933c 100644 --- a/tests/pipelines/cogvideo/test_cogvideox_image2video.py +++ b/tests/pipelines/cogvideo/test_cogvideox_image2video.py @@ -19,7 +19,7 @@ import numpy as np import torch from PIL import Image -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKLCogVideoX, CogVideoXImageToVideoPipeline, CogVideoXTransformer3DModel, DDIMScheduler from diffusers.utils import load_image @@ -113,7 +113,8 @@ def get_dummy_components(self): torch.manual_seed(0) scheduler = DDIMScheduler() - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") components = { @@ -237,6 +238,9 @@ def test_attention_slicing_forward_pass( return components = self.get_dummy_components() + for key in components: + if "text_encoder" in key and hasattr(components[key], "eval"): + components[key].eval() pipe = self.pipeline_class(**components) for component in pipe.components.values(): if hasattr(component, "set_default_attn_processor"): diff --git a/tests/pipelines/cogvideo/test_cogvideox_video2video.py b/tests/pipelines/cogvideo/test_cogvideox_video2video.py index 3a1da7c4e7f7..60424ad2a04e 100644 --- a/tests/pipelines/cogvideo/test_cogvideox_video2video.py +++ b/tests/pipelines/cogvideo/test_cogvideox_video2video.py @@ -18,7 +18,7 @@ import numpy as np import torch from PIL import Image -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel, CogVideoXVideoToVideoPipeline, DDIMScheduler @@ -99,7 +99,8 @@ def get_dummy_components(self): torch.manual_seed(0) scheduler = DDIMScheduler() - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") components = { diff --git a/tests/pipelines/cogview3/test_cogview3plus.py b/tests/pipelines/cogview3/test_cogview3plus.py index 819d4b952fc7..374cb6a2a295 100644 --- a/tests/pipelines/cogview3/test_cogview3plus.py +++ b/tests/pipelines/cogview3/test_cogview3plus.py @@ -18,7 +18,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKL, CogVideoXDDIMScheduler, CogView3PlusPipeline, CogView3PlusTransformer2DModel @@ -89,7 +89,8 @@ def get_dummy_components(self): torch.manual_seed(0) scheduler = CogVideoXDDIMScheduler() - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") components = { diff --git a/tests/pipelines/cogview4/test_cogview4.py b/tests/pipelines/cogview4/test_cogview4.py index a1f0fc7a715b..5f71b1b296d9 100644 --- a/tests/pipelines/cogview4/test_cogview4.py +++ b/tests/pipelines/cogview4/test_cogview4.py @@ -108,7 +108,7 @@ def get_dummy_inputs(self, device, seed=0): generator = torch.Generator(device=device).manual_seed(seed) inputs = { "prompt": "dance monkey", - "negative_prompt": "", + "negative_prompt": "bad", "generator": generator, "num_inference_steps": 2, "guidance_scale": 6.0, diff --git a/tests/pipelines/consisid/test_consisid.py b/tests/pipelines/consisid/test_consisid.py index 4fd9e536cddc..b427eeea1d8c 100644 --- a/tests/pipelines/consisid/test_consisid.py +++ b/tests/pipelines/consisid/test_consisid.py @@ -19,7 +19,7 @@ import numpy as np import torch from PIL import Image -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKLCogVideoX, ConsisIDPipeline, ConsisIDTransformer3DModel, DDIMScheduler from diffusers.utils import load_image @@ -122,7 +122,8 @@ def get_dummy_components(self): torch.manual_seed(0) scheduler = DDIMScheduler() - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") components = { @@ -248,6 +249,9 @@ def test_attention_slicing_forward_pass( return components = self.get_dummy_components() + for key in components: + if "text_encoder" in key and hasattr(components[key], "eval"): + components[key].eval() pipe = self.pipeline_class(**components) for component in pipe.components.values(): if hasattr(component, "set_default_attn_processor"): diff --git a/tests/pipelines/controlnet_flux/test_controlnet_flux.py b/tests/pipelines/controlnet_flux/test_controlnet_flux.py index 0895d9de3581..8607cd6944d9 100644 --- a/tests/pipelines/controlnet_flux/test_controlnet_flux.py +++ b/tests/pipelines/controlnet_flux/test_controlnet_flux.py @@ -19,7 +19,7 @@ import numpy as np import torch from huggingface_hub import hf_hub_download -from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast +from transformers import AutoConfig, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast from diffusers import ( AutoencoderKL, @@ -97,7 +97,8 @@ def get_dummy_components(self): text_encoder = CLIPTextModel(clip_text_encoder_config) torch.manual_seed(0) - text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_2 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = T5TokenizerFast.from_pretrained("hf-internal-testing/tiny-random-t5") diff --git a/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py b/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py index 3d8378a5786d..a4749188dfd8 100644 --- a/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py +++ b/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py @@ -2,7 +2,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel from diffusers import ( AutoencoderKL, @@ -13,9 +13,7 @@ ) from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import ( - torch_device, -) +from ...testing_utils import torch_device from ..test_pipelines_common import PipelineTesterMixin, check_qkv_fused_layers_exist @@ -70,7 +68,8 @@ def get_dummy_components(self): text_encoder = CLIPTextModel(clip_text_encoder_config) torch.manual_seed(0) - text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_2 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") diff --git a/tests/pipelines/controlnet_flux/test_controlnet_flux_inpaint.py b/tests/pipelines/controlnet_flux/test_controlnet_flux_inpaint.py index 3ba475deb8a8..6eb560d90848 100644 --- a/tests/pipelines/controlnet_flux/test_controlnet_flux_inpaint.py +++ b/tests/pipelines/controlnet_flux/test_controlnet_flux_inpaint.py @@ -3,15 +3,7 @@ import numpy as np import torch - -# torch_device, # {{ edit_1 }} Removed unused import -from transformers import ( - AutoTokenizer, - CLIPTextConfig, - CLIPTextModel, - CLIPTokenizer, - T5EncoderModel, -) +from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel from diffusers import ( AutoencoderKL, @@ -22,11 +14,7 @@ ) from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import ( - enable_full_determinism, - floats_tensor, - torch_device, -) +from ...testing_utils import enable_full_determinism, floats_tensor, torch_device from ..test_pipelines_common import PipelineTesterMixin @@ -85,7 +73,8 @@ def get_dummy_components(self): text_encoder = CLIPTextModel(clip_text_encoder_config) torch.manual_seed(0) - text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_2 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") diff --git a/tests/pipelines/controlnet_hunyuandit/test_controlnet_hunyuandit.py b/tests/pipelines/controlnet_hunyuandit/test_controlnet_hunyuandit.py index bf31f2abcffb..034ef56b0fd3 100644 --- a/tests/pipelines/controlnet_hunyuandit/test_controlnet_hunyuandit.py +++ b/tests/pipelines/controlnet_hunyuandit/test_controlnet_hunyuandit.py @@ -18,7 +18,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, BertModel, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, BertModel, T5EncoderModel from diffusers import ( AutoencoderKL, @@ -96,7 +96,10 @@ def get_dummy_components(self): scheduler = DDPMScheduler() text_encoder = BertModel.from_pretrained("hf-internal-testing/tiny-random-BertModel") tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-BertModel") - text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + + torch.manual_seed(0) + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_2 = T5EncoderModel(config) tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") components = { diff --git a/tests/pipelines/controlnet_sd3/test_controlnet_inpaint_sd3.py b/tests/pipelines/controlnet_sd3/test_controlnet_inpaint_sd3.py index 34c34b7a2ce7..072f9aa405d9 100644 --- a/tests/pipelines/controlnet_sd3/test_controlnet_inpaint_sd3.py +++ b/tests/pipelines/controlnet_sd3/test_controlnet_inpaint_sd3.py @@ -17,7 +17,14 @@ import numpy as np import torch -from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderModel +from transformers import ( + AutoConfig, + AutoTokenizer, + CLIPTextConfig, + CLIPTextModelWithProjection, + CLIPTokenizer, + T5EncoderModel, +) from diffusers import ( AutoencoderKL, @@ -28,10 +35,7 @@ from diffusers.models import SD3ControlNetModel from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import ( - enable_full_determinism, - torch_device, -) +from ...testing_utils import enable_full_determinism, torch_device from ..test_pipelines_common import PipelineTesterMixin @@ -103,7 +107,8 @@ def get_dummy_components(self): text_encoder_2 = CLIPTextModelWithProjection(clip_text_encoder_config) torch.manual_seed(0) - text_encoder_3 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_3 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") diff --git a/tests/pipelines/controlnet_sd3/test_controlnet_sd3.py b/tests/pipelines/controlnet_sd3/test_controlnet_sd3.py index 2b6cf8d1e8be..46b08cf1f00b 100644 --- a/tests/pipelines/controlnet_sd3/test_controlnet_sd3.py +++ b/tests/pipelines/controlnet_sd3/test_controlnet_sd3.py @@ -15,11 +15,17 @@ import gc import unittest -from typing import Optional import numpy as np import torch -from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderModel +from transformers import ( + AutoConfig, + AutoTokenizer, + CLIPTextConfig, + CLIPTextModelWithProjection, + CLIPTokenizer, + T5EncoderModel, +) from diffusers import ( AutoencoderKL, @@ -63,7 +69,7 @@ class StableDiffusion3ControlNetPipelineFastTests(unittest.TestCase, PipelineTes test_group_offloading = True def get_dummy_components( - self, num_controlnet_layers: int = 3, qk_norm: Optional[str] = "rms_norm", use_dual_attention=False + self, num_controlnet_layers: int = 3, qk_norm: str | None = "rms_norm", use_dual_attention=False ): torch.manual_seed(0) transformer = SD3Transformer2DModel( @@ -118,7 +124,8 @@ def get_dummy_components( text_encoder_2 = CLIPTextModelWithProjection(clip_text_encoder_config) torch.manual_seed(0) - text_encoder_3 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_3 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") diff --git a/tests/pipelines/cosmos/cosmos_guardrail.py b/tests/pipelines/cosmos/cosmos_guardrail.py index c9ef597fdb36..e1d667608ce6 100644 --- a/tests/pipelines/cosmos/cosmos_guardrail.py +++ b/tests/pipelines/cosmos/cosmos_guardrail.py @@ -14,7 +14,6 @@ # ===== This file is an implementation of a dummy guardrail for the fast tests ===== -from typing import Union import numpy as np import torch @@ -35,7 +34,7 @@ def check_text_safety(self, prompt: str) -> bool: def check_video_safety(self, frames: np.ndarray) -> np.ndarray: return frames - def to(self, device: Union[str, torch.device] = None, dtype: torch.dtype = None): + def to(self, device: str | torch.device = None, dtype: torch.dtype = None): module = super().to(device=device, dtype=dtype) return module diff --git a/tests/pipelines/cosmos/test_cosmos.py b/tests/pipelines/cosmos/test_cosmos.py index 32eea9c98c2c..3aa92dca1677 100644 --- a/tests/pipelines/cosmos/test_cosmos.py +++ b/tests/pipelines/cosmos/test_cosmos.py @@ -20,7 +20,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKLCosmos, CosmosTextToWorldPipeline, CosmosTransformer3DModel, EDMEulerScheduler @@ -107,7 +107,8 @@ def get_dummy_components(self): rho=7.0, final_sigmas_type="sigma_min", ) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") components = { @@ -232,6 +233,9 @@ def test_attention_slicing_forward_pass( return components = self.get_dummy_components() + for key in components: + if "text_encoder" in key and hasattr(components[key], "eval"): + components[key].eval() pipe = self.pipeline_class(**components) for component in pipe.components.values(): if hasattr(component, "set_default_attn_processor"): diff --git a/tests/pipelines/cosmos/test_cosmos2_5_transfer.py b/tests/pipelines/cosmos/test_cosmos2_5_transfer.py new file mode 100644 index 000000000000..c098a46a8f8a --- /dev/null +++ b/tests/pipelines/cosmos/test_cosmos2_5_transfer.py @@ -0,0 +1,440 @@ +# Copyright 2025 The HuggingFace Team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import json +import os +import tempfile +import unittest + +import numpy as np +import torch +from transformers import Qwen2_5_VLConfig, Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer + +from diffusers import ( + AutoencoderKLWan, + Cosmos2_5_TransferPipeline, + CosmosControlNetModel, + CosmosTransformer3DModel, + UniPCMultistepScheduler, +) + +from ...testing_utils import enable_full_determinism, torch_device +from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS +from ..test_pipelines_common import PipelineTesterMixin, to_np +from .cosmos_guardrail import DummyCosmosSafetyChecker + + +enable_full_determinism() + + +class Cosmos2_5_TransferWrapper(Cosmos2_5_TransferPipeline): + @staticmethod + def from_pretrained(*args, **kwargs): + if "safety_checker" not in kwargs or kwargs["safety_checker"] is None: + safety_checker = DummyCosmosSafetyChecker() + device_map = kwargs.get("device_map", "cpu") + torch_dtype = kwargs.get("torch_dtype") + if device_map is not None or torch_dtype is not None: + safety_checker = safety_checker.to(device_map, dtype=torch_dtype) + kwargs["safety_checker"] = safety_checker + return Cosmos2_5_TransferPipeline.from_pretrained(*args, **kwargs) + + +class Cosmos2_5_TransferPipelineFastTests(PipelineTesterMixin, unittest.TestCase): + pipeline_class = Cosmos2_5_TransferWrapper + params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} + batch_params = TEXT_TO_IMAGE_BATCH_PARAMS.union({"controls"}) + image_params = TEXT_TO_IMAGE_IMAGE_PARAMS + image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS + required_optional_params = frozenset( + [ + "num_inference_steps", + "generator", + "latents", + "return_dict", + "callback_on_step_end", + "callback_on_step_end_tensor_inputs", + ] + ) + supports_dduf = False + test_xformers_attention = False + test_layerwise_casting = True + test_group_offloading = True + + def get_dummy_components(self): + torch.manual_seed(0) + # Transformer with img_context support for Transfer2.5 + transformer = CosmosTransformer3DModel( + in_channels=16 + 1, + out_channels=16, + num_attention_heads=2, + attention_head_dim=16, + num_layers=2, + mlp_ratio=2, + text_embed_dim=32, + adaln_lora_dim=4, + max_size=(4, 32, 32), + patch_size=(1, 2, 2), + rope_scale=(2.0, 1.0, 1.0), + concat_padding_mask=True, + extra_pos_embed_type="learnable", + controlnet_block_every_n=1, + img_context_dim_in=32, + img_context_num_tokens=4, + img_context_dim_out=32, + ) + + torch.manual_seed(0) + controlnet = CosmosControlNetModel( + n_controlnet_blocks=2, + in_channels=16 + 1 + 1, # control latent channels + condition_mask + padding_mask + latent_channels=16 + 1 + 1, # base latent channels (16) + condition_mask (1) + padding_mask (1) = 18 + model_channels=32, + num_attention_heads=2, + attention_head_dim=16, + mlp_ratio=2, + text_embed_dim=32, + adaln_lora_dim=4, + patch_size=(1, 2, 2), + max_size=(4, 32, 32), + rope_scale=(2.0, 1.0, 1.0), + extra_pos_embed_type="learnable", # Match transformer's config + img_context_dim_in=32, + img_context_dim_out=32, + use_crossattn_projection=False, # Test doesn't need this projection + ) + + torch.manual_seed(0) + vae = AutoencoderKLWan( + base_dim=3, + z_dim=16, + dim_mult=[1, 1, 1, 1], + num_res_blocks=1, + temperal_downsample=[False, True, True], + ) + + torch.manual_seed(0) + scheduler = UniPCMultistepScheduler() + + torch.manual_seed(0) + config = Qwen2_5_VLConfig( + text_config={ + "hidden_size": 16, + "intermediate_size": 16, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "rope_scaling": { + "mrope_section": [1, 1, 2], + "rope_type": "default", + "type": "default", + }, + "rope_theta": 1000000.0, + }, + vision_config={ + "depth": 2, + "hidden_size": 16, + "intermediate_size": 16, + "num_heads": 2, + "out_hidden_size": 16, + }, + hidden_size=16, + vocab_size=152064, + vision_end_token_id=151653, + vision_start_token_id=151652, + vision_token_id=151654, + ) + text_encoder = Qwen2_5_VLForConditionalGeneration(config) + tokenizer = Qwen2Tokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") + + components = { + "transformer": transformer, + "controlnet": controlnet, + "vae": vae, + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + "safety_checker": DummyCosmosSafetyChecker(), + } + return components + + def get_dummy_inputs(self, device, seed=0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device=device).manual_seed(seed) + + controls_generator = torch.Generator(device="cpu").manual_seed(seed) + + inputs = { + "prompt": "dance monkey", + "negative_prompt": "bad quality", + "controls": [torch.randn(3, 32, 32, generator=controls_generator) for _ in range(5)], + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 3.0, + "height": 32, + "width": 32, + "num_frames": 3, + "num_frames_per_chunk": 16, + "max_sequence_length": 16, + "output_type": "pt", + } + + return inputs + + def test_components_function(self): + init_components = self.get_dummy_components() + init_components = {k: v for k, v in init_components.items() if not isinstance(v, (str, int, float))} + pipe = self.pipeline_class(**init_components) + self.assertTrue(hasattr(pipe, "components")) + self.assertTrue(set(pipe.components.keys()) == set(init_components.keys())) + + def test_inference(self): + device = "cpu" + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + video = pipe(**inputs).frames + generated_video = video[0] + self.assertEqual(generated_video.shape, (3, 3, 32, 32)) + self.assertTrue(torch.isfinite(generated_video).all()) + + def test_inference_autoregressive_multi_chunk(self): + device = "cpu" + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + inputs["num_frames"] = 5 + inputs["num_frames_per_chunk"] = 3 + inputs["num_ar_conditional_frames"] = 1 + + video = pipe(**inputs).frames + generated_video = video[0] + self.assertEqual(generated_video.shape, (5, 3, 32, 32)) + self.assertTrue(torch.isfinite(generated_video).all()) + + def test_inference_autoregressive_multi_chunk_no_condition_frames(self): + device = "cpu" + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + inputs["num_frames"] = 5 + inputs["num_frames_per_chunk"] = 3 + inputs["num_ar_conditional_frames"] = 0 + + video = pipe(**inputs).frames + generated_video = video[0] + self.assertEqual(generated_video.shape, (5, 3, 32, 32)) + self.assertTrue(torch.isfinite(generated_video).all()) + + def test_num_frames_per_chunk_above_rope_raises(self): + device = "cpu" + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + inputs["num_frames_per_chunk"] = 17 + + with self.assertRaisesRegex(ValueError, "too large for RoPE setting"): + pipe(**inputs) + + def test_inference_with_controls(self): + """Test inference with control inputs (ControlNet).""" + device = "cpu" + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + inputs["controls"] = [torch.randn(3, 32, 32) for _ in range(5)] # list of 5 frames (C, H, W) + inputs["controls_conditioning_scale"] = 1.0 + inputs["num_frames"] = None + + video = pipe(**inputs).frames + generated_video = video[0] + self.assertEqual(generated_video.shape, (5, 3, 32, 32)) + self.assertTrue(torch.isfinite(generated_video).all()) + + def test_callback_inputs(self): + sig = inspect.signature(self.pipeline_class.__call__) + has_callback_tensor_inputs = "callback_on_step_end_tensor_inputs" in sig.parameters + has_callback_step_end = "callback_on_step_end" in sig.parameters + + if not (has_callback_tensor_inputs and has_callback_step_end): + return + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + self.assertTrue( + hasattr(pipe, "_callback_tensor_inputs"), + f" {self.pipeline_class} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs", + ) + + def callback_inputs_subset(pipe, i, t, callback_kwargs): + for tensor_name in callback_kwargs.keys(): + assert tensor_name in pipe._callback_tensor_inputs + return callback_kwargs + + def callback_inputs_all(pipe, i, t, callback_kwargs): + for tensor_name in pipe._callback_tensor_inputs: + assert tensor_name in callback_kwargs + for tensor_name in callback_kwargs.keys(): + assert tensor_name in pipe._callback_tensor_inputs + return callback_kwargs + + inputs = self.get_dummy_inputs(torch_device) + + inputs["callback_on_step_end"] = callback_inputs_subset + inputs["callback_on_step_end_tensor_inputs"] = ["latents"] + _ = pipe(**inputs)[0] + + inputs["callback_on_step_end"] = callback_inputs_all + inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs + _ = pipe(**inputs)[0] + + def callback_inputs_change_tensor(pipe, i, t, callback_kwargs): + is_last = i == (pipe.num_timesteps - 1) + if is_last: + callback_kwargs["latents"] = torch.zeros_like(callback_kwargs["latents"]) + return callback_kwargs + + inputs["callback_on_step_end"] = callback_inputs_change_tensor + inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs + output = pipe(**inputs)[0] + assert output.abs().sum() < 1e10 + + def test_inference_batch_single_identical(self): + self._test_inference_batch_single_identical(batch_size=2, expected_max_diff=1e-2) + + def test_attention_slicing_forward_pass( + self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3 + ): + if not getattr(self, "test_attention_slicing", True): + return + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + for component in pipe.components.values(): + if hasattr(component, "set_default_attn_processor"): + component.set_default_attn_processor() + pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + + generator_device = "cpu" + inputs = self.get_dummy_inputs(generator_device) + output_without_slicing = pipe(**inputs)[0] + + pipe.enable_attention_slicing(slice_size=1) + inputs = self.get_dummy_inputs(generator_device) + output_with_slicing1 = pipe(**inputs)[0] + + pipe.enable_attention_slicing(slice_size=2) + inputs = self.get_dummy_inputs(generator_device) + output_with_slicing2 = pipe(**inputs)[0] + + if test_max_difference: + max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max() + max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max() + self.assertLess( + max(max_diff1, max_diff2), + expected_max_diff, + "Attention slicing should not affect the inference results", + ) + + def test_serialization_with_variants(self): + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + model_components = [ + component_name + for component_name, component in pipe.components.items() + if isinstance(component, torch.nn.Module) + ] + # Remove components that aren't saved as standard diffusers models + if "safety_checker" in model_components: + model_components.remove("safety_checker") + variant = "fp16" + + with tempfile.TemporaryDirectory() as tmpdir: + pipe.save_pretrained(tmpdir, variant=variant, safe_serialization=False) + + with open(f"{tmpdir}/model_index.json", "r") as f: + config = json.load(f) + + for subfolder in os.listdir(tmpdir): + if not os.path.isfile(subfolder) and subfolder in model_components: + folder_path = os.path.join(tmpdir, subfolder) + is_folder = os.path.isdir(folder_path) and subfolder in config + assert is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)) + + def test_torch_dtype_dict(self): + components = self.get_dummy_components() + if not components: + self.skipTest("No dummy components defined.") + + pipe = self.pipeline_class(**components) + + specified_key = next(iter(components.keys())) + + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: + pipe.save_pretrained(tmpdirname, safe_serialization=False) + torch_dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} + loaded_pipe = self.pipeline_class.from_pretrained( + tmpdirname, safety_checker=DummyCosmosSafetyChecker(), torch_dtype=torch_dtype_dict + ) + + for name, component in loaded_pipe.components.items(): + # Skip components that are not loaded from disk or have special handling + if name == "safety_checker": + continue + if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): + expected_dtype = torch_dtype_dict.get(name, torch_dtype_dict.get("default", torch.float32)) + self.assertEqual( + component.dtype, + expected_dtype, + f"Component '{name}' has dtype {component.dtype} but expected {expected_dtype}", + ) + + def test_save_load_optional_components(self, expected_max_difference=1e-4): + self.pipeline_class._optional_components.remove("safety_checker") + super().test_save_load_optional_components(expected_max_difference=expected_max_difference) + self.pipeline_class._optional_components.append("safety_checker") + + @unittest.skip( + "The pipeline should not be runnable without a safety checker. The test creates a pipeline without passing in " + "a safety checker, which makes the pipeline default to the actual Cosmos Guardrail. The Cosmos Guardrail is " + "too large and slow to run on CI." + ) + def test_encode_prompt_works_in_isolation(self): + pass diff --git a/tests/pipelines/cosmos/test_cosmos2_text2image.py b/tests/pipelines/cosmos/test_cosmos2_text2image.py index 8e3c5e4c29f4..71c61eff0054 100644 --- a/tests/pipelines/cosmos/test_cosmos2_text2image.py +++ b/tests/pipelines/cosmos/test_cosmos2_text2image.py @@ -20,7 +20,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import ( AutoencoderKLWan, @@ -95,7 +95,8 @@ def get_dummy_components(self): torch.manual_seed(0) scheduler = FlowMatchEulerDiscreteScheduler(use_karras_sigmas=True) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") components = { diff --git a/tests/pipelines/cosmos/test_cosmos2_video2world.py b/tests/pipelines/cosmos/test_cosmos2_video2world.py index b0ca0e160d98..1b814257a30a 100644 --- a/tests/pipelines/cosmos/test_cosmos2_video2world.py +++ b/tests/pipelines/cosmos/test_cosmos2_video2world.py @@ -21,7 +21,7 @@ import numpy as np import PIL.Image import torch -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import ( AutoencoderKLWan, @@ -96,7 +96,8 @@ def get_dummy_components(self): torch.manual_seed(0) scheduler = FlowMatchEulerDiscreteScheduler(use_karras_sigmas=True) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") components = { diff --git a/tests/pipelines/cosmos/test_cosmos_video2world.py b/tests/pipelines/cosmos/test_cosmos_video2world.py index 2633c2007ac2..925c3b448677 100644 --- a/tests/pipelines/cosmos/test_cosmos_video2world.py +++ b/tests/pipelines/cosmos/test_cosmos_video2world.py @@ -21,7 +21,7 @@ import numpy as np import PIL.Image import torch -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKLCosmos, CosmosTransformer3DModel, CosmosVideoToWorldPipeline, EDMEulerScheduler @@ -108,7 +108,8 @@ def get_dummy_components(self): rho=7.0, final_sigmas_type="sigma_min", ) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") components = { @@ -245,6 +246,9 @@ def test_attention_slicing_forward_pass( return components = self.get_dummy_components() + for key in components: + if "text_encoder" in key and hasattr(components[key], "eval"): + components[key].eval() pipe = self.pipeline_class(**components) for component in pipe.components.values(): if hasattr(component, "set_default_attn_processor"): diff --git a/tests/pipelines/deepfloyd_if/__init__.py b/tests/pipelines/deepfloyd_if/__init__.py index d47374b07e22..d8f72bd2ca5b 100644 --- a/tests/pipelines/deepfloyd_if/__init__.py +++ b/tests/pipelines/deepfloyd_if/__init__.py @@ -2,7 +2,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import DDPMScheduler, UNet2DConditionModel from diffusers.models.attention_processor import AttnAddedKVProcessor @@ -18,7 +18,8 @@ class IFPipelineTesterMixin: def _get_dummy_components(self): torch.manual_seed(0) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) torch.manual_seed(0) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") @@ -75,7 +76,8 @@ def _get_dummy_components(self): def _get_superresolution_dummy_components(self): torch.manual_seed(0) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) torch.manual_seed(0) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") @@ -250,6 +252,9 @@ def _test_save_load_optional_components(self): # This should be handled in the base test and then this method can be removed. def _test_save_load_local(self): components = self.get_dummy_components() + for key in components: + if "text_encoder" in key and hasattr(components[key], "eval"): + components[key].eval() pipe = self.pipeline_class(**components) pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) diff --git a/tests/pipelines/deepfloyd_if/test_if.py b/tests/pipelines/deepfloyd_if/test_if.py index e1870ddcbae9..0fd1391decd0 100644 --- a/tests/pipelines/deepfloyd_if/test_if.py +++ b/tests/pipelines/deepfloyd_if/test_if.py @@ -18,9 +18,7 @@ import torch -from diffusers import ( - IFPipeline, -) +from diffusers import IFPipeline from diffusers.models.attention_processor import AttnAddedKVProcessor from diffusers.utils.import_utils import is_xformers_available diff --git a/tests/pipelines/flux/test_pipeline_flux.py b/tests/pipelines/flux/test_pipeline_flux.py index 74499bfa607a..13336f0cde9b 100644 --- a/tests/pipelines/flux/test_pipeline_flux.py +++ b/tests/pipelines/flux/test_pipeline_flux.py @@ -4,7 +4,7 @@ import numpy as np import torch from huggingface_hub import hf_hub_download -from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel from diffusers import ( AutoencoderKL, @@ -27,6 +27,7 @@ FasterCacheTesterMixin, FirstBlockCacheTesterMixin, FluxIPAdapterTesterMixin, + MagCacheTesterMixin, PipelineTesterMixin, PyramidAttentionBroadcastTesterMixin, TaylorSeerCacheTesterMixin, @@ -41,6 +42,7 @@ class FluxPipelineFastTests( FasterCacheTesterMixin, FirstBlockCacheTesterMixin, TaylorSeerCacheTesterMixin, + MagCacheTesterMixin, unittest.TestCase, ): pipeline_class = FluxPipeline @@ -91,7 +93,8 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): text_encoder = CLIPTextModel(clip_text_encoder_config) torch.manual_seed(0) - text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_2 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") diff --git a/tests/pipelines/flux/test_pipeline_flux_control.py b/tests/pipelines/flux/test_pipeline_flux_control.py index 7e966470a336..44efca9b9f0e 100644 --- a/tests/pipelines/flux/test_pipeline_flux_control.py +++ b/tests/pipelines/flux/test_pipeline_flux_control.py @@ -3,7 +3,7 @@ import numpy as np import torch from PIL import Image -from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxControlPipeline, FluxTransformer2DModel @@ -53,7 +53,8 @@ def get_dummy_components(self): text_encoder = CLIPTextModel(clip_text_encoder_config) torch.manual_seed(0) - text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_2 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") diff --git a/tests/pipelines/flux/test_pipeline_flux_control_img2img.py b/tests/pipelines/flux/test_pipeline_flux_control_img2img.py index e56136f2e91b..0f0bc0934115 100644 --- a/tests/pipelines/flux/test_pipeline_flux_control_img2img.py +++ b/tests/pipelines/flux/test_pipeline_flux_control_img2img.py @@ -3,7 +3,7 @@ import numpy as np import torch from PIL import Image -from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel from diffusers import ( AutoencoderKL, @@ -57,7 +57,8 @@ def get_dummy_components(self): text_encoder = CLIPTextModel(clip_text_encoder_config) torch.manual_seed(0) - text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_2 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") diff --git a/tests/pipelines/flux/test_pipeline_flux_control_inpaint.py b/tests/pipelines/flux/test_pipeline_flux_control_inpaint.py index e42c5fc2aab5..ae2b6b829e54 100644 --- a/tests/pipelines/flux/test_pipeline_flux_control_inpaint.py +++ b/tests/pipelines/flux/test_pipeline_flux_control_inpaint.py @@ -3,7 +3,7 @@ import numpy as np import torch from PIL import Image -from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel from diffusers import ( AutoencoderKL, @@ -58,7 +58,8 @@ def get_dummy_components(self): text_encoder = CLIPTextModel(clip_text_encoder_config) torch.manual_seed(0) - text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_2 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") diff --git a/tests/pipelines/flux/test_pipeline_flux_fill.py b/tests/pipelines/flux/test_pipeline_flux_fill.py index 25a4a3354820..42cd1efad495 100644 --- a/tests/pipelines/flux/test_pipeline_flux_fill.py +++ b/tests/pipelines/flux/test_pipeline_flux_fill.py @@ -3,7 +3,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxFillPipeline, FluxTransformer2DModel @@ -58,7 +58,8 @@ def get_dummy_components(self): text_encoder = CLIPTextModel(clip_text_encoder_config) torch.manual_seed(0) - text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_2 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") diff --git a/tests/pipelines/flux/test_pipeline_flux_img2img.py b/tests/pipelines/flux/test_pipeline_flux_img2img.py index 6f435760aef5..00587905d337 100644 --- a/tests/pipelines/flux/test_pipeline_flux_img2img.py +++ b/tests/pipelines/flux/test_pipeline_flux_img2img.py @@ -3,7 +3,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxImg2ImgPipeline, FluxTransformer2DModel @@ -55,7 +55,8 @@ def get_dummy_components(self): text_encoder = CLIPTextModel(clip_text_encoder_config) torch.manual_seed(0) - text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_2 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") diff --git a/tests/pipelines/flux/test_pipeline_flux_inpaint.py b/tests/pipelines/flux/test_pipeline_flux_inpaint.py index 6324ff236e10..14edb9e441b5 100644 --- a/tests/pipelines/flux/test_pipeline_flux_inpaint.py +++ b/tests/pipelines/flux/test_pipeline_flux_inpaint.py @@ -3,7 +3,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxInpaintPipeline, FluxTransformer2DModel @@ -55,7 +55,8 @@ def get_dummy_components(self): text_encoder = CLIPTextModel(clip_text_encoder_config) torch.manual_seed(0) - text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_2 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") diff --git a/tests/pipelines/flux/test_pipeline_flux_kontext.py b/tests/pipelines/flux/test_pipeline_flux_kontext.py index 5c78964ea54f..1c018f14b522 100644 --- a/tests/pipelines/flux/test_pipeline_flux_kontext.py +++ b/tests/pipelines/flux/test_pipeline_flux_kontext.py @@ -3,7 +3,7 @@ import numpy as np import PIL.Image import torch -from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel from diffusers import ( AutoencoderKL, @@ -79,7 +79,8 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): text_encoder = CLIPTextModel(clip_text_encoder_config) torch.manual_seed(0) - text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_2 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") diff --git a/tests/pipelines/flux/test_pipeline_flux_kontext_inpaint.py b/tests/pipelines/flux/test_pipeline_flux_kontext_inpaint.py index 9a2e32056dcb..b5f8570ebd1a 100644 --- a/tests/pipelines/flux/test_pipeline_flux_kontext_inpaint.py +++ b/tests/pipelines/flux/test_pipeline_flux_kontext_inpaint.py @@ -3,7 +3,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel from diffusers import ( AutoencoderKL, @@ -79,7 +79,8 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): text_encoder = CLIPTextModel(clip_text_encoder_config) torch.manual_seed(0) - text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_2 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") diff --git a/tests/pipelines/flux2/test_pipeline_flux2_klein.py b/tests/pipelines/flux2/test_pipeline_flux2_klein.py new file mode 100644 index 000000000000..8ed9bf3d1e91 --- /dev/null +++ b/tests/pipelines/flux2/test_pipeline_flux2_klein.py @@ -0,0 +1,183 @@ +import unittest + +import numpy as np +import torch +from PIL import Image +from transformers import Qwen2TokenizerFast, Qwen3Config, Qwen3ForCausalLM + +from diffusers import ( + AutoencoderKLFlux2, + FlowMatchEulerDiscreteScheduler, + Flux2KleinPipeline, + Flux2Transformer2DModel, +) + +from ...testing_utils import torch_device +from ..test_pipelines_common import PipelineTesterMixin, check_qkv_fused_layers_exist + + +class Flux2KleinPipelineFastTests(PipelineTesterMixin, unittest.TestCase): + pipeline_class = Flux2KleinPipeline + params = frozenset(["prompt", "height", "width", "guidance_scale", "prompt_embeds"]) + batch_params = frozenset(["prompt"]) + + test_xformers_attention = False + test_layerwise_casting = True + test_group_offloading = True + + supports_dduf = False + + def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): + torch.manual_seed(0) + transformer = Flux2Transformer2DModel( + patch_size=1, + in_channels=4, + num_layers=num_layers, + num_single_layers=num_single_layers, + attention_head_dim=16, + num_attention_heads=2, + joint_attention_dim=16, + timestep_guidance_channels=256, + axes_dims_rope=[4, 4, 4, 4], + guidance_embeds=False, + ) + + # Create minimal Qwen3 config + config = Qwen3Config( + intermediate_size=16, + hidden_size=16, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=2, + vocab_size=151936, + max_position_embeddings=512, + ) + torch.manual_seed(0) + text_encoder = Qwen3ForCausalLM(config) + + # Use a simple tokenizer for testing + tokenizer = Qwen2TokenizerFast.from_pretrained( + "hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration" + ) + + torch.manual_seed(0) + vae = AutoencoderKLFlux2( + sample_size=32, + in_channels=3, + out_channels=3, + down_block_types=("DownEncoderBlock2D",), + up_block_types=("UpDecoderBlock2D",), + block_out_channels=(4,), + layers_per_block=1, + latent_channels=1, + norm_num_groups=1, + use_quant_conv=False, + use_post_quant_conv=False, + ) + + scheduler = FlowMatchEulerDiscreteScheduler() + + return { + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + "transformer": transformer, + "vae": vae, + } + + def get_dummy_inputs(self, device, seed=0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device="cpu").manual_seed(seed) + + inputs = { + "prompt": "a dog is dancing", + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 4.0, + "height": 8, + "width": 8, + "max_sequence_length": 64, + "output_type": "np", + "text_encoder_out_layers": (1,), + } + return inputs + + def test_fused_qkv_projections(self): + device = "cpu" # ensure determinism for the device-dependent torch.Generator + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + image = pipe(**inputs).images + original_image_slice = image[0, -3:, -3:, -1] + + pipe.transformer.fuse_qkv_projections() + self.assertTrue( + check_qkv_fused_layers_exist(pipe.transformer, ["to_qkv"]), + ("Something wrong with the fused attention layers. Expected all the attention projections to be fused."), + ) + + inputs = self.get_dummy_inputs(device) + image = pipe(**inputs).images + image_slice_fused = image[0, -3:, -3:, -1] + + pipe.transformer.unfuse_qkv_projections() + inputs = self.get_dummy_inputs(device) + image = pipe(**inputs).images + image_slice_disabled = image[0, -3:, -3:, -1] + + self.assertTrue( + np.allclose(original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3), + ("Fusion of QKV projections shouldn't affect the outputs."), + ) + self.assertTrue( + np.allclose(image_slice_fused, image_slice_disabled, atol=1e-3, rtol=1e-3), + ("Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled."), + ) + self.assertTrue( + np.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), + ("Original outputs should match when fused QKV projections are disabled."), + ) + + def test_image_output_shape(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + inputs = self.get_dummy_inputs(torch_device) + + height_width_pairs = [(32, 32), (72, 57)] + for height, width in height_width_pairs: + expected_height = height - height % (pipe.vae_scale_factor * 2) + expected_width = width - width % (pipe.vae_scale_factor * 2) + + inputs.update({"height": height, "width": width}) + image = pipe(**inputs).images[0] + output_height, output_width, _ = image.shape + self.assertEqual( + (output_height, output_width), + (expected_height, expected_width), + f"Output shape {image.shape} does not match expected shape {(expected_height, expected_width)}", + ) + + def test_image_input(self): + device = "cpu" + pipe = self.pipeline_class(**self.get_dummy_components()).to(device) + inputs = self.get_dummy_inputs(device) + + inputs["image"] = Image.new("RGB", (64, 64)) + image = pipe(**inputs).images.flatten() + generated_slice = np.concatenate([image[:8], image[-8:]]) + # fmt: off + expected_slice = np.array( + [ + 0.8255048 , 0.66054785, 0.6643694 , 0.67462724, 0.5494932 , 0.3480271 , 0.52535003, 0.44510138, 0.23549396, 0.21372932, 0.21166152, 0.63198495, 0.49942136, 0.39147034, 0.49156153, 0.3713916 + ] + ) + # fmt: on + assert np.allclose(expected_slice, generated_slice, atol=1e-4, rtol=1e-4) + + @unittest.skip("Needs to be revisited") + def test_encode_prompt_works_in_isolation(self): + pass diff --git a/tests/pipelines/glm_image/__init__.py b/tests/pipelines/glm_image/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/pipelines/glm_image/test_glm_image.py b/tests/pipelines/glm_image/test_glm_image.py new file mode 100644 index 000000000000..f17776554b50 --- /dev/null +++ b/tests/pipelines/glm_image/test_glm_image.py @@ -0,0 +1,397 @@ +# Copyright 2025 The HuggingFace Team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import numpy as np +import torch +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel + +from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, GlmImagePipeline, GlmImageTransformer2DModel +from diffusers.utils import is_transformers_version + +from ...testing_utils import enable_full_determinism, require_torch_accelerator, require_transformers_version_greater +from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS +from ..test_pipelines_common import PipelineTesterMixin + + +if is_transformers_version(">=", "5.0.0.dev0"): + from transformers import GlmImageConfig, GlmImageForConditionalGeneration, GlmImageProcessor + + +enable_full_determinism() + + +@require_transformers_version_greater("4.57.4") +@require_torch_accelerator +class GlmImagePipelineFastTests(PipelineTesterMixin, unittest.TestCase): + pipeline_class = GlmImagePipeline + params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs", "negative_prompt"} + batch_params = TEXT_TO_IMAGE_BATCH_PARAMS + image_params = TEXT_TO_IMAGE_IMAGE_PARAMS + image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS + required_optional_params = frozenset( + [ + "num_inference_steps", + "generator", + "latents", + "return_dict", + "callback_on_step_end", + "callback_on_step_end_tensor_inputs", + ] + ) + test_xformers_attention = False + test_attention_slicing = False + supports_dduf = False + + def get_dummy_components(self): + torch.manual_seed(0) + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) + tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") + + glm_config = GlmImageConfig( + text_config={ + "vocab_size": 168064, + "hidden_size": 32, + "intermediate_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "max_position_embeddings": 512, + "vision_vocab_size": 128, + "rope_parameters": {"mrope_section": (4, 2, 2)}, + }, + vision_config={ + "depth": 2, + "hidden_size": 32, + "num_heads": 2, + "image_size": 32, + "patch_size": 8, + "intermediate_size": 32, + }, + vq_config={"embed_dim": 32, "num_embeddings": 128, "latent_channels": 32}, + ) + + torch.manual_seed(0) + vision_language_encoder = GlmImageForConditionalGeneration(glm_config) + + processor = GlmImageProcessor.from_pretrained("zai-org/GLM-Image", subfolder="processor") + + torch.manual_seed(0) + # For GLM-Image, the relationship between components must satisfy: + # patch_size × vae_scale_factor = 16 (since AR tokens are upsampled 2× from d32) + transformer = GlmImageTransformer2DModel( + patch_size=2, + in_channels=4, + out_channels=4, + num_layers=2, + attention_head_dim=8, + num_attention_heads=2, + text_embed_dim=text_encoder.config.hidden_size, + time_embed_dim=16, + condition_dim=8, + prior_vq_quantizer_codebook_size=128, + ) + + torch.manual_seed(0) + vae = AutoencoderKL( + block_out_channels=(4, 8, 16, 16), + in_channels=3, + out_channels=3, + down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"], + up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"], + latent_channels=4, + norm_num_groups=4, + sample_size=128, + latents_mean=[0.0] * 4, + latents_std=[1.0] * 4, + ) + + scheduler = FlowMatchEulerDiscreteScheduler() + + components = { + "tokenizer": tokenizer, + "processor": processor, + "text_encoder": text_encoder, + "vision_language_encoder": vision_language_encoder, + "vae": vae, + "transformer": transformer, + "scheduler": scheduler, + } + + return components + + def get_dummy_inputs(self, device, seed=0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device=device).manual_seed(seed) + + height, width = 32, 32 + + inputs = { + "prompt": "A photo of a cat", + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 1.5, + "height": height, + "width": width, + "max_sequence_length": 16, + "output_type": "pt", + } + + return inputs + + def test_inference(self): + device = "cpu" + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + image = pipe(**inputs).images[0] + generated_slice = image.flatten() + generated_slice = np.concatenate([generated_slice[:8], generated_slice[-8:]]) + + # fmt: off + expected_slice = np.array( + [ + 0.5849247, 0.50278825, 0.45747858, 0.45895284, 0.43804976, 0.47044256, 0.5239665, 0.47904694, 0.3323419, 0.38725388, 0.28505728, 0.3161863, 0.35026982, 0.37546024, 0.4090118, 0.46629113 + ] + ) + # fmt: on + + self.assertEqual(image.shape, (3, 32, 32)) + self.assertTrue(np.allclose(expected_slice, generated_slice, atol=1e-4, rtol=1e-4)) + + def test_inference_batch_single_identical(self): + """Test that batch=1 produces consistent results with the same seed.""" + device = "cpu" + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + # Run twice with same seed + inputs1 = self.get_dummy_inputs(device, seed=42) + inputs2 = self.get_dummy_inputs(device, seed=42) + + image1 = pipe(**inputs1).images[0] + image2 = pipe(**inputs2).images[0] + + self.assertTrue(torch.allclose(image1, image2, atol=1e-4)) + + def test_inference_batch_multiple_prompts(self): + """Test batch processing with multiple prompts.""" + device = "cpu" + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + generator = torch.Generator(device=device).manual_seed(42) + height, width = 32, 32 + + inputs = { + "prompt": ["A photo of a cat", "A photo of a dog"], + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 1.5, + "height": height, + "width": width, + "max_sequence_length": 16, + "output_type": "pt", + } + + images = pipe(**inputs).images + + # Should return 2 images + self.assertEqual(len(images), 2) + self.assertEqual(images[0].shape, (3, 32, 32)) + self.assertEqual(images[1].shape, (3, 32, 32)) + + def test_num_images_per_prompt(self): + """Test generating multiple images per prompt.""" + device = "cpu" + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + generator = torch.Generator(device=device).manual_seed(42) + height, width = 32, 32 + + inputs = { + "prompt": "A photo of a cat", + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 1.5, + "height": height, + "width": width, + "max_sequence_length": 16, + "output_type": "pt", + "num_images_per_prompt": 2, + } + + images = pipe(**inputs).images + + # Should return 2 images for single prompt + self.assertEqual(len(images), 2) + self.assertEqual(images[0].shape, (3, 32, 32)) + self.assertEqual(images[1].shape, (3, 32, 32)) + + def test_batch_with_num_images_per_prompt(self): + """Test batch prompts with num_images_per_prompt > 1.""" + device = "cpu" + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + generator = torch.Generator(device=device).manual_seed(42) + height, width = 32, 32 + + inputs = { + "prompt": ["A photo of a cat", "A photo of a dog"], + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 1.5, + "height": height, + "width": width, + "max_sequence_length": 16, + "output_type": "pt", + "num_images_per_prompt": 2, + } + + images = pipe(**inputs).images + + # Should return 4 images (2 prompts × 2 images per prompt) + self.assertEqual(len(images), 4) + + def test_prompt_with_prior_token_ids(self): + """Test that prompt and prior_token_ids can be provided together. + + When both are given, the AR generation step is skipped (prior_token_ids is used + directly) and prompt is used to generate prompt_embeds via the glyph encoder. + """ + device = "cpu" + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + height, width = 32, 32 + + # Step 1: Run with prompt only to get prior_token_ids from AR model + generator = torch.Generator(device=device).manual_seed(0) + prior_token_ids, _, _ = pipe.generate_prior_tokens( + prompt="A photo of a cat", + height=height, + width=width, + device=torch.device(device), + generator=torch.Generator(device=device).manual_seed(0), + ) + + # Step 2: Run with both prompt and prior_token_ids — should not raise + generator = torch.Generator(device=device).manual_seed(0) + inputs_both = { + "prompt": "A photo of a cat", + "prior_token_ids": prior_token_ids, + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 1.5, + "height": height, + "width": width, + "max_sequence_length": 16, + "output_type": "pt", + } + images = pipe(**inputs_both).images + self.assertEqual(len(images), 1) + self.assertEqual(images[0].shape, (3, 32, 32)) + + def test_check_inputs_rejects_invalid_combinations(self): + """Test that check_inputs correctly rejects invalid input combinations.""" + device = "cpu" + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + + height, width = 32, 32 + + # Neither prompt nor prior_token_ids → error + with self.assertRaises(ValueError): + pipe.check_inputs( + prompt=None, + height=height, + width=width, + callback_on_step_end_tensor_inputs=None, + prompt_embeds=torch.randn(1, 16, 32), + ) + + # prior_token_ids alone without prompt or prompt_embeds → error + with self.assertRaises(ValueError): + pipe.check_inputs( + prompt=None, + height=height, + width=width, + callback_on_step_end_tensor_inputs=None, + prior_token_ids=torch.randint(0, 100, (1, 64)), + ) + + # prompt + prompt_embeds together → error + with self.assertRaises(ValueError): + pipe.check_inputs( + prompt="A cat", + height=height, + width=width, + callback_on_step_end_tensor_inputs=None, + prompt_embeds=torch.randn(1, 16, 32), + ) + + @unittest.skip("Needs to be revisited.") + def test_encode_prompt_works_in_isolation(self): + pass + + @unittest.skip("Needs to be revisited.") + def test_pipeline_level_group_offloading_inference(self): + pass + + @unittest.skip( + "Follow set of tests are relaxed because this pipeline doesn't guarantee same outputs for the same inputs in consecutive runs." + ) + def test_dict_tuple_outputs_equivalent(self): + pass + + @unittest.skip("Skipped") + def test_cpu_offload_forward_pass_twice(self): + pass + + @unittest.skip("Skipped") + def test_sequential_offload_forward_pass_twice(self): + pass + + @unittest.skip("Skipped") + def test_float16_inference(self): + pass + + @unittest.skip("Skipped") + def test_save_load_float16(self): + pass + + @unittest.skip("Skipped") + def test_save_load_local(self): + pass diff --git a/tests/pipelines/helios/__init__.py b/tests/pipelines/helios/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/pipelines/helios/test_helios.py b/tests/pipelines/helios/test_helios.py new file mode 100644 index 000000000000..b8ee99085036 --- /dev/null +++ b/tests/pipelines/helios/test_helios.py @@ -0,0 +1,172 @@ +# Copyright 2025 The HuggingFace Team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import gc +import unittest + +import torch +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel + +from diffusers import AutoencoderKLWan, HeliosPipeline, HeliosScheduler, HeliosTransformer3DModel + +from ...testing_utils import ( + backend_empty_cache, + enable_full_determinism, + require_torch_accelerator, + slow, + torch_device, +) +from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS +from ..test_pipelines_common import PipelineTesterMixin + + +enable_full_determinism() + + +class HeliosPipelineFastTests(PipelineTesterMixin, unittest.TestCase): + pipeline_class = HeliosPipeline + params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} + batch_params = TEXT_TO_IMAGE_BATCH_PARAMS + image_params = TEXT_TO_IMAGE_IMAGE_PARAMS + image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS + required_optional_params = frozenset( + [ + "num_inference_steps", + "generator", + "latents", + "return_dict", + "callback_on_step_end", + "callback_on_step_end_tensor_inputs", + ] + ) + test_xformers_attention = False + supports_dduf = False + + def get_dummy_components(self): + torch.manual_seed(0) + vae = AutoencoderKLWan( + base_dim=3, + z_dim=16, + dim_mult=[1, 1, 1, 1], + num_res_blocks=1, + temperal_downsample=[False, True, True], + ) + + torch.manual_seed(0) + scheduler = HeliosScheduler(stage_range=[0, 1], stages=1, use_dynamic_shifting=True) + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) + tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") + + torch.manual_seed(0) + transformer = HeliosTransformer3DModel( + patch_size=(1, 2, 2), + num_attention_heads=2, + attention_head_dim=12, + in_channels=16, + out_channels=16, + text_dim=32, + freq_dim=256, + ffn_dim=32, + num_layers=2, + cross_attn_norm=True, + qk_norm="rms_norm_across_heads", + rope_dim=(4, 4, 4), + has_multi_term_memory_patch=True, + guidance_cross_attn=True, + zero_history_timestep=True, + is_amplify_history=False, + ) + + components = { + "transformer": transformer, + "vae": vae, + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + } + return components + + def get_dummy_inputs(self, device, seed=0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device=device).manual_seed(seed) + inputs = { + "prompt": "dance monkey", + "negative_prompt": "negative", + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 1.0, + "height": 16, + "width": 16, + "num_frames": 9, + "max_sequence_length": 16, + "output_type": "pt", + } + return inputs + + def test_inference(self): + device = "cpu" + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + video = pipe(**inputs).frames + generated_video = video[0] + self.assertEqual(generated_video.shape, (33, 3, 16, 16)) + + # fmt: off + expected_slice = torch.tensor([0.4529, 0.4527, 0.4499, 0.4542, 0.4528, 0.4524, 0.4531, 0.4534, 0.5328, + 0.5340, 0.5012, 0.5135, 0.5322, 0.5203, 0.5144, 0.5101]) + # fmt: on + + generated_slice = generated_video.flatten() + generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) + self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=1e-3)) + + # Override to set a more lenient max diff threshold. + def test_save_load_float16(self): + super().test_save_load_float16(expected_max_diff=0.03) + + @unittest.skip("Test not supported") + def test_attention_slicing_forward_pass(self): + pass + + @unittest.skip("Optional components not applicable for Helios") + def test_save_load_optional_components(self): + pass + + +@slow +@require_torch_accelerator +class HeliosPipelineIntegrationTests(unittest.TestCase): + prompt = "A painting of a squirrel eating a burger." + + def setUp(self): + super().setUp() + gc.collect() + backend_empty_cache(torch_device) + + def tearDown(self): + super().tearDown() + gc.collect() + backend_empty_cache(torch_device) + + @unittest.skip("TODO: test needs to be implemented") + def test_helios(self): + pass diff --git a/tests/pipelines/hidream_image/test_pipeline_hidream.py b/tests/pipelines/hidream_image/test_pipeline_hidream.py index ddf39ba4c1e6..1d51872d5526 100644 --- a/tests/pipelines/hidream_image/test_pipeline_hidream.py +++ b/tests/pipelines/hidream_image/test_pipeline_hidream.py @@ -18,6 +18,7 @@ import numpy as np import torch from transformers import ( + AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModelWithProjection, @@ -94,7 +95,8 @@ def get_dummy_components(self): text_encoder_2 = CLIPTextModelWithProjection(clip_text_encoder_config) torch.manual_seed(0) - text_encoder_3 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_3 = T5EncoderModel(config) torch.manual_seed(0) text_encoder_4 = LlamaForCausalLM.from_pretrained("hf-internal-testing/tiny-random-LlamaForCausalLM") @@ -149,12 +151,12 @@ def test_inference(self): self.assertEqual(generated_image.shape, (128, 128, 3)) # fmt: off - expected_slice = np.array([0.4507, 0.5256, 0.4205, 0.5791, 0.4848, 0.4831, 0.4443, 0.5107, 0.6586, 0.3163, 0.7318, 0.5933, 0.6252, 0.5512, 0.5357, 0.5983]) + expected_slice = np.array([0.4501, 0.5256, 0.4207, 0.5783, 0.4842, 0.4833, 0.4441, 0.5112, 0.6587, 0.3169, 0.7308, 0.5927, 0.6251, 0.5509, 0.5355, 0.5969]) # fmt: on generated_slice = generated_image.flatten() generated_slice = np.concatenate([generated_slice[:8], generated_slice[-8:]]) - self.assertTrue(np.allclose(generated_slice, expected_slice, atol=1e-3)) + self.assertTrue(np.allclose(generated_slice, expected_slice, atol=5e-3)) def test_inference_batch_single_identical(self): super().test_inference_batch_single_identical(expected_max_diff=3e-4) diff --git a/tests/pipelines/hunyuan_image_21/test_hunyuanimage.py b/tests/pipelines/hunyuan_image_21/test_hunyuanimage.py index e4b2c686b8b1..7cb6b0913bf6 100644 --- a/tests/pipelines/hunyuan_image_21/test_hunyuanimage.py +++ b/tests/pipelines/hunyuan_image_21/test_hunyuanimage.py @@ -34,11 +34,7 @@ ) from ...testing_utils import enable_full_determinism -from ..test_pipelines_common import ( - FirstBlockCacheTesterMixin, - PipelineTesterMixin, - to_np, -) +from ..test_pipelines_common import FirstBlockCacheTesterMixin, PipelineTesterMixin, to_np enable_full_determinism() @@ -227,7 +223,7 @@ def test_inference_guider(self): self.assertEqual(generated_image.shape, (3, 16, 16)) expected_slice_np = np.array( - [0.61494756, 0.49616697, 0.60327923, 0.6115793, 0.49047345, 0.56977504, 0.53066164, 0.58880305, 0.5570612] + [0.6068114, 0.48716035, 0.5984431, 0.60241306, 0.48849544, 0.5624479, 0.53696984, 0.58964247, 0.54248774] ) output_slice = generated_image[0, -3:, -3:].flatten().cpu().numpy() diff --git a/tests/pipelines/hunyuan_video/test_hunyuan_image2video.py b/tests/pipelines/hunyuan_video/test_hunyuan_image2video.py index 27b5bde31050..2a28e5e42f7d 100644 --- a/tests/pipelines/hunyuan_video/test_hunyuan_image2video.py +++ b/tests/pipelines/hunyuan_video/test_hunyuan_image2video.py @@ -233,7 +233,7 @@ def test_inference(self): self.assertEqual(generated_video.shape, (5, 3, 16, 16)) # fmt: off - expected_slice = torch.tensor([0.444, 0.479, 0.4485, 0.5752, 0.3539, 0.1548, 0.2706, 0.3593, 0.5323, 0.6635, 0.6795, 0.5255, 0.5091, 0.345, 0.4276, 0.4128]) + expected_slice = torch.tensor([0.4441, 0.4790, 0.4485, 0.5748, 0.3539, 0.1553, 0.2707, 0.3594, 0.5331, 0.6645, 0.6799, 0.5257, 0.5092, 0.3450, 0.4276, 0.4127]) # fmt: on generated_slice = generated_video.flatten() @@ -371,13 +371,13 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2): # TODO(aryan): Create a dummy gemma model with smol vocab size @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_consistent(self): pass @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_single_identical(self): pass diff --git a/tests/pipelines/hunyuan_video/test_hunyuan_skyreels_image2video.py b/tests/pipelines/hunyuan_video/test_hunyuan_skyreels_image2video.py index 7ebe797febfa..0ac15d4c405a 100644 --- a/tests/pipelines/hunyuan_video/test_hunyuan_skyreels_image2video.py +++ b/tests/pipelines/hunyuan_video/test_hunyuan_skyreels_image2video.py @@ -333,13 +333,13 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2): # TODO(aryan): Create a dummy gemma model with smol vocab size @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_consistent(self): pass @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_single_identical(self): pass diff --git a/tests/pipelines/hunyuan_video/test_hunyuan_video.py b/tests/pipelines/hunyuan_video/test_hunyuan_video.py index 57a6daebad1f..719bb9e67fbf 100644 --- a/tests/pipelines/hunyuan_video/test_hunyuan_video.py +++ b/tests/pipelines/hunyuan_video/test_hunyuan_video.py @@ -346,13 +346,13 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2): # TODO(aryan): Create a dummy gemma model with smol vocab size @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_consistent(self): pass @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_single_identical(self): pass diff --git a/tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py b/tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py index 51c258b15c38..f7f8908b6f8f 100644 --- a/tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py +++ b/tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py @@ -392,13 +392,13 @@ def test_sequential_offload_forward_pass_twice(self): # TODO(aryan): Create a dummy gemma model with smol vocab size @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_consistent(self): pass @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_single_identical(self): pass diff --git a/tests/pipelines/hunyuan_video1_5/test_hunyuan_1_5.py b/tests/pipelines/hunyuan_video1_5/test_hunyuan_1_5.py index 993c7ef6e4bb..de20148105bf 100644 --- a/tests/pipelines/hunyuan_video1_5/test_hunyuan_1_5.py +++ b/tests/pipelines/hunyuan_video1_5/test_hunyuan_1_5.py @@ -15,7 +15,14 @@ import unittest import torch -from transformers import ByT5Tokenizer, Qwen2_5_VLTextConfig, Qwen2_5_VLTextModel, Qwen2Tokenizer, T5EncoderModel +from transformers import ( + AutoConfig, + ByT5Tokenizer, + Qwen2_5_VLTextConfig, + Qwen2_5_VLTextModel, + Qwen2Tokenizer, + T5EncoderModel, +) from diffusers import ( AutoencoderKLHunyuanVideo15, @@ -114,7 +121,8 @@ def get_dummy_components(self, num_layers: int = 1): tokenizer = Qwen2Tokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") torch.manual_seed(0) - text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_2 = T5EncoderModel(config) tokenizer_2 = ByT5Tokenizer() guider = ClassifierFreeGuidance(guidance_scale=1.0) diff --git a/tests/pipelines/hunyuandit/test_hunyuan_dit.py b/tests/pipelines/hunyuandit/test_hunyuan_dit.py index 2a329f10bc80..ba57b6a3599a 100644 --- a/tests/pipelines/hunyuandit/test_hunyuan_dit.py +++ b/tests/pipelines/hunyuandit/test_hunyuan_dit.py @@ -19,7 +19,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, BertModel, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, BertModel, T5EncoderModel from diffusers import AutoencoderKL, DDPMScheduler, HunyuanDiT2DModel, HunyuanDiTPipeline @@ -74,7 +74,9 @@ def get_dummy_components(self): scheduler = DDPMScheduler() text_encoder = BertModel.from_pretrained("hf-internal-testing/tiny-random-BertModel") tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-BertModel") - text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + torch.manual_seed(0) + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_2 = T5EncoderModel(config) tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") components = { diff --git a/tests/pipelines/kandinsky/test_kandinsky_combined.py b/tests/pipelines/kandinsky/test_kandinsky_combined.py index eba897659700..fdb36b433a94 100644 --- a/tests/pipelines/kandinsky/test_kandinsky_combined.py +++ b/tests/pipelines/kandinsky/test_kandinsky_combined.py @@ -34,9 +34,7 @@ class KandinskyPipelineCombinedFastTests(PipelineTesterMixin, unittest.TestCase): pipeline_class = KandinskyCombinedPipeline - params = [ - "prompt", - ] + params = ["prompt"] batch_params = ["prompt", "negative_prompt"] required_optional_params = [ "generator", @@ -148,6 +146,10 @@ def test_float16_inference(self): def test_dict_tuple_outputs_equivalent(self): super().test_dict_tuple_outputs_equivalent(expected_max_difference=5e-4) + @unittest.skip("Test not supported.") + def test_pipeline_with_accelerator_device_map(self): + pass + class KandinskyPipelineImg2ImgCombinedFastTests(PipelineTesterMixin, unittest.TestCase): pipeline_class = KandinskyImg2ImgCombinedPipeline @@ -264,6 +266,10 @@ def test_dict_tuple_outputs_equivalent(self): def test_save_load_optional_components(self): super().test_save_load_optional_components(expected_max_difference=5e-4) + @unittest.skip("Test not supported.") + def test_pipeline_with_accelerator_device_map(self): + pass + class KandinskyPipelineInpaintCombinedFastTests(PipelineTesterMixin, unittest.TestCase): pipeline_class = KandinskyInpaintCombinedPipeline @@ -384,3 +390,7 @@ def test_save_load_optional_components(self): def test_save_load_local(self): super().test_save_load_local(expected_max_difference=5e-3) + + @unittest.skip("Test not supported.") + def test_pipeline_with_accelerator_device_map(self): + pass diff --git a/tests/pipelines/kandinsky2_2/test_kandinsky_combined.py b/tests/pipelines/kandinsky2_2/test_kandinsky_combined.py index 62f5853da9a5..30aa2896089b 100644 --- a/tests/pipelines/kandinsky2_2/test_kandinsky_combined.py +++ b/tests/pipelines/kandinsky2_2/test_kandinsky_combined.py @@ -36,9 +36,7 @@ class KandinskyV22PipelineCombinedFastTests(PipelineTesterMixin, unittest.TestCase): pipeline_class = KandinskyV22CombinedPipeline - params = [ - "prompt", - ] + params = ["prompt"] batch_params = ["prompt", "negative_prompt"] required_optional_params = [ "generator", @@ -70,12 +68,7 @@ def get_dummy_components(self): def get_dummy_inputs(self, device, seed=0): prior_dummy = PriorDummies() inputs = prior_dummy.get_dummy_inputs(device=device, seed=seed) - inputs.update( - { - "height": 64, - "width": 64, - } - ) + inputs.update({"height": 64, "width": 64}) return inputs def test_kandinsky(self): @@ -155,12 +148,18 @@ def test_save_load_local(self): def test_save_load_optional_components(self): super().test_save_load_optional_components(expected_max_difference=5e-3) + @unittest.skip("Test not supported.") def test_callback_inputs(self): pass + @unittest.skip("Test not supported.") def test_callback_cfg(self): pass + @unittest.skip("Test not supported.") + def test_pipeline_with_accelerator_device_map(self): + pass + class KandinskyV22PipelineImg2ImgCombinedFastTests(PipelineTesterMixin, unittest.TestCase): pipeline_class = KandinskyV22Img2ImgCombinedPipeline @@ -279,12 +278,18 @@ def test_save_load_optional_components(self): def save_load_local(self): super().test_save_load_local(expected_max_difference=5e-3) + @unittest.skip("Test not supported.") def test_callback_inputs(self): pass + @unittest.skip("Test not supported.") def test_callback_cfg(self): pass + @unittest.skip("Test not supported.") + def test_pipeline_with_accelerator_device_map(self): + pass + class KandinskyV22PipelineInpaintCombinedFastTests(PipelineTesterMixin, unittest.TestCase): pipeline_class = KandinskyV22InpaintCombinedPipeline @@ -411,3 +416,7 @@ def test_callback_inputs(self): def test_callback_cfg(self): pass + + @unittest.skip("`device_map` is not yet supported for connected pipelines.") + def test_pipeline_with_accelerator_device_map(self): + pass diff --git a/tests/pipelines/kandinsky2_2/test_kandinsky_inpaint.py b/tests/pipelines/kandinsky2_2/test_kandinsky_inpaint.py index 8a693e9c2dd0..5659699fc7aa 100644 --- a/tests/pipelines/kandinsky2_2/test_kandinsky_inpaint.py +++ b/tests/pipelines/kandinsky2_2/test_kandinsky_inpaint.py @@ -248,6 +248,9 @@ def test_inference_batch_single_identical(self): def test_float16_inference(self): super().test_float16_inference(expected_max_diff=5e-1) + def test_save_load_dduf(self): + super().test_save_load_dduf(atol=1e-3, rtol=1e-3) + @is_flaky() def test_model_cpu_offload_forward_pass(self): super().test_inference_batch_single_identical(expected_max_diff=8e-4) @@ -293,6 +296,9 @@ def callback_inputs_test(pipe, i, t, callback_kwargs): output = pipe(**inputs)[0] assert output.abs().sum() == 0 + def test_pipeline_with_accelerator_device_map(self): + super().test_pipeline_with_accelerator_device_map(expected_max_difference=5e-3) + @slow @require_torch_accelerator diff --git a/tests/pipelines/kandinsky3/test_kandinsky3.py b/tests/pipelines/kandinsky3/test_kandinsky3.py index 55500f729bbb..7deae194fb4b 100644 --- a/tests/pipelines/kandinsky3/test_kandinsky3.py +++ b/tests/pipelines/kandinsky3/test_kandinsky3.py @@ -19,7 +19,7 @@ import numpy as np import torch from PIL import Image -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import ( AutoPipelineForImage2Image, @@ -108,7 +108,8 @@ def get_dummy_components(self, time_cond_proj_dim=None): torch.manual_seed(0) movq = self.dummy_movq torch.manual_seed(0) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config).eval() torch.manual_seed(0) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") @@ -155,9 +156,9 @@ def test_kandinsky3(self): assert image.shape == (1, 16, 16, 3) - expected_slice = np.array([0.3768, 0.4373, 0.4865, 0.4890, 0.4299, 0.5122, 0.4921, 0.4924, 0.5599]) + expected_slice = np.array([0.3944, 0.3680, 0.4842, 0.5333, 0.4412, 0.4812, 0.5089, 0.5381, 0.5578]) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2, ( + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-1, ( f" expected_slice {expected_slice}, but got {image_slice.flatten()}" ) diff --git a/tests/pipelines/kandinsky3/test_kandinsky3_img2img.py b/tests/pipelines/kandinsky3/test_kandinsky3_img2img.py index 503fdb242dff..469300e67788 100644 --- a/tests/pipelines/kandinsky3/test_kandinsky3_img2img.py +++ b/tests/pipelines/kandinsky3/test_kandinsky3_img2img.py @@ -20,7 +20,7 @@ import numpy as np import torch from PIL import Image -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import ( AutoPipelineForImage2Image, @@ -119,7 +119,8 @@ def get_dummy_components(self, time_cond_proj_dim=None): torch.manual_seed(0) movq = self.dummy_movq torch.manual_seed(0) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config).eval() torch.manual_seed(0) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") @@ -155,10 +156,7 @@ def get_dummy_inputs(self, device, seed=0): return inputs def test_dict_tuple_outputs_equivalent(self): - expected_slice = None - if torch_device == "cpu": - expected_slice = np.array([0.5762, 0.6112, 0.4150, 0.6018, 0.6167, 0.4626, 0.5426, 0.5641, 0.6536]) - super().test_dict_tuple_outputs_equivalent(expected_slice=expected_slice) + super().test_dict_tuple_outputs_equivalent() def test_kandinsky3_img2img(self): device = "cpu" @@ -177,11 +175,9 @@ def test_kandinsky3_img2img(self): assert image.shape == (1, 64, 64, 3) - expected_slice = np.array( - [0.576259, 0.6132097, 0.41703486, 0.603196, 0.62062526, 0.4655338, 0.5434324, 0.5660727, 0.65433365] - ) + expected_slice = np.array([0.5725, 0.6248, 0.4355, 0.5732, 0.6105, 0.5267, 0.5470, 0.5512, 0.6618]) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2, ( + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-1, ( f" expected_slice {expected_slice}, but got {image_slice.flatten()}" ) @@ -191,6 +187,12 @@ def test_float16_inference(self): def test_inference_batch_single_identical(self): super().test_inference_batch_single_identical(expected_max_diff=1e-2) + def test_save_load_dduf(self): + super().test_save_load_dduf(atol=1e-3, rtol=1e-3) + + def test_pipeline_with_accelerator_device_map(self): + super().test_pipeline_with_accelerator_device_map(expected_max_difference=5e-3) + @slow @require_torch_accelerator diff --git a/tests/pipelines/latte/test_latte.py b/tests/pipelines/latte/test_latte.py index a40d4bf8eede..873c06e11c5b 100644 --- a/tests/pipelines/latte/test_latte.py +++ b/tests/pipelines/latte/test_latte.py @@ -20,7 +20,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import ( AutoencoderKL, @@ -109,7 +109,8 @@ def get_dummy_components(self, num_layers: int = 1): vae = AutoencoderKL() scheduler = DDIMScheduler() - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") diff --git a/tests/pipelines/ltx/test_ltx.py b/tests/pipelines/ltx/test_ltx.py index aaf4161b51fb..9836551d30a1 100644 --- a/tests/pipelines/ltx/test_ltx.py +++ b/tests/pipelines/ltx/test_ltx.py @@ -17,7 +17,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKLLTXVideo, FlowMatchEulerDiscreteScheduler, LTXPipeline, LTXVideoTransformer3DModel @@ -88,7 +88,8 @@ def get_dummy_components(self, num_layers: int = 1): torch.manual_seed(0) scheduler = FlowMatchEulerDiscreteScheduler() - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") components = { diff --git a/tests/pipelines/ltx/test_ltx_condition.py b/tests/pipelines/ltx/test_ltx_condition.py index f5dfb0186209..b469662241fc 100644 --- a/tests/pipelines/ltx/test_ltx_condition.py +++ b/tests/pipelines/ltx/test_ltx_condition.py @@ -17,7 +17,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import ( AutoencoderKLLTXVideo, @@ -92,7 +92,8 @@ def get_dummy_components(self): torch.manual_seed(0) scheduler = FlowMatchEulerDiscreteScheduler() - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") components = { diff --git a/tests/pipelines/ltx/test_ltx_image2video.py b/tests/pipelines/ltx/test_ltx_image2video.py index 2702993d4a59..7407c8bef5ea 100644 --- a/tests/pipelines/ltx/test_ltx_image2video.py +++ b/tests/pipelines/ltx/test_ltx_image2video.py @@ -17,7 +17,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import ( AutoencoderKLLTXVideo, @@ -91,7 +91,8 @@ def get_dummy_components(self): torch.manual_seed(0) scheduler = FlowMatchEulerDiscreteScheduler() - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") components = { diff --git a/tests/pipelines/ltx2/__init__.py b/tests/pipelines/ltx2/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/pipelines/ltx2/test_ltx2.py b/tests/pipelines/ltx2/test_ltx2.py new file mode 100644 index 000000000000..7d1a3bfc9987 --- /dev/null +++ b/tests/pipelines/ltx2/test_ltx2.py @@ -0,0 +1,289 @@ +# Copyright 2025 The HuggingFace Team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import torch +from transformers import AutoTokenizer, Gemma3ForConditionalGeneration + +from diffusers import ( + AutoencoderKLLTX2Audio, + AutoencoderKLLTX2Video, + FlowMatchEulerDiscreteScheduler, + LTX2Pipeline, + LTX2VideoTransformer3DModel, +) +from diffusers.pipelines.ltx2 import LTX2TextConnectors +from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder + +from ...testing_utils import enable_full_determinism +from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS +from ..test_pipelines_common import PipelineTesterMixin + + +enable_full_determinism() + + +class LTX2PipelineFastTests(PipelineTesterMixin, unittest.TestCase): + pipeline_class = LTX2Pipeline + params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} + batch_params = TEXT_TO_IMAGE_BATCH_PARAMS + image_params = TEXT_TO_IMAGE_IMAGE_PARAMS + image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS + required_optional_params = frozenset( + [ + "num_inference_steps", + "generator", + "latents", + "audio_latents", + "output_type", + "return_dict", + "callback_on_step_end", + "callback_on_step_end_tensor_inputs", + ] + ) + test_attention_slicing = False + test_xformers_attention = False + supports_dduf = False + + base_text_encoder_ckpt_id = "hf-internal-testing/tiny-gemma3" + + def get_dummy_components(self): + tokenizer = AutoTokenizer.from_pretrained(self.base_text_encoder_ckpt_id) + text_encoder = Gemma3ForConditionalGeneration.from_pretrained(self.base_text_encoder_ckpt_id) + + torch.manual_seed(0) + transformer = LTX2VideoTransformer3DModel( + in_channels=4, + out_channels=4, + patch_size=1, + patch_size_t=1, + num_attention_heads=2, + attention_head_dim=8, + cross_attention_dim=16, + audio_in_channels=4, + audio_out_channels=4, + audio_num_attention_heads=2, + audio_attention_head_dim=4, + audio_cross_attention_dim=8, + num_layers=2, + qk_norm="rms_norm_across_heads", + caption_channels=text_encoder.config.text_config.hidden_size, + rope_double_precision=False, + rope_type="split", + ) + + torch.manual_seed(0) + connectors = LTX2TextConnectors( + caption_channels=text_encoder.config.text_config.hidden_size, + text_proj_in_factor=text_encoder.config.text_config.num_hidden_layers + 1, + video_connector_num_attention_heads=4, + video_connector_attention_head_dim=8, + video_connector_num_layers=1, + video_connector_num_learnable_registers=None, + audio_connector_num_attention_heads=4, + audio_connector_attention_head_dim=8, + audio_connector_num_layers=1, + audio_connector_num_learnable_registers=None, + connector_rope_base_seq_len=32, + rope_theta=10000.0, + rope_double_precision=False, + causal_temporal_positioning=False, + rope_type="split", + ) + + torch.manual_seed(0) + vae = AutoencoderKLLTX2Video( + in_channels=3, + out_channels=3, + latent_channels=4, + block_out_channels=(8,), + decoder_block_out_channels=(8,), + layers_per_block=(1,), + decoder_layers_per_block=(1, 1), + spatio_temporal_scaling=(True,), + decoder_spatio_temporal_scaling=(True,), + decoder_inject_noise=(False, False), + downsample_type=("spatial",), + upsample_residual=(False,), + upsample_factor=(1,), + timestep_conditioning=False, + patch_size=1, + patch_size_t=1, + encoder_causal=True, + decoder_causal=False, + ) + vae.use_framewise_encoding = False + vae.use_framewise_decoding = False + + torch.manual_seed(0) + audio_vae = AutoencoderKLLTX2Audio( + base_channels=4, + output_channels=2, + ch_mult=(1,), + num_res_blocks=1, + attn_resolutions=None, + in_channels=2, + resolution=32, + latent_channels=2, + norm_type="pixel", + causality_axis="height", + dropout=0.0, + mid_block_add_attention=False, + sample_rate=16000, + mel_hop_length=160, + is_causal=True, + mel_bins=8, + ) + + torch.manual_seed(0) + vocoder = LTX2Vocoder( + in_channels=audio_vae.config.output_channels * audio_vae.config.mel_bins, + hidden_channels=32, + out_channels=2, + upsample_kernel_sizes=[4, 4], + upsample_factors=[2, 2], + resnet_kernel_sizes=[3], + resnet_dilations=[[1, 3, 5]], + leaky_relu_negative_slope=0.1, + output_sampling_rate=16000, + ) + + scheduler = FlowMatchEulerDiscreteScheduler() + + components = { + "transformer": transformer, + "vae": vae, + "audio_vae": audio_vae, + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + "connectors": connectors, + "vocoder": vocoder, + } + + return components + + def get_dummy_inputs(self, device, seed=0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device=device).manual_seed(seed) + + inputs = { + "prompt": "a robot dancing", + "negative_prompt": "", + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 1.0, + "height": 32, + "width": 32, + "num_frames": 5, + "frame_rate": 25.0, + "max_sequence_length": 16, + "output_type": "pt", + } + + return inputs + + def test_inference(self): + device = "cpu" + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + output = pipe(**inputs) + video = output.frames + audio = output.audio + + self.assertEqual(video.shape, (1, 5, 3, 32, 32)) + self.assertEqual(audio.shape[0], 1) + self.assertEqual(audio.shape[1], components["vocoder"].config.out_channels) + + # fmt: off + expected_video_slice = torch.tensor( + [ + 0.4331, 0.6203, 0.3245, 0.7294, 0.4822, 0.5703, 0.2999, 0.7700, 0.4961, 0.4242, 0.4581, 0.4351, 0.1137, 0.4437, 0.6304, 0.3184 + ] + ) + expected_audio_slice = torch.tensor( + [ + 0.0263, 0.0528, 0.1217, 0.1104, 0.1632, 0.1072, 0.1789, 0.0949, 0.0672, -0.0069, 0.0688, 0.0097, 0.0808, 0.1231, 0.0986, 0.0739 + ] + ) + # fmt: on + + video = video.flatten() + audio = audio.flatten() + generated_video_slice = torch.cat([video[:8], video[-8:]]) + generated_audio_slice = torch.cat([audio[:8], audio[-8:]]) + + assert torch.allclose(expected_video_slice, generated_video_slice, atol=1e-4, rtol=1e-4) + assert torch.allclose(expected_audio_slice, generated_audio_slice, atol=1e-4, rtol=1e-4) + + def test_two_stages_inference(self): + device = "cpu" + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + inputs["output_type"] = "latent" + first_stage_output = pipe(**inputs) + video_latent = first_stage_output.frames + audio_latent = first_stage_output.audio + + self.assertEqual(video_latent.shape, (1, 4, 3, 16, 16)) + self.assertEqual(audio_latent.shape, (1, 2, 5, 2)) + self.assertEqual(audio_latent.shape[1], components["vocoder"].config.out_channels) + + inputs["latents"] = video_latent + inputs["audio_latents"] = audio_latent + inputs["output_type"] = "pt" + second_stage_output = pipe(**inputs) + video = second_stage_output.frames + audio = second_stage_output.audio + + self.assertEqual(video.shape, (1, 5, 3, 32, 32)) + self.assertEqual(audio.shape[0], 1) + self.assertEqual(audio.shape[1], components["vocoder"].config.out_channels) + + # fmt: off + expected_video_slice = torch.tensor( + [ + 0.5514, 0.5943, 0.4260, 0.5971, 0.4306, 0.6369, 0.3124, 0.6964, 0.5419, 0.2412, 0.3882, 0.4504, 0.1941, 0.3404, 0.6037, 0.2464 + ] + ) + expected_audio_slice = torch.tensor( + [ + 0.0252, 0.0526, 0.1211, 0.1119, 0.1638, 0.1042, 0.1776, 0.0948, 0.0672, -0.0069, 0.0688, 0.0097, 0.0808, 0.1231, 0.0986, 0.0739 + ] + ) + # fmt: on + + video = video.flatten() + audio = audio.flatten() + generated_video_slice = torch.cat([video[:8], video[-8:]]) + generated_audio_slice = torch.cat([audio[:8], audio[-8:]]) + + assert torch.allclose(expected_video_slice, generated_video_slice, atol=1e-4, rtol=1e-4) + assert torch.allclose(expected_audio_slice, generated_audio_slice, atol=1e-4, rtol=1e-4) + + def test_inference_batch_single_identical(self): + self._test_inference_batch_single_identical(batch_size=2, expected_max_diff=2e-2) diff --git a/tests/pipelines/ltx2/test_ltx2_image2video.py b/tests/pipelines/ltx2/test_ltx2_image2video.py new file mode 100644 index 000000000000..92c000c7bf7c --- /dev/null +++ b/tests/pipelines/ltx2/test_ltx2_image2video.py @@ -0,0 +1,356 @@ +# Copyright 2025 The HuggingFace Team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import torch +from transformers import AutoTokenizer, Gemma3ForConditionalGeneration + +from diffusers import ( + AutoencoderKLLTX2Audio, + AutoencoderKLLTX2Video, + FlowMatchEulerDiscreteScheduler, + LTX2ImageToVideoPipeline, + LTX2VideoTransformer3DModel, +) +from diffusers.pipelines.ltx2 import LTX2LatentUpsamplePipeline, LTX2TextConnectors +from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel +from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder + +from ...testing_utils import enable_full_determinism +from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS +from ..test_pipelines_common import PipelineTesterMixin + + +enable_full_determinism() + + +class LTX2ImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase): + pipeline_class = LTX2ImageToVideoPipeline + params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} + batch_params = TEXT_TO_IMAGE_BATCH_PARAMS.union({"image"}) + image_params = TEXT_TO_IMAGE_IMAGE_PARAMS + image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS + required_optional_params = frozenset( + [ + "num_inference_steps", + "generator", + "latents", + "audio_latents", + "return_dict", + "callback_on_step_end", + "callback_on_step_end_tensor_inputs", + ] + ) + test_attention_slicing = False + test_xformers_attention = False + supports_dduf = False + + base_text_encoder_ckpt_id = "hf-internal-testing/tiny-gemma3" + + def get_dummy_components(self): + tokenizer = AutoTokenizer.from_pretrained(self.base_text_encoder_ckpt_id) + text_encoder = Gemma3ForConditionalGeneration.from_pretrained(self.base_text_encoder_ckpt_id) + + torch.manual_seed(0) + transformer = LTX2VideoTransformer3DModel( + in_channels=4, + out_channels=4, + patch_size=1, + patch_size_t=1, + num_attention_heads=2, + attention_head_dim=8, + cross_attention_dim=16, + audio_in_channels=4, + audio_out_channels=4, + audio_num_attention_heads=2, + audio_attention_head_dim=4, + audio_cross_attention_dim=8, + num_layers=2, + qk_norm="rms_norm_across_heads", + caption_channels=text_encoder.config.text_config.hidden_size, + rope_double_precision=False, + rope_type="split", + ) + + torch.manual_seed(0) + connectors = LTX2TextConnectors( + caption_channels=text_encoder.config.text_config.hidden_size, + text_proj_in_factor=text_encoder.config.text_config.num_hidden_layers + 1, + video_connector_num_attention_heads=4, + video_connector_attention_head_dim=8, + video_connector_num_layers=1, + video_connector_num_learnable_registers=None, + audio_connector_num_attention_heads=4, + audio_connector_attention_head_dim=8, + audio_connector_num_layers=1, + audio_connector_num_learnable_registers=None, + connector_rope_base_seq_len=32, + rope_theta=10000.0, + rope_double_precision=False, + causal_temporal_positioning=False, + rope_type="split", + ) + + torch.manual_seed(0) + vae = AutoencoderKLLTX2Video( + in_channels=3, + out_channels=3, + latent_channels=4, + block_out_channels=(8,), + decoder_block_out_channels=(8,), + layers_per_block=(1,), + decoder_layers_per_block=(1, 1), + spatio_temporal_scaling=(True,), + decoder_spatio_temporal_scaling=(True,), + decoder_inject_noise=(False, False), + downsample_type=("spatial",), + upsample_residual=(False,), + upsample_factor=(1,), + timestep_conditioning=False, + patch_size=1, + patch_size_t=1, + encoder_causal=True, + decoder_causal=False, + ) + vae.use_framewise_encoding = False + vae.use_framewise_decoding = False + + torch.manual_seed(0) + audio_vae = AutoencoderKLLTX2Audio( + base_channels=4, + output_channels=2, + ch_mult=(1,), + num_res_blocks=1, + attn_resolutions=None, + in_channels=2, + resolution=32, + latent_channels=2, + norm_type="pixel", + causality_axis="height", + dropout=0.0, + mid_block_add_attention=False, + sample_rate=16000, + mel_hop_length=160, + is_causal=True, + mel_bins=8, + ) + + torch.manual_seed(0) + vocoder = LTX2Vocoder( + in_channels=audio_vae.config.output_channels * audio_vae.config.mel_bins, + hidden_channels=32, + out_channels=2, + upsample_kernel_sizes=[4, 4], + upsample_factors=[2, 2], + resnet_kernel_sizes=[3], + resnet_dilations=[[1, 3, 5]], + leaky_relu_negative_slope=0.1, + output_sampling_rate=16000, + ) + + scheduler = FlowMatchEulerDiscreteScheduler() + + components = { + "transformer": transformer, + "vae": vae, + "audio_vae": audio_vae, + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + "connectors": connectors, + "vocoder": vocoder, + } + + return components + + def get_dummy_upsample_component(self, in_channels=4, mid_channels=32, num_blocks_per_stage=1): + upsampler = LTX2LatentUpsamplerModel( + in_channels=in_channels, + mid_channels=mid_channels, + num_blocks_per_stage=num_blocks_per_stage, + ) + + return upsampler + + def get_dummy_inputs(self, device, seed=0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device=device).manual_seed(seed) + + image = torch.rand((1, 3, 32, 32), generator=generator, device=device) + + inputs = { + "image": image, + "prompt": "a robot dancing", + "negative_prompt": "", + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 1.0, + "height": 32, + "width": 32, + "num_frames": 5, + "frame_rate": 25.0, + "max_sequence_length": 16, + "output_type": "pt", + } + + return inputs + + def test_inference(self): + device = "cpu" + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + output = pipe(**inputs) + video = output.frames + audio = output.audio + + self.assertEqual(video.shape, (1, 5, 3, 32, 32)) + self.assertEqual(audio.shape[0], 1) + self.assertEqual(audio.shape[1], components["vocoder"].config.out_channels) + + # fmt: off + expected_video_slice = torch.tensor( + [ + 0.3573, 0.8382, 0.3581, 0.6114, 0.3682, 0.7969, 0.2552, 0.6399, 0.3113, 0.1497, 0.3249, 0.5395, 0.3498, 0.4526, 0.4536, 0.4555 + ] + ) + expected_audio_slice = torch.tensor( + [ + 0.0294, 0.0498, 0.1269, 0.1135, 0.1639, 0.1116, 0.1730, 0.0931, 0.0672, -0.0069, 0.0688, 0.0097, 0.0808, 0.1231, 0.0986, 0.0739 + ] + ) + # fmt: on + + video = video.flatten() + audio = audio.flatten() + generated_video_slice = torch.cat([video[:8], video[-8:]]) + generated_audio_slice = torch.cat([audio[:8], audio[-8:]]) + + assert torch.allclose(expected_video_slice, generated_video_slice, atol=1e-4, rtol=1e-4) + assert torch.allclose(expected_audio_slice, generated_audio_slice, atol=1e-4, rtol=1e-4) + + def test_two_stages_inference(self): + device = "cpu" + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + inputs["output_type"] = "latent" + first_stage_output = pipe(**inputs) + video_latent = first_stage_output.frames + audio_latent = first_stage_output.audio + + self.assertEqual(video_latent.shape, (1, 4, 3, 16, 16)) + self.assertEqual(audio_latent.shape, (1, 2, 5, 2)) + self.assertEqual(audio_latent.shape[1], components["vocoder"].config.out_channels) + + inputs["latents"] = video_latent + inputs["audio_latents"] = audio_latent + inputs["output_type"] = "pt" + second_stage_output = pipe(**inputs) + video = second_stage_output.frames + audio = second_stage_output.audio + + self.assertEqual(video.shape, (1, 5, 3, 32, 32)) + self.assertEqual(audio.shape[0], 1) + self.assertEqual(audio.shape[1], components["vocoder"].config.out_channels) + + # fmt: off + expected_video_slice = torch.tensor( + [ + 0.2665, 0.6915, 0.2939, 0.6767, 0.2552, 0.6215, 0.1765, 0.6248, 0.2800, 0.2356, 0.3480, 0.5395, 0.3190, 0.4128, 0.4784, 0.4086 + ] + ) + expected_audio_slice = torch.tensor( + [ + 0.0273, 0.0490, 0.1253, 0.1129, 0.1655, 0.1057, 0.1707, 0.0943, 0.0672, -0.0069, 0.0688, 0.0097, 0.0808, 0.1231, 0.0986, 0.0739 + ] + ) + # fmt: on + + video = video.flatten() + audio = audio.flatten() + generated_video_slice = torch.cat([video[:8], video[-8:]]) + generated_audio_slice = torch.cat([audio[:8], audio[-8:]]) + + assert torch.allclose(expected_video_slice, generated_video_slice, atol=1e-4, rtol=1e-4) + assert torch.allclose(expected_audio_slice, generated_audio_slice, atol=1e-4, rtol=1e-4) + + def test_two_stages_inference_with_upsampler(self): + device = "cpu" + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + inputs["output_type"] = "latent" + first_stage_output = pipe(**inputs) + video_latent = first_stage_output.frames + audio_latent = first_stage_output.audio + + self.assertEqual(video_latent.shape, (1, 4, 3, 16, 16)) + self.assertEqual(audio_latent.shape, (1, 2, 5, 2)) + self.assertEqual(audio_latent.shape[1], components["vocoder"].config.out_channels) + + upsampler = self.get_dummy_upsample_component(in_channels=video_latent.shape[1]) + upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=upsampler) + upscaled_video_latent = upsample_pipe(latents=video_latent, output_type="latent", return_dict=False)[0] + self.assertEqual(upscaled_video_latent.shape, (1, 4, 3, 32, 32)) + + inputs["latents"] = upscaled_video_latent + inputs["audio_latents"] = audio_latent + inputs["output_type"] = "pt" + second_stage_output = pipe(**inputs) + video = second_stage_output.frames + audio = second_stage_output.audio + + self.assertEqual(video.shape, (1, 5, 3, 64, 64)) + self.assertEqual(audio.shape[0], 1) + self.assertEqual(audio.shape[1], components["vocoder"].config.out_channels) + + # fmt: off + expected_video_slice = torch.tensor( + [ + 0.4497, 0.6757, 0.4219, 0.7686, 0.4525, 0.6483, 0.3969, 0.7404, 0.3541, 0.3039, 0.4592, 0.3521, 0.3665, 0.2785, 0.3336, 0.3079 + ] + ) + expected_audio_slice = torch.tensor( + [ + 0.0271, 0.0492, 0.1249, 0.1126, 0.1661, 0.1060, 0.1717, 0.0944, 0.0672, -0.0069, 0.0688, 0.0097, 0.0808, 0.1231, 0.0986, 0.0739 + ] + ) + # fmt: on + + video = video.flatten() + audio = audio.flatten() + generated_video_slice = torch.cat([video[:8], video[-8:]]) + generated_audio_slice = torch.cat([audio[:8], audio[-8:]]) + + assert torch.allclose(expected_video_slice, generated_video_slice, atol=1e-4, rtol=1e-4) + assert torch.allclose(expected_audio_slice, generated_audio_slice, atol=1e-4, rtol=1e-4) + + def test_inference_batch_single_identical(self): + self._test_inference_batch_single_identical(batch_size=2, expected_max_diff=2e-2) diff --git a/tests/pipelines/mochi/test_mochi.py b/tests/pipelines/mochi/test_mochi.py index 5615720a9343..df6f33f79cb2 100644 --- a/tests/pipelines/mochi/test_mochi.py +++ b/tests/pipelines/mochi/test_mochi.py @@ -18,7 +18,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKLMochi, FlowMatchEulerDiscreteScheduler, MochiPipeline, MochiTransformer3DModel @@ -89,7 +89,8 @@ def get_dummy_components(self, num_layers: int = 2): torch.manual_seed(0) scheduler = FlowMatchEulerDiscreteScheduler() - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") components = { @@ -207,6 +208,9 @@ def test_attention_slicing_forward_pass( return components = self.get_dummy_components() + for key in components: + if "text_encoder" in key and hasattr(components[key], "eval"): + components[key].eval() pipe = self.pipeline_class(**components) for component in pipe.components.values(): if hasattr(component, "set_default_attn_processor"): diff --git a/tests/pipelines/pag/test_pag_hunyuan_dit.py b/tests/pipelines/pag/test_pag_hunyuan_dit.py index f268a614f85c..38686ee448de 100644 --- a/tests/pipelines/pag/test_pag_hunyuan_dit.py +++ b/tests/pipelines/pag/test_pag_hunyuan_dit.py @@ -19,7 +19,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, BertModel, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, BertModel, T5EncoderModel from diffusers import ( AutoencoderKL, @@ -67,7 +67,9 @@ def get_dummy_components(self): scheduler = DDPMScheduler() text_encoder = BertModel.from_pretrained("hf-internal-testing/tiny-random-BertModel") tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-BertModel") - text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + torch.manual_seed(0) + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_2 = T5EncoderModel(config) tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") components = { diff --git a/tests/pipelines/pag/test_pag_pixart_sigma.py b/tests/pipelines/pag/test_pag_pixart_sigma.py index c04ebad08fdc..9bc2f6eed395 100644 --- a/tests/pipelines/pag/test_pag_pixart_sigma.py +++ b/tests/pipelines/pag/test_pag_pixart_sigma.py @@ -19,7 +19,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel import diffusers from diffusers import ( @@ -80,7 +80,8 @@ def get_dummy_components(self): vae = AutoencoderKL() scheduler = DDIMScheduler() - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") diff --git a/tests/pipelines/pag/test_pag_sana.py b/tests/pipelines/pag/test_pag_sana.py index 5408595c729d..366267dc092b 100644 --- a/tests/pipelines/pag/test_pag_sana.py +++ b/tests/pipelines/pag/test_pag_sana.py @@ -325,13 +325,13 @@ def test_pag_applied_layers(self): # TODO(aryan): Create a dummy gemma model with smol vocab size @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_consistent(self): pass @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_single_identical(self): pass diff --git a/tests/pipelines/pag/test_pag_sd3.py b/tests/pipelines/pag/test_pag_sd3.py index 26e6ca099286..7f755ea8e170 100644 --- a/tests/pipelines/pag/test_pag_sd3.py +++ b/tests/pipelines/pag/test_pag_sd3.py @@ -3,7 +3,14 @@ import numpy as np import torch -from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderModel +from transformers import ( + AutoConfig, + AutoTokenizer, + CLIPTextConfig, + CLIPTextModelWithProjection, + CLIPTokenizer, + T5EncoderModel, +) from diffusers import ( AutoencoderKL, @@ -73,7 +80,9 @@ def get_dummy_components(self): torch.manual_seed(0) text_encoder_2 = CLIPTextModelWithProjection(clip_text_encoder_config) - text_encoder_3 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + torch.manual_seed(0) + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_3 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") diff --git a/tests/pipelines/pag/test_pag_sd3_img2img.py b/tests/pipelines/pag/test_pag_sd3_img2img.py index 19a36e283de4..e4146b87803c 100644 --- a/tests/pipelines/pag/test_pag_sd3_img2img.py +++ b/tests/pipelines/pag/test_pag_sd3_img2img.py @@ -5,7 +5,14 @@ import numpy as np import torch -from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderModel +from transformers import ( + AutoConfig, + AutoTokenizer, + CLIPTextConfig, + CLIPTextModelWithProjection, + CLIPTokenizer, + T5EncoderModel, +) from diffusers import ( AutoencoderKL, @@ -84,7 +91,9 @@ def get_dummy_components(self): torch.manual_seed(0) text_encoder_2 = CLIPTextModelWithProjection(clip_text_encoder_config) - text_encoder_3 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + torch.manual_seed(0) + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_3 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") diff --git a/tests/pipelines/pixart_alpha/test_pixart.py b/tests/pipelines/pixart_alpha/test_pixart.py index fd41c9887dcc..037a9f44f31e 100644 --- a/tests/pipelines/pixart_alpha/test_pixart.py +++ b/tests/pipelines/pixart_alpha/test_pixart.py @@ -19,7 +19,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import ( AutoencoderKL, @@ -77,7 +77,10 @@ def get_dummy_components(self): vae = AutoencoderKL() scheduler = DDIMScheduler() - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + + torch.manual_seed(0) + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") diff --git a/tests/pipelines/pixart_sigma/test_pixart.py b/tests/pipelines/pixart_sigma/test_pixart.py index 6e8535062a79..51eebadd0ed0 100644 --- a/tests/pipelines/pixart_sigma/test_pixart.py +++ b/tests/pipelines/pixart_sigma/test_pixart.py @@ -19,7 +19,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import ( AutoencoderKL, @@ -83,7 +83,10 @@ def get_dummy_components(self): vae = AutoencoderKL() scheduler = DDIMScheduler() - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + + torch.manual_seed(0) + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") diff --git a/tests/pipelines/prx/test_pipeline_prx.py b/tests/pipelines/prx/test_pipeline_prx.py index 46c6a5760e22..20c2727b3156 100644 --- a/tests/pipelines/prx/test_pipeline_prx.py +++ b/tests/pipelines/prx/test_pipeline_prx.py @@ -1,7 +1,6 @@ import unittest import numpy as np -import pytest import torch from transformers import AutoTokenizer from transformers.models.t5gemma.configuration_t5gemma import T5GemmaConfig, T5GemmaModuleConfig @@ -11,17 +10,11 @@ from diffusers.models.transformers.transformer_prx import PRXTransformer2DModel from diffusers.pipelines.prx.pipeline_prx import PRXPipeline from diffusers.schedulers import FlowMatchEulerDiscreteScheduler -from diffusers.utils import is_transformers_version from ..pipeline_params import TEXT_TO_IMAGE_PARAMS from ..test_pipelines_common import PipelineTesterMixin -@pytest.mark.xfail( - condition=is_transformers_version(">", "4.57.1"), - reason="See https://github.com/huggingface/diffusers/pull/12456#issuecomment-3424228544", - strict=False, -) class PRXPipelineFastTests(PipelineTesterMixin, unittest.TestCase): pipeline_class = PRXPipeline params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} @@ -99,7 +92,7 @@ def get_dummy_components(self): } encoder_config = T5GemmaModuleConfig(**encoder_params) text_encoder_config = T5GemmaConfig(encoder=encoder_config, is_encoder_decoder=False, **encoder_params) - text_encoder = T5GemmaEncoder(text_encoder_config) + text_encoder = T5GemmaEncoder(text_encoder_config.encoder) return { "transformer": transformer, @@ -263,3 +256,27 @@ def test_inference_with_autoencoder_dc(self): expected_image = torch.zeros(3, 32, 32) max_diff = np.abs(generated_image - expected_image).max() self.assertLessEqual(max_diff, 1e10) + + @unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.") + def test_save_load_dduf(self): + pass + + @unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.") + def test_loading_with_variants(self): + pass + + @unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.") + def test_pipeline_with_accelerator_device_map(self): + pass + + @unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.") + def test_save_load_local(self): + pass + + @unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.") + def test_save_load_optional_components(self): + pass + + @unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.") + def test_torch_dtype_dict(self): + pass diff --git a/tests/pipelines/qwenimage/test_qwenimage.py b/tests/pipelines/qwenimage/test_qwenimage.py index 8ebfe7d08bc1..ff53d1c234c7 100644 --- a/tests/pipelines/qwenimage/test_qwenimage.py +++ b/tests/pipelines/qwenimage/test_qwenimage.py @@ -113,7 +113,7 @@ def get_dummy_components(self): vision_start_token_id=151652, vision_token_id=151654, ) - text_encoder = Qwen2_5_VLForConditionalGeneration(config) + text_encoder = Qwen2_5_VLForConditionalGeneration(config).eval() tokenizer = Qwen2Tokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") components = { @@ -160,12 +160,12 @@ def test_inference(self): self.assertEqual(generated_image.shape, (3, 32, 32)) # fmt: off - expected_slice = torch.tensor([0.56331, 0.63677, 0.6015, 0.56369, 0.58166, 0.55277, 0.57176, 0.63261, 0.41466, 0.35561, 0.56229, 0.48334, 0.49714, 0.52622, 0.40872, 0.50208]) + expected_slice = torch.tensor([0.5633, 0.6368, 0.6015, 0.5637, 0.5817, 0.5528, 0.5718, 0.6326, 0.4147, 0.3556, 0.5623, 0.4833, 0.4971, 0.5262, 0.4087, 0.5021]) # fmt: on generated_slice = generated_image.flatten() generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) - self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=1e-3)) + self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=5e-3)) def test_inference_batch_single_identical(self): self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-1) diff --git a/tests/pipelines/qwenimage/test_qwenimage_controlnet.py b/tests/pipelines/qwenimage/test_qwenimage_controlnet.py index 188106b49b84..59a2dd497184 100644 --- a/tests/pipelines/qwenimage/test_qwenimage_controlnet.py +++ b/tests/pipelines/qwenimage/test_qwenimage_controlnet.py @@ -211,7 +211,7 @@ def test_qwen_controlnet(self): generated_slice = generated_image.flatten() generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) - self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=1e-3)) + self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=5e-3)) def test_qwen_controlnet_multicondition(self): device = "cpu" @@ -255,7 +255,7 @@ def test_qwen_controlnet_multicondition(self): generated_slice = generated_image.flatten() generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) - self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=1e-3)) + self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=5e-3)) def test_attention_slicing_forward_pass( self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3 diff --git a/tests/pipelines/qwenimage/test_qwenimage_edit.py b/tests/pipelines/qwenimage/test_qwenimage_edit.py index 058548cf5f1b..8184c895f825 100644 --- a/tests/pipelines/qwenimage/test_qwenimage_edit.py +++ b/tests/pipelines/qwenimage/test_qwenimage_edit.py @@ -115,7 +115,7 @@ def get_dummy_components(self): vision_start_token_id=151652, vision_token_id=151654, ) - text_encoder = Qwen2_5_VLForConditionalGeneration(config) + text_encoder = Qwen2_5_VLForConditionalGeneration(config).eval() tokenizer = Qwen2Tokenizer.from_pretrained(tiny_ckpt_id) components = { @@ -163,12 +163,12 @@ def test_inference(self): self.assertEqual(generated_image.shape, (3, 32, 32)) # fmt: off - expected_slice = torch.tensor([[0.5637, 0.6341, 0.6001, 0.5620, 0.5794, 0.5498, 0.5757, 0.6389, 0.4174, 0.3597, 0.5649, 0.4894, 0.4969, 0.5255, 0.4083, 0.4986]]) + expected_slice = torch.tensor([0.5637, 0.6341, 0.6001, 0.5620, 0.5794, 0.5498, 0.5757, 0.6389, 0.4174, 0.3597, 0.5649, 0.4894, 0.4969, 0.5255, 0.4083, 0.4986]) # fmt: on generated_slice = generated_image.flatten() generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) - self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=1e-3)) + self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=5e-3)) def test_inference_batch_single_identical(self): self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-1) diff --git a/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py b/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py index 6faf34728286..e8bc694ced84 100644 --- a/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py +++ b/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py @@ -164,7 +164,7 @@ def test_inference(self): self.assertEqual(generated_image.shape, (3, 32, 32)) # fmt: off - expected_slice = torch.tensor([[0.5637, 0.6341, 0.6001, 0.5620, 0.5794, 0.5498, 0.5757, 0.6389, 0.4174, 0.3597, 0.5649, 0.4894, 0.4969, 0.5255, 0.4083, 0.4986]]) + expected_slice = torch.tensor([0.5640, 0.6339, 0.5997, 0.5607, 0.5799, 0.5496, 0.5760, 0.6393, 0.4172, 0.3595, 0.5655, 0.4896, 0.4971, 0.5255, 0.4088, 0.4987]) # fmt: on generated_slice = generated_image.flatten() diff --git a/tests/pipelines/sana/test_sana.py b/tests/pipelines/sana/test_sana.py index f23303c966e5..fc533dd7bf5c 100644 --- a/tests/pipelines/sana/test_sana.py +++ b/tests/pipelines/sana/test_sana.py @@ -290,13 +290,13 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2): # TODO(aryan): Create a dummy gemma model with smol vocab size @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_consistent(self): pass @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_single_identical(self): pass diff --git a/tests/pipelines/sana/test_sana_controlnet.py b/tests/pipelines/sana/test_sana_controlnet.py index df14d935edf5..0f432a998ddc 100644 --- a/tests/pipelines/sana/test_sana_controlnet.py +++ b/tests/pipelines/sana/test_sana_controlnet.py @@ -309,13 +309,13 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2): # TODO(aryan): Create a dummy gemma model with smol vocab size @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_consistent(self): pass @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_single_identical(self): pass diff --git a/tests/pipelines/sana/test_sana_sprint.py b/tests/pipelines/sana/test_sana_sprint.py index 0d45205ea8c7..9f88d0f344ae 100644 --- a/tests/pipelines/sana/test_sana_sprint.py +++ b/tests/pipelines/sana/test_sana_sprint.py @@ -283,13 +283,13 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2): # TODO(aryan): Create a dummy gemma model with smol vocab size @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_consistent(self): pass @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_single_identical(self): pass diff --git a/tests/pipelines/sana/test_sana_sprint_img2img.py b/tests/pipelines/sana/test_sana_sprint_img2img.py index 5de5c7f44606..312e2b2b8080 100644 --- a/tests/pipelines/sana/test_sana_sprint_img2img.py +++ b/tests/pipelines/sana/test_sana_sprint_img2img.py @@ -295,13 +295,13 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2): # TODO(aryan): Create a dummy gemma model with smol vocab size @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_consistent(self): pass @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_single_identical(self): pass diff --git a/tests/pipelines/sana_video/test_sana_video.py b/tests/pipelines/sana_video/test_sana_video.py index 9f360a942a64..dd897a848be8 100644 --- a/tests/pipelines/sana_video/test_sana_video.py +++ b/tests/pipelines/sana_video/test_sana_video.py @@ -185,13 +185,13 @@ def test_save_load_local(self, expected_max_difference=5e-4): # TODO(aryan): Create a dummy gemma model with smol vocab size @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_consistent(self): pass @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_single_identical(self): pass diff --git a/tests/pipelines/sana_video/test_sana_video_i2v.py b/tests/pipelines/sana_video/test_sana_video_i2v.py index 36a646ca528f..2232968eb2ca 100644 --- a/tests/pipelines/sana_video/test_sana_video_i2v.py +++ b/tests/pipelines/sana_video/test_sana_video_i2v.py @@ -196,13 +196,13 @@ def test_save_load_local(self, expected_max_difference=5e-4): # TODO(aryan): Create a dummy gemma model with smol vocab size @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_consistent(self): pass @unittest.skip( - "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_single_identical(self): pass diff --git a/tests/pipelines/skyreels_v2/test_skyreels_v2.py b/tests/pipelines/skyreels_v2/test_skyreels_v2.py index 1bcec877c30d..b6adb3cc1f2c 100644 --- a/tests/pipelines/skyreels_v2/test_skyreels_v2.py +++ b/tests/pipelines/skyreels_v2/test_skyreels_v2.py @@ -18,20 +18,11 @@ import torch from transformers import AutoTokenizer, T5EncoderModel -from diffusers import ( - AutoencoderKLWan, - SkyReelsV2Pipeline, - SkyReelsV2Transformer3DModel, - UniPCMultistepScheduler, -) - -from ...testing_utils import ( - enable_full_determinism, -) +from diffusers import AutoencoderKLWan, SkyReelsV2Pipeline, SkyReelsV2Transformer3DModel, UniPCMultistepScheduler + +from ...testing_utils import enable_full_determinism from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS -from ..test_pipelines_common import ( - PipelineTesterMixin, -) +from ..test_pipelines_common import PipelineTesterMixin enable_full_determinism() diff --git a/tests/pipelines/skyreels_v2/test_skyreels_v2_df.py b/tests/pipelines/skyreels_v2/test_skyreels_v2_df.py index 74235d59efd6..213d085a6dca 100644 --- a/tests/pipelines/skyreels_v2/test_skyreels_v2_df.py +++ b/tests/pipelines/skyreels_v2/test_skyreels_v2_df.py @@ -25,13 +25,9 @@ UniPCMultistepScheduler, ) -from ...testing_utils import ( - enable_full_determinism, -) +from ...testing_utils import enable_full_determinism from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS -from ..test_pipelines_common import ( - PipelineTesterMixin, -) +from ..test_pipelines_common import PipelineTesterMixin enable_full_determinism() diff --git a/tests/pipelines/skyreels_v2/test_skyreels_v2_df_image_to_video.py b/tests/pipelines/skyreels_v2/test_skyreels_v2_df_image_to_video.py index f0cbc710df05..02daa61395c3 100644 --- a/tests/pipelines/skyreels_v2/test_skyreels_v2_df_image_to_video.py +++ b/tests/pipelines/skyreels_v2/test_skyreels_v2_df_image_to_video.py @@ -17,10 +17,7 @@ import numpy as np import torch from PIL import Image -from transformers import ( - AutoTokenizer, - T5EncoderModel, -) +from transformers import AutoTokenizer, T5EncoderModel from diffusers import ( AutoencoderKLWan, diff --git a/tests/pipelines/skyreels_v2/test_skyreels_v2_df_video_to_video.py b/tests/pipelines/skyreels_v2/test_skyreels_v2_df_video_to_video.py index 1b0b23318e63..cbfb86f7696e 100644 --- a/tests/pipelines/skyreels_v2/test_skyreels_v2_df_video_to_video.py +++ b/tests/pipelines/skyreels_v2/test_skyreels_v2_df_video_to_video.py @@ -27,14 +27,9 @@ UniPCMultistepScheduler, ) -from ...testing_utils import ( - enable_full_determinism, - torch_device, -) +from ...testing_utils import enable_full_determinism, torch_device from ..pipeline_params import TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS -from ..test_pipelines_common import ( - PipelineTesterMixin, -) +from ..test_pipelines_common import PipelineTesterMixin enable_full_determinism() diff --git a/tests/pipelines/stable_audio/test_stable_audio.py b/tests/pipelines/stable_audio/test_stable_audio.py index dd03f4d07f07..492aa92252de 100644 --- a/tests/pipelines/stable_audio/test_stable_audio.py +++ b/tests/pipelines/stable_audio/test_stable_audio.py @@ -19,10 +19,7 @@ import numpy as np import torch -from transformers import ( - T5EncoderModel, - T5Tokenizer, -) +from transformers import AutoConfig, T5EncoderModel, T5Tokenizer from diffusers import ( AutoencoderOobleck, @@ -111,7 +108,8 @@ def get_dummy_components(self): ) torch.manual_seed(0) t5_repo_id = "hf-internal-testing/tiny-random-T5ForConditionalGeneration" - text_encoder = T5EncoderModel.from_pretrained(t5_repo_id) + config = AutoConfig.from_pretrained(t5_repo_id) + text_encoder = T5EncoderModel(config) tokenizer = T5Tokenizer.from_pretrained(t5_repo_id, truncation=True, model_max_length=25) torch.manual_seed(0) diff --git a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py index 3ccefe3de35d..200c832d0941 100644 --- a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py +++ b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py @@ -3,7 +3,14 @@ import numpy as np import torch -from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderModel +from transformers import ( + AutoConfig, + AutoTokenizer, + CLIPTextConfig, + CLIPTextModelWithProjection, + CLIPTokenizer, + T5EncoderModel, +) from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, SD3Transformer2DModel, StableDiffusion3Pipeline @@ -72,7 +79,9 @@ def get_dummy_components(self): torch.manual_seed(0) text_encoder_2 = CLIPTextModelWithProjection(clip_text_encoder_config) - text_encoder_3 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + torch.manual_seed(0) + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_3 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") diff --git a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_img2img.py b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_img2img.py index 9025b1060c9e..3f46b341a09e 100644 --- a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_img2img.py +++ b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_img2img.py @@ -4,7 +4,14 @@ import numpy as np import torch -from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderModel +from transformers import ( + AutoConfig, + AutoTokenizer, + CLIPTextConfig, + CLIPTextModelWithProjection, + CLIPTokenizer, + T5EncoderModel, +) from diffusers import ( AutoencoderKL, @@ -73,7 +80,9 @@ def get_dummy_components(self): torch.manual_seed(0) text_encoder_2 = CLIPTextModelWithProjection(clip_text_encoder_config) - text_encoder_3 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + torch.manual_seed(0) + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_3 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") diff --git a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_inpaint.py b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_inpaint.py index 628930340294..a90ca21a801b 100644 --- a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_inpaint.py +++ b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_inpaint.py @@ -3,7 +3,14 @@ import numpy as np import torch -from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderModel +from transformers import ( + AutoConfig, + AutoTokenizer, + CLIPTextConfig, + CLIPTextModelWithProjection, + CLIPTokenizer, + T5EncoderModel, +) from diffusers import ( AutoencoderKL, @@ -73,7 +80,9 @@ def get_dummy_components(self): torch.manual_seed(0) text_encoder_2 = CLIPTextModelWithProjection(clip_text_encoder_config) - text_encoder_3 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + torch.manual_seed(0) + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_3 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") diff --git a/tests/pipelines/test_pipelines_common.py b/tests/pipelines/test_pipelines_common.py index 7db5f4da89ca..af3573ce84cb 100644 --- a/tests/pipelines/test_pipelines_common.py +++ b/tests/pipelines/test_pipelines_common.py @@ -5,7 +5,7 @@ import tempfile import unittest import uuid -from typing import Any, Callable, Dict, Union +from typing import Any, Callable, Dict import numpy as np import PIL.Image @@ -35,6 +35,7 @@ from diffusers.hooks import apply_group_offloading from diffusers.hooks.faster_cache import FasterCacheBlockHook, FasterCacheDenoiserHook from diffusers.hooks.first_block_cache import FirstBlockCacheConfig +from diffusers.hooks.mag_cache import MagCacheConfig from diffusers.hooks.pyramid_attention_broadcast import PyramidAttentionBroadcastHook from diffusers.hooks.taylorseer_cache import TaylorSeerCacheConfig from diffusers.image_processor import VaeImageProcessor @@ -1071,7 +1072,7 @@ def get_generator(self, seed): return generator @property - def pipeline_class(self) -> Union[Callable, DiffusionPipeline]: + def pipeline_class(self) -> Callable | DiffusionPipeline: raise NotImplementedError( "You need to set the attribute `pipeline_class = ClassNameOfPipeline` in the child test class. " "See existing pipeline tests for reference." @@ -1156,6 +1157,9 @@ def tearDown(self): def test_save_load_local(self, expected_max_difference=5e-4): components = self.get_dummy_components() + for key in components: + if "text_encoder" in key and hasattr(components[key], "eval"): + components[key].eval() pipe = self.pipeline_class(**components) for component in pipe.components.values(): if hasattr(component, "set_default_attn_processor"): @@ -1294,6 +1298,9 @@ def _test_inference_batch_single_identical( additional_params_copy_to_batched_inputs=["num_inference_steps"], ): components = self.get_dummy_components() + for key in components: + if "text_encoder" in key and hasattr(components[key], "eval"): + components[key].eval() pipe = self.pipeline_class(**components) for components in pipe.components.values(): if hasattr(components, "set_default_attn_processor"): @@ -1344,6 +1351,9 @@ def _test_inference_batch_single_identical( def test_dict_tuple_outputs_equivalent(self, expected_slice=None, expected_max_difference=1e-4): components = self.get_dummy_components() + for key in components: + if "text_encoder" in key and hasattr(components[key], "eval"): + components[key].eval() pipe = self.pipeline_class(**components) for component in pipe.components.values(): if hasattr(component, "set_default_attn_processor"): @@ -1476,6 +1486,9 @@ def test_save_load_optional_components(self, expected_max_difference=1e-4): if not self.pipeline_class._optional_components: return components = self.get_dummy_components() + for key in components: + if "text_encoder" in key and hasattr(components[key], "eval"): + components[key].eval() pipe = self.pipeline_class(**components) for component in pipe.components.values(): if hasattr(component, "set_default_attn_processor"): @@ -1556,6 +1569,9 @@ def _test_attention_slicing_forward_pass( return components = self.get_dummy_components() + for key in components: + if "text_encoder" in key and hasattr(components[key], "eval"): + components[key].eval() pipe = self.pipeline_class(**components) for component in pipe.components.values(): if hasattr(component, "set_default_attn_processor"): @@ -2064,7 +2080,16 @@ def is_nan(tensor): for component_name in model_components_pipe: pipe_component = model_components_pipe[component_name] pipe_loaded_component = model_components_pipe_loaded[component_name] - for p1, p2 in zip(pipe_component.parameters(), pipe_loaded_component.parameters()): + + model_loaded_params = dict(pipe_loaded_component.named_parameters()) + model_original_params = dict(pipe_component.named_parameters()) + + for name, p1 in model_original_params.items(): + # Skip tied weights that aren't saved with variants (transformers v5 behavior) + if name not in model_loaded_params: + continue + + p2 = model_loaded_params[name] # nan check for luminanext (mps). if not (is_nan(p1) and is_nan(p2)): self.assertTrue(torch.equal(p1, p2)) @@ -2088,6 +2113,9 @@ def test_encode_prompt_works_in_isolation(self, extra_required_param_value_dict= return components = self.get_dummy_components() + for key in components: + if "text_encoder" in key and hasattr(components[key], "eval"): + components[key].eval() # We initialize the pipeline with only text encoders and tokenizers, # mimicking a real-world scenario. @@ -2219,6 +2247,9 @@ def test_save_load_dduf(self, atol=1e-4, rtol=1e-4): from huggingface_hub import export_folder_as_dduf components = self.get_dummy_components() + for key in components: + if "text_encoder" in key and hasattr(components[key], "eval"): + components[key].eval() pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) @@ -2354,9 +2385,13 @@ def test_torch_dtype_dict(self): f"Component '{name}' has dtype {component.dtype} but expected {expected_dtype}", ) - @require_torch_accelerator def test_pipeline_with_accelerator_device_map(self, expected_max_difference=1e-4): components = self.get_dummy_components() + # Set text encoders to eval mode to match from_pretrained behavior + # This ensures deterministic outputs when models are loaded with device_map + for key in components: + if "text_encoder" in key and hasattr(components[key], "eval"): + components[key].eval() pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) @@ -2405,7 +2440,11 @@ def test_pipeline_level_group_offloading_sanity_checks(self): if name not in [exclude_module_name] and isinstance(component, torch.nn.Module): # `component.device` prints the `onload_device` type. We should probably override the # `device` property in `ModelMixin`. - component_device = next(component.parameters())[0].device + # Skip modules with no parameters (e.g., dummy safety checkers with only buffers) + params = list(component.parameters()) + if not params: + continue + component_device = params[0].device self.assertTrue(torch.device(component_device).type == torch.device(offload_device).type) @require_torch_accelerator @@ -2676,6 +2715,9 @@ def test_pyramid_attention_broadcast_inference(self, expected_atol: float = 0.2) device = "cpu" # ensure determinism for the device-dependent torch.Generator num_layers = 2 components = self.get_dummy_components(num_layers=num_layers) + for key in components: + if "text_encoder" in key and hasattr(components[key], "eval"): + components[key].eval() pipe = self.pipeline_class(**components) pipe = pipe.to(device) pipe.set_progress_bar_config(disable=None) @@ -2976,6 +3018,59 @@ def run_forward(pipe): ) +class MagCacheTesterMixin: + mag_cache_config = MagCacheConfig( + threshold=0.06, + max_skip_steps=3, + retention_ratio=0.2, + num_inference_steps=50, + mag_ratios=torch.ones(50), + ) + + def test_mag_cache_inference(self, expected_atol: float = 0.1): + device = "cpu" + + def create_pipe(): + torch.manual_seed(0) + num_layers = 2 + components = self.get_dummy_components(num_layers=num_layers) + pipe = self.pipeline_class(**components) + pipe = pipe.to(device) + pipe.set_progress_bar_config(disable=None) + return pipe + + def run_forward(pipe): + torch.manual_seed(0) + inputs = self.get_dummy_inputs(device) + # Match the config steps + inputs["num_inference_steps"] = 50 + return pipe(**inputs)[0] + + # 1. Run inference without MagCache (Baseline) + pipe = create_pipe() + output = run_forward(pipe).flatten() + original_image_slice = np.concatenate((output[:8], output[-8:])) + + # 2. Run inference with MagCache ENABLED + pipe = create_pipe() + pipe.transformer.enable_cache(self.mag_cache_config) + output = run_forward(pipe).flatten() + image_slice_enabled = np.concatenate((output[:8], output[-8:])) + + # 3. Run inference with MagCache DISABLED + pipe.transformer.disable_cache() + output = run_forward(pipe).flatten() + image_slice_disabled = np.concatenate((output[:8], output[-8:])) + + assert np.allclose(original_image_slice, image_slice_enabled, atol=expected_atol), ( + "MagCache outputs should not differ too much from baseline." + ) + + assert np.allclose(original_image_slice, image_slice_disabled, atol=1e-4), ( + "Outputs after disabling cache should match original inference exactly." + ) + + # Some models (e.g. unCLIP) are extremely likely to significantly deviate depending on which hardware is used. # This helper function is used to check that the image doesn't deviate on average more than 10 pixels from a # reference image. diff --git a/tests/pipelines/visualcloze/test_pipeline_visualcloze_combined.py b/tests/pipelines/visualcloze/test_pipeline_visualcloze_combined.py index 00ae0441fe99..bfbcd141d9b3 100644 --- a/tests/pipelines/visualcloze/test_pipeline_visualcloze_combined.py +++ b/tests/pipelines/visualcloze/test_pipeline_visualcloze_combined.py @@ -5,7 +5,7 @@ import numpy as np import torch from PIL import Image -from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel import diffusers from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxTransformer2DModel, VisualClozePipeline @@ -77,7 +77,8 @@ def get_dummy_components(self): text_encoder = CLIPTextModel(clip_text_encoder_config) torch.manual_seed(0) - text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_2 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") @@ -215,7 +216,7 @@ def test_different_task_prompts(self, expected_min_diff=1e-1): def test_callback_cfg(self): pass - def test_save_load_local(self, expected_max_difference=5e-4): + def test_save_load_local(self, expected_max_difference=1e-3): components = self.get_dummy_components() pipe = self.pipeline_class(**components) for component in pipe.components.values(): @@ -260,6 +261,9 @@ def test_save_load_optional_components(self, expected_max_difference=1e-4): if not hasattr(self.pipeline_class, "_optional_components"): return components = self.get_dummy_components() + for key in components: + if "text_encoder" in key and hasattr(components[key], "eval"): + components[key].eval() pipe = self.pipeline_class(**components) for component in pipe.components.values(): if hasattr(component, "set_default_attn_processor"): @@ -342,3 +346,7 @@ def test_save_load_float16(self, expected_max_diff=1e-2): self.assertLess( max_diff, expected_max_diff, "The output of the fp16 pipeline changed after saving and loading." ) + + @unittest.skip("Test not supported.") + def test_pipeline_with_accelerator_device_map(self): + pass diff --git a/tests/pipelines/visualcloze/test_pipeline_visualcloze_generation.py b/tests/pipelines/visualcloze/test_pipeline_visualcloze_generation.py index ab6b3ca5c587..5640f2f634d3 100644 --- a/tests/pipelines/visualcloze/test_pipeline_visualcloze_generation.py +++ b/tests/pipelines/visualcloze/test_pipeline_visualcloze_generation.py @@ -5,7 +5,7 @@ import numpy as np import torch from PIL import Image -from transformers import AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel import diffusers from diffusers import ( @@ -79,7 +79,8 @@ def get_dummy_components(self): text_encoder = CLIPTextModel(clip_text_encoder_config) torch.manual_seed(0) - text_encoder_2 = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder_2 = T5EncoderModel(config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") @@ -310,3 +311,7 @@ def test_save_load_float16(self, expected_max_diff=1e-2): @unittest.skip("Skipped due to missing layout_prompt. Needs further investigation.") def test_encode_prompt_works_in_isolation(self, extra_required_param_value_dict=None, atol=0.0001, rtol=0.0001): pass + + @unittest.skip("Needs to be revisited later.") + def test_pipeline_with_accelerator_device_map(self, expected_max_difference=0.0001): + pass diff --git a/tests/pipelines/wan/test_wan.py b/tests/pipelines/wan/test_wan.py index 106a7b294646..958e1b8c8eaf 100644 --- a/tests/pipelines/wan/test_wan.py +++ b/tests/pipelines/wan/test_wan.py @@ -18,7 +18,7 @@ import numpy as np import torch -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler, WanPipeline, WanTransformer3DModel @@ -68,7 +68,8 @@ def get_dummy_components(self): torch.manual_seed(0) # TODO: impl FlowDPMSolverMultistepScheduler scheduler = FlowMatchEulerDiscreteScheduler(shift=7.0) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") torch.manual_seed(0) diff --git a/tests/pipelines/wan/test_wan_22.py b/tests/pipelines/wan/test_wan_22.py index 56ef5ceb97ed..fd17ca414af4 100644 --- a/tests/pipelines/wan/test_wan_22.py +++ b/tests/pipelines/wan/test_wan_22.py @@ -17,14 +17,11 @@ import numpy as np import torch -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKLWan, UniPCMultistepScheduler, WanPipeline, WanTransformer3DModel -from ...testing_utils import ( - enable_full_determinism, - torch_device, -) +from ...testing_utils import enable_full_determinism, torch_device from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS from ..test_pipelines_common import PipelineTesterMixin @@ -63,7 +60,8 @@ def get_dummy_components(self): torch.manual_seed(0) scheduler = UniPCMultistepScheduler(prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=3.0) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") torch.manual_seed(0) @@ -235,7 +233,8 @@ def get_dummy_components(self): torch.manual_seed(0) scheduler = UniPCMultistepScheduler(prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=3.0) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") torch.manual_seed(0) diff --git a/tests/pipelines/wan/test_wan_22_image_to_video.py b/tests/pipelines/wan/test_wan_22_image_to_video.py index 6294d62044f3..4634047ebb73 100644 --- a/tests/pipelines/wan/test_wan_22_image_to_video.py +++ b/tests/pipelines/wan/test_wan_22_image_to_video.py @@ -18,7 +18,7 @@ import numpy as np import torch from PIL import Image -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKLWan, UniPCMultistepScheduler, WanImageToVideoPipeline, WanTransformer3DModel @@ -64,7 +64,8 @@ def get_dummy_components(self): torch.manual_seed(0) scheduler = UniPCMultistepScheduler(prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=3.0) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") torch.manual_seed(0) @@ -248,7 +249,8 @@ def get_dummy_components(self): torch.manual_seed(0) scheduler = UniPCMultistepScheduler(prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=3.0) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") torch.manual_seed(0) diff --git a/tests/pipelines/wan/test_wan_animate.py b/tests/pipelines/wan/test_wan_animate.py index d6d1b09f3620..5d634fb71849 100644 --- a/tests/pipelines/wan/test_wan_animate.py +++ b/tests/pipelines/wan/test_wan_animate.py @@ -19,6 +19,7 @@ import torch from PIL import Image from transformers import ( + AutoConfig, AutoTokenizer, CLIPImageProcessor, CLIPVisionConfig, @@ -78,7 +79,8 @@ def get_dummy_components(self): torch.manual_seed(0) scheduler = FlowMatchEulerDiscreteScheduler(shift=7.0) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") torch.manual_seed(0) diff --git a/tests/pipelines/wan/test_wan_image_to_video.py b/tests/pipelines/wan/test_wan_image_to_video.py index 07a9142f2553..7ed263abdcb5 100644 --- a/tests/pipelines/wan/test_wan_image_to_video.py +++ b/tests/pipelines/wan/test_wan_image_to_video.py @@ -19,6 +19,7 @@ import torch from PIL import Image from transformers import ( + AutoConfig, AutoTokenizer, CLIPImageProcessor, CLIPVisionConfig, @@ -68,7 +69,8 @@ def get_dummy_components(self): torch.manual_seed(0) # TODO: impl FlowDPMSolverMultistepScheduler scheduler = FlowMatchEulerDiscreteScheduler(shift=7.0) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") torch.manual_seed(0) @@ -239,7 +241,8 @@ def get_dummy_components(self): torch.manual_seed(0) # TODO: impl FlowDPMSolverMultistepScheduler scheduler = FlowMatchEulerDiscreteScheduler(shift=7.0) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") torch.manual_seed(0) diff --git a/tests/pipelines/wan/test_wan_vace.py b/tests/pipelines/wan/test_wan_vace.py index fe078c0deb8a..53becce1685d 100644 --- a/tests/pipelines/wan/test_wan_vace.py +++ b/tests/pipelines/wan/test_wan_vace.py @@ -18,7 +18,7 @@ import numpy as np import torch from PIL import Image -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import ( AutoencoderKLWan, @@ -67,7 +67,8 @@ def get_dummy_components(self): torch.manual_seed(0) scheduler = FlowMatchEulerDiscreteScheduler(shift=7.0) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") torch.manual_seed(0) diff --git a/tests/pipelines/wan/test_wan_video_to_video.py b/tests/pipelines/wan/test_wan_video_to_video.py index 27ada121ca48..3804e972b97f 100644 --- a/tests/pipelines/wan/test_wan_video_to_video.py +++ b/tests/pipelines/wan/test_wan_video_to_video.py @@ -16,7 +16,7 @@ import torch from PIL import Image -from transformers import AutoTokenizer, T5EncoderModel +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKLWan, UniPCMultistepScheduler, WanTransformer3DModel, WanVideoToVideoPipeline @@ -62,7 +62,8 @@ def get_dummy_components(self): torch.manual_seed(0) scheduler = UniPCMultistepScheduler(flow_shift=3.0) - text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") torch.manual_seed(0) diff --git a/tests/pipelines/z_image/test_z_image_inpaint.py b/tests/pipelines/z_image/test_z_image_inpaint.py new file mode 100644 index 000000000000..e904a4e44bd7 --- /dev/null +++ b/tests/pipelines/z_image/test_z_image_inpaint.py @@ -0,0 +1,396 @@ +# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import gc +import os +import unittest + +import numpy as np +import torch +from transformers import Qwen2Tokenizer, Qwen3Config, Qwen3Model + +from diffusers import ( + AutoencoderKL, + FlowMatchEulerDiscreteScheduler, + ZImageInpaintPipeline, + ZImageTransformer2DModel, +) +from diffusers.utils.testing_utils import floats_tensor + +from ...testing_utils import torch_device +from ..pipeline_params import ( + IMAGE_TO_IMAGE_IMAGE_PARAMS, + TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, + TEXT_GUIDED_IMAGE_INPAINTING_PARAMS, +) +from ..test_pipelines_common import PipelineTesterMixin, to_np + + +# Z-Image requires torch.use_deterministic_algorithms(False) due to complex64 RoPE operations +# Cannot use enable_full_determinism() which sets it to True +# Note: Z-Image does not support FP16 inference due to complex64 RoPE embeddings +os.environ["CUDA_LAUNCH_BLOCKING"] = "1" +os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" +torch.use_deterministic_algorithms(False) +torch.backends.cudnn.deterministic = True +torch.backends.cudnn.benchmark = False +if hasattr(torch.backends, "cuda"): + torch.backends.cuda.matmul.allow_tf32 = False + + +class ZImageInpaintPipelineFastTests(PipelineTesterMixin, unittest.TestCase): + pipeline_class = ZImageInpaintPipeline + params = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {"cross_attention_kwargs"} + batch_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS + image_params = frozenset(["image", "mask_image"]) + image_latents_params = IMAGE_TO_IMAGE_IMAGE_PARAMS + required_optional_params = frozenset( + [ + "num_inference_steps", + "strength", + "generator", + "latents", + "return_dict", + "callback_on_step_end", + "callback_on_step_end_tensor_inputs", + ] + ) + supports_dduf = False + test_xformers_attention = False + test_layerwise_casting = True + test_group_offloading = True + + def setUp(self): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.synchronize() + torch.manual_seed(0) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(0) + + def tearDown(self): + super().tearDown() + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.synchronize() + torch.manual_seed(0) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(0) + + def get_dummy_components(self): + torch.manual_seed(0) + transformer = ZImageTransformer2DModel( + all_patch_size=(2,), + all_f_patch_size=(1,), + in_channels=16, + dim=32, + n_layers=2, + n_refiner_layers=1, + n_heads=2, + n_kv_heads=2, + norm_eps=1e-5, + qk_norm=True, + cap_feat_dim=16, + rope_theta=256.0, + t_scale=1000.0, + axes_dims=[8, 4, 4], + axes_lens=[256, 32, 32], + ) + # `x_pad_token` and `cap_pad_token` are initialized with `torch.empty` which contains + # uninitialized memory. Set them to known values for deterministic test behavior. + with torch.no_grad(): + transformer.x_pad_token.copy_(torch.ones_like(transformer.x_pad_token.data)) + transformer.cap_pad_token.copy_(torch.ones_like(transformer.cap_pad_token.data)) + + torch.manual_seed(0) + vae = AutoencoderKL( + in_channels=3, + out_channels=3, + down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"], + up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"], + block_out_channels=[32, 64], + layers_per_block=1, + latent_channels=16, + norm_num_groups=32, + sample_size=32, + scaling_factor=0.3611, + shift_factor=0.1159, + ) + + torch.manual_seed(0) + scheduler = FlowMatchEulerDiscreteScheduler() + + torch.manual_seed(0) + config = Qwen3Config( + hidden_size=16, + intermediate_size=16, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=2, + vocab_size=151936, + max_position_embeddings=512, + ) + text_encoder = Qwen3Model(config) + tokenizer = Qwen2Tokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") + + components = { + "transformer": transformer, + "vae": vae, + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + } + return components + + def get_dummy_inputs(self, device, seed=0): + import random + + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device=device).manual_seed(seed) + + image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) + # Create mask: 1 = inpaint region, 0 = preserve region + mask_image = torch.zeros((1, 1, 32, 32), device=device) + mask_image[:, :, 8:24, 8:24] = 1.0 # Inpaint center region + + inputs = { + "prompt": "dance monkey", + "negative_prompt": "bad quality", + "image": image, + "mask_image": mask_image, + "strength": 1.0, + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 3.0, + "cfg_normalization": False, + "cfg_truncation": 1.0, + "height": 32, + "width": 32, + "max_sequence_length": 16, + "output_type": "np", + } + + return inputs + + def test_inference(self): + device = "cpu" + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + image = pipe(**inputs).images + generated_image = image[0] + self.assertEqual(generated_image.shape, (32, 32, 3)) + + def test_inference_batch_single_identical(self): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.synchronize() + torch.manual_seed(0) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(0) + self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-1) + + def test_num_images_per_prompt(self): + import inspect + + sig = inspect.signature(self.pipeline_class.__call__) + + if "num_images_per_prompt" not in sig.parameters: + return + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + + batch_sizes = [1, 2] + num_images_per_prompts = [1, 2] + + for batch_size in batch_sizes: + for num_images_per_prompt in num_images_per_prompts: + inputs = self.get_dummy_inputs(torch_device) + + for key in inputs.keys(): + if key in self.batch_params: + inputs[key] = batch_size * [inputs[key]] + + images = pipe(**inputs, num_images_per_prompt=num_images_per_prompt)[0] + + assert images.shape[0] == batch_size * num_images_per_prompt + + del pipe + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.synchronize() + + def test_attention_slicing_forward_pass( + self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3 + ): + if not self.test_attention_slicing: + return + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + for component in pipe.components.values(): + if hasattr(component, "set_default_attn_processor"): + component.set_default_attn_processor() + pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + + generator_device = "cpu" + inputs = self.get_dummy_inputs(generator_device) + output_without_slicing = pipe(**inputs)[0] + + pipe.enable_attention_slicing(slice_size=1) + inputs = self.get_dummy_inputs(generator_device) + output_with_slicing1 = pipe(**inputs)[0] + + pipe.enable_attention_slicing(slice_size=2) + inputs = self.get_dummy_inputs(generator_device) + output_with_slicing2 = pipe(**inputs)[0] + + if test_max_difference: + max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max() + max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max() + self.assertLess( + max(max_diff1, max_diff2), + expected_max_diff, + "Attention slicing should not affect the inference results", + ) + + def test_vae_tiling(self, expected_diff_max: float = 0.7): + import random + + generator_device = "cpu" + components = self.get_dummy_components() + + pipe = self.pipeline_class(**components) + pipe.to("cpu") + pipe.set_progress_bar_config(disable=None) + + # Without tiling + inputs = self.get_dummy_inputs(generator_device) + inputs["height"] = inputs["width"] = 128 + # Generate a larger image for the input + inputs["image"] = floats_tensor((1, 3, 128, 128), rng=random.Random(0)).to("cpu") + # Generate a larger mask for the input + mask = torch.zeros((1, 1, 128, 128), device="cpu") + mask[:, :, 32:96, 32:96] = 1.0 + inputs["mask_image"] = mask + output_without_tiling = pipe(**inputs)[0] + + # With tiling (standard AutoencoderKL doesn't accept parameters) + pipe.vae.enable_tiling() + inputs = self.get_dummy_inputs(generator_device) + inputs["height"] = inputs["width"] = 128 + inputs["image"] = floats_tensor((1, 3, 128, 128), rng=random.Random(0)).to("cpu") + inputs["mask_image"] = mask + output_with_tiling = pipe(**inputs)[0] + + self.assertLess( + (to_np(output_without_tiling) - to_np(output_with_tiling)).max(), + expected_diff_max, + "VAE tiling should not affect the inference results", + ) + + def test_pipeline_with_accelerator_device_map(self, expected_max_difference=1e-3): + # Z-Image RoPE embeddings (complex64) have slightly higher numerical tolerance + # Inpainting mask blending adds additional numerical variance + super().test_pipeline_with_accelerator_device_map(expected_max_difference=expected_max_difference) + + def test_group_offloading_inference(self): + # Block-level offloading conflicts with RoPE cache. Pipeline-level offloading (tested separately) works fine. + self.skipTest("Using test_pipeline_level_group_offloading_inference instead") + + def test_save_load_float16(self, expected_max_diff=1e-2): + # Z-Image does not support FP16 due to complex64 RoPE embeddings + self.skipTest("Z-Image does not support FP16 inference") + + def test_float16_inference(self, expected_max_diff=5e-2): + # Z-Image does not support FP16 due to complex64 RoPE embeddings + self.skipTest("Z-Image does not support FP16 inference") + + def test_strength_parameter(self): + """Test that strength parameter affects the output correctly.""" + device = "cpu" + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + # Test with different strength values + inputs_low_strength = self.get_dummy_inputs(device) + inputs_low_strength["strength"] = 0.2 + + inputs_high_strength = self.get_dummy_inputs(device) + inputs_high_strength["strength"] = 0.8 + + # Both should complete without errors + output_low = pipe(**inputs_low_strength).images[0] + output_high = pipe(**inputs_high_strength).images[0] + + # Outputs should be different (different amount of transformation) + self.assertFalse(np.allclose(output_low, output_high, atol=1e-3)) + + def test_invalid_strength(self): + """Test that invalid strength values raise appropriate errors.""" + device = "cpu" + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + + inputs = self.get_dummy_inputs(device) + + # Test strength < 0 + inputs["strength"] = -0.1 + with self.assertRaises(ValueError): + pipe(**inputs) + + # Test strength > 1 + inputs["strength"] = 1.5 + with self.assertRaises(ValueError): + pipe(**inputs) + + def test_mask_inpainting(self): + """Test that the mask properly controls which regions are inpainted.""" + device = "cpu" + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + # Generate with full mask (inpaint everything) + inputs_full = self.get_dummy_inputs(device) + inputs_full["mask_image"] = torch.ones((1, 1, 32, 32), device=device) + + # Generate with no mask (preserve everything) + inputs_none = self.get_dummy_inputs(device) + inputs_none["mask_image"] = torch.zeros((1, 1, 32, 32), device=device) + + # Both should complete without errors + output_full = pipe(**inputs_full).images[0] + output_none = pipe(**inputs_none).images[0] + + # Outputs should be different (full inpaint vs preserve) + self.assertFalse(np.allclose(output_full, output_none, atol=1e-3)) diff --git a/tests/quantization/bnb/test_mixed_int8.py b/tests/quantization/bnb/test_mixed_int8.py index fde3966dec97..031fdc9f9e27 100644 --- a/tests/quantization/bnb/test_mixed_int8.py +++ b/tests/quantization/bnb/test_mixed_int8.py @@ -288,31 +288,29 @@ def test_config_from_pretrained(self): self.assertTrue(linear.weight.__class__ == bnb.nn.Int8Params) self.assertTrue(hasattr(linear.weight, "SCB")) + @require_bitsandbytes_version_greater("0.48.0") def test_device_and_dtype_assignment(self): r""" Test whether trying to cast (or assigning a device to) a model after converting it in 8-bit will throw an error. Checks also if other models are casted correctly. """ - with self.assertRaises(ValueError): - # Tries with `str` - self.model_8bit.to("cpu") with self.assertRaises(ValueError): # Tries with a `dtype`` self.model_8bit.to(torch.float16) - with self.assertRaises(ValueError): - # Tries with a `device` - self.model_8bit.to(torch.device(f"{torch_device}:0")) - with self.assertRaises(ValueError): # Tries with a `device` self.model_8bit.float() with self.assertRaises(ValueError): - # Tries with a `device` + # Tries with a `dtype` self.model_8bit.half() + # This should work with 0.48.0 + self.model_8bit.to("cpu") + self.model_8bit.to(torch.device(f"{torch_device}:0")) + # Test if we did not break anything self.model_fp16 = self.model_fp16.to(dtype=torch.float32, device=torch_device) input_dict_for_transformer = self.get_dummy_inputs() @@ -837,7 +835,7 @@ def test_serialization_sharded(self): @require_torch_version_greater_equal("2.6.0") -@require_bitsandbytes_version_greater("0.45.5") +@require_bitsandbytes_version_greater("0.48.0") class Bnb8BitCompileTests(QuantCompileTests, unittest.TestCase): @property def quantization_config(self): @@ -848,7 +846,7 @@ def quantization_config(self): ) @pytest.mark.xfail( - reason="Test fails because of an offloading problem from Accelerate with confusion in hooks." + reason="Test fails because of a type change when recompiling." " Test passes without recompilation context manager. Refer to https://github.com/huggingface/diffusers/pull/12002/files#r2240462757 for details." ) def test_torch_compile(self): @@ -858,6 +856,5 @@ def test_torch_compile(self): def test_torch_compile_with_cpu_offload(self): super()._test_torch_compile_with_cpu_offload(torch_dtype=torch.float16) - @pytest.mark.xfail(reason="Test fails because of an offloading problem from Accelerate with confusion in hooks.") def test_torch_compile_with_group_offload_leaf(self): super()._test_torch_compile_with_group_offload_leaf(torch_dtype=torch.float16, use_stream=True) diff --git a/tests/quantization/torchao/test_torchao.py b/tests/quantization/torchao/test_torchao.py index e6bfc2530a5a..a722eaece4d1 100644 --- a/tests/quantization/torchao/test_torchao.py +++ b/tests/quantization/torchao/test_torchao.py @@ -74,7 +74,7 @@ @require_torch @require_torch_accelerator -@require_torchao_version_greater_or_equal("0.7.0") +@require_torchao_version_greater_or_equal("0.14.0") class TorchAoConfigTest(unittest.TestCase): def test_to_dict(self): """ @@ -132,7 +132,7 @@ def test_repr(self): # Slices for these tests have been obtained on our aws-g6e-xlarge-plus runners @require_torch @require_torch_accelerator -@require_torchao_version_greater_or_equal("0.7.0") +@require_torchao_version_greater_or_equal("0.14.0") class TorchAoTest(unittest.TestCase): def tearDown(self): gc.collect() @@ -256,9 +256,12 @@ def test_quantization(self): # Cutlass fails to initialize for below # ("float8dq_e4m3_row", np.array([0, 0, 0, 0, 0, 0, 0, 0, 0])), # ===== - ("fp4", np.array([0.4668, 0.5195, 0.5547, 0.4199, 0.4434, 0.6445, 0.4316, 0.4531, 0.5625])), - ("fp6", np.array([0.4668, 0.5195, 0.5547, 0.4199, 0.4434, 0.6445, 0.4316, 0.4531, 0.5625])), ]) + if version.parse(importlib.metadata.version("torchao")) <= version.Version("0.14.1"): + QUANTIZATION_TYPES_TO_TEST.extend([ + ("fp4", np.array([0.4668, 0.5195, 0.5547, 0.4199, 0.4434, 0.6445, 0.4316, 0.4531, 0.5625])), + ("fp6", np.array([0.4668, 0.5195, 0.5547, 0.4199, 0.4434, 0.6445, 0.4316, 0.4531, 0.5625])), + ]) # fmt: on for quantization_name, expected_slice in QUANTIZATION_TYPES_TO_TEST: @@ -271,6 +274,34 @@ def test_quantization(self): ) self._test_quant_type(quantization_config, expected_slice, model_id) + @unittest.skip("Skipping floatx quantization tests") + def test_floatx_quantization(self): + for model_id in ["hf-internal-testing/tiny-flux-pipe", "hf-internal-testing/tiny-flux-sharded"]: + if TorchAoConfig._is_xpu_or_cuda_capability_atleast_8_9(): + if version.parse(importlib.metadata.version("torchao")) <= version.Version("0.14.1"): + quantization_config = TorchAoConfig(quant_type="fp4", modules_to_not_convert=["x_embedder"]) + self._test_quant_type( + quantization_config, + np.array( + [ + 0.4648, + 0.5195, + 0.5547, + 0.4180, + 0.4434, + 0.6445, + 0.4316, + 0.4531, + 0.5625, + ] + ), + model_id, + ) + else: + # Make sure the correct error is thrown + with self.assertRaisesRegex(ValueError, "Please downgrade"): + quantization_config = TorchAoConfig(quant_type="fp4", modules_to_not_convert=["x_embedder"]) + def test_int4wo_quant_bfloat16_conversion(self): """ Tests whether the dtype of model will be modified to bfloat16 for int4 weight-only quantization. @@ -556,7 +587,7 @@ def test_aobase_config(self): # Slices for these tests have been obtained on our aws-g6e-xlarge-plus runners @require_torch @require_torch_accelerator -@require_torchao_version_greater_or_equal("0.7.0") +@require_torchao_version_greater_or_equal("0.14.0") class TorchAoSerializationTest(unittest.TestCase): model_name = "hf-internal-testing/tiny-flux-pipe" @@ -667,23 +698,22 @@ def test_aobase_config(self): self._check_serialization_expected_slice(quant_method, quant_method_kwargs, expected_slice, device) -@require_torchao_version_greater_or_equal("0.7.0") +@require_torchao_version_greater_or_equal("0.14.0") class TorchAoCompileTest(QuantCompileTests, unittest.TestCase): @property def quantization_config(self): return PipelineQuantizationConfig( - quant_mapping={ - "transformer": TorchAoConfig(quant_type="int8_weight_only"), - }, + quant_mapping={"transformer": TorchAoConfig(Int8WeightOnlyConfig())}, ) - @unittest.skip( - "Changing the device of AQT tensor with module._apply (called from doing module.to() in accelerate) does not work " - "when compiling." - ) def test_torch_compile_with_cpu_offload(self): + pipe = self._init_pipeline(self.quantization_config, torch.bfloat16) + pipe.enable_model_cpu_offload() + # No compilation because it fails with: # RuntimeError: _apply(): Couldn't swap Linear.weight - super().test_torch_compile_with_cpu_offload() + + # small resolutions to ensure speedy execution. + pipe("a dog", num_inference_steps=2, max_sequence_length=16, height=256, width=256) @parameterized.expand([False, True]) @unittest.skip( @@ -714,7 +744,7 @@ def test_torch_compile_with_group_offload_leaf(self, use_stream): # Slices for these tests have been obtained on our aws-g6e-xlarge-plus runners @require_torch @require_torch_accelerator -@require_torchao_version_greater_or_equal("0.7.0") +@require_torchao_version_greater_or_equal("0.14.0") @slow @nightly class SlowTorchAoTests(unittest.TestCase): @@ -794,8 +824,11 @@ def test_quantization(self): if TorchAoConfig._is_xpu_or_cuda_capability_atleast_8_9(): QUANTIZATION_TYPES_TO_TEST.extend([ ("float8wo_e4m3", np.array([0.0546, 0.0722, 0.1328, 0.0468, 0.0585, 0.1367, 0.0605, 0.0703, 0.1328, 0.0625, 0.0703, 0.1445, 0.0585, 0.0703, 0.1406, 0.0605, 0.3496, 0.7109, 0.4843, 0.4042, 0.7226, 0.5000, 0.4160, 0.7031, 0.4824, 0.3886, 0.6757, 0.4667, 0.3710, 0.6679, 0.4902, 0.4238])), - ("fp5_e3m1", np.array([0.0527, 0.0762, 0.1309, 0.0449, 0.0645, 0.1328, 0.0566, 0.0723, 0.125, 0.0566, 0.0703, 0.1328, 0.0566, 0.0742, 0.1348, 0.0566, 0.3633, 0.7617, 0.5273, 0.4277, 0.7891, 0.5469, 0.4375, 0.8008, 0.5586, 0.4336, 0.7383, 0.5156, 0.3906, 0.6992, 0.5156, 0.4375])), ]) + if version.parse(importlib.metadata.version("torchao")) <= version.Version("0.14.1"): + QUANTIZATION_TYPES_TO_TEST.extend([ + ("fp5_e3m1", np.array([0.0527, 0.0762, 0.1309, 0.0449, 0.0645, 0.1328, 0.0566, 0.0723, 0.125, 0.0566, 0.0703, 0.1328, 0.0566, 0.0742, 0.1348, 0.0566, 0.3633, 0.7617, 0.5273, 0.4277, 0.7891, 0.5469, 0.4375, 0.8008, 0.5586, 0.4336, 0.7383, 0.5156, 0.3906, 0.6992, 0.5156, 0.4375])), + ]) # fmt: on for quantization_name, expected_slice in QUANTIZATION_TYPES_TO_TEST: @@ -873,7 +906,7 @@ def test_memory_footprint_int8wo(self): @require_torch @require_torch_accelerator -@require_torchao_version_greater_or_equal("0.7.0") +@require_torchao_version_greater_or_equal("0.14.0") @slow @nightly class SlowTorchAoPreserializedModelTests(unittest.TestCase): diff --git a/tests/remote/test_remote_decode.py b/tests/remote/test_remote_decode.py index 27170cba0835..e48ddcb17438 100644 --- a/tests/remote/test_remote_decode.py +++ b/tests/remote/test_remote_decode.py @@ -14,7 +14,6 @@ # limitations under the License. import unittest -from typing import Tuple, Union import numpy as np import PIL.Image @@ -44,13 +43,13 @@ class RemoteAutoencoderKLMixin: - shape: Tuple[int, ...] = None - out_hw: Tuple[int, int] = None + shape: tuple[int, ...] = None + out_hw: tuple[int, int] = None endpoint: str = None dtype: torch.dtype = None scaling_factor: float = None shift_factor: float = None - processor_cls: Union[VaeImageProcessor, VideoProcessor] = None + processor_cls: VaeImageProcessor | VideoProcessor = None output_pil_slice: torch.Tensor = None output_pt_slice: torch.Tensor = None partial_postprocess_return_pt_slice: torch.Tensor = None diff --git a/tests/testing_utils.py b/tests/testing_utils.py index 4550813259af..53c1b8aa26ce 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -18,7 +18,7 @@ from contextlib import contextmanager from io import BytesIO, StringIO from pathlib import Path -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple, Union +from typing import TYPE_CHECKING, Any, Callable, Tuple, Union import numpy as np import PIL.Image @@ -38,6 +38,7 @@ is_gguf_available, is_kernels_available, is_note_seq_available, + is_nvidia_modelopt_version, is_onnx_available, is_opencv_available, is_optimum_quanto_available, @@ -130,6 +131,59 @@ def torch_all_close(a, b, *args, **kwargs): return True +def assert_tensors_close( + actual: "torch.Tensor", + expected: "torch.Tensor", + atol: float = 1e-5, + rtol: float = 1e-5, + msg: str = "", +) -> None: + """ + Assert that two tensors are close within tolerance. + + Uses the same formula as torch.allclose: |actual - expected| <= atol + rtol * |expected| + Provides concise, actionable error messages without dumping full tensors. + + Args: + actual: The actual tensor from the computation. + expected: The expected tensor to compare against. + atol: Absolute tolerance. + rtol: Relative tolerance. + msg: Optional message prefix for the assertion error. + + Raises: + AssertionError: If tensors have different shapes or values exceed tolerance. + + Example: + >>> assert_tensors_close(output, expected_output, atol=1e-5, rtol=1e-5, msg="Forward pass") + """ + if not is_torch_available(): + raise ValueError("PyTorch needs to be installed to use this function.") + + if actual.shape != expected.shape: + raise AssertionError(f"{msg} Shape mismatch: actual {actual.shape} vs expected {expected.shape}") + + if not torch.allclose(actual, expected, atol=atol, rtol=rtol): + abs_diff = (actual - expected).abs() + max_diff = abs_diff.max().item() + + flat_idx = abs_diff.argmax().item() + max_idx = tuple(idx.item() for idx in torch.unravel_index(torch.tensor(flat_idx), actual.shape)) + + threshold = atol + rtol * expected.abs() + mismatched = (abs_diff > threshold).sum().item() + total = actual.numel() + + raise AssertionError( + f"{msg}\n" + f"Tensors not close! Mismatched elements: {mismatched}/{total} ({100 * mismatched / total:.1f}%)\n" + f" Max diff: {max_diff:.6e} at index {max_idx}\n" + f" Actual: {actual.flatten()[flat_idx].item():.6e}\n" + f" Expected: {expected.flatten()[flat_idx].item():.6e}\n" + f" atol: {atol:.6e}, rtol: {rtol:.6e}" + ) + + def numpy_cosine_similarity_distance(a, b): similarity = np.dot(a, b) / (norm(a) * norm(b)) distance = 1.0 - similarity.mean() @@ -241,7 +295,6 @@ def parse_flag_from_env(key, default=False): _run_slow_tests = parse_flag_from_env("RUN_SLOW", default=False) _run_nightly_tests = parse_flag_from_env("RUN_NIGHTLY", default=False) -_run_compile_tests = parse_flag_from_env("RUN_COMPILE", default=False) def floats_tensor(shape, scale=1.0, rng=None, name=None): @@ -282,12 +335,155 @@ def nightly(test_case): def is_torch_compile(test_case): """ - Decorator marking a test that runs compile tests in the diffusers CI. + Decorator marking a test as a torch.compile test. These tests can be filtered using: + pytest -m "not compile" to skip + pytest -m compile to run only these tests + """ + return pytest.mark.compile(test_case) + + +def is_single_file(test_case): + """ + Decorator marking a test as a single file loading test. These tests can be filtered using: + pytest -m "not single_file" to skip + pytest -m single_file to run only these tests + """ + return pytest.mark.single_file(test_case) + + +def is_lora(test_case): + """ + Decorator marking a test as a LoRA test. These tests can be filtered using: + pytest -m "not lora" to skip + pytest -m lora to run only these tests + """ + return pytest.mark.lora(test_case) + + +def is_ip_adapter(test_case): + """ + Decorator marking a test as an IP Adapter test. These tests can be filtered using: + pytest -m "not ip_adapter" to skip + pytest -m ip_adapter to run only these tests + """ + return pytest.mark.ip_adapter(test_case) + + +def is_training(test_case): + """ + Decorator marking a test as a training test. These tests can be filtered using: + pytest -m "not training" to skip + pytest -m training to run only these tests + """ + return pytest.mark.training(test_case) + + +def is_attention(test_case): + """ + Decorator marking a test as an attention test. These tests can be filtered using: + pytest -m "not attention" to skip + pytest -m attention to run only these tests + """ + return pytest.mark.attention(test_case) + + +def is_memory(test_case): + """ + Decorator marking a test as a memory optimization test. These tests can be filtered using: + pytest -m "not memory" to skip + pytest -m memory to run only these tests + """ + return pytest.mark.memory(test_case) + + +def is_cpu_offload(test_case): + """ + Decorator marking a test as a CPU offload test. These tests can be filtered using: + pytest -m "not cpu_offload" to skip + pytest -m cpu_offload to run only these tests + """ + return pytest.mark.cpu_offload(test_case) + + +def is_group_offload(test_case): + """ + Decorator marking a test as a group offload test. These tests can be filtered using: + pytest -m "not group_offload" to skip + pytest -m group_offload to run only these tests + """ + return pytest.mark.group_offload(test_case) + + +def is_quantization(test_case): + """ + Decorator marking a test as a quantization test. These tests can be filtered using: + pytest -m "not quantization" to skip + pytest -m quantization to run only these tests + """ + return pytest.mark.quantization(test_case) + + +def is_bitsandbytes(test_case): + """ + Decorator marking a test as a BitsAndBytes quantization test. These tests can be filtered using: + pytest -m "not bitsandbytes" to skip + pytest -m bitsandbytes to run only these tests + """ + return pytest.mark.bitsandbytes(test_case) + + +def is_quanto(test_case): + """ + Decorator marking a test as a Quanto quantization test. These tests can be filtered using: + pytest -m "not quanto" to skip + pytest -m quanto to run only these tests + """ + return pytest.mark.quanto(test_case) + + +def is_torchao(test_case): + """ + Decorator marking a test as a TorchAO quantization test. These tests can be filtered using: + pytest -m "not torchao" to skip + pytest -m torchao to run only these tests + """ + return pytest.mark.torchao(test_case) + + +def is_gguf(test_case): + """ + Decorator marking a test as a GGUF quantization test. These tests can be filtered using: + pytest -m "not gguf" to skip + pytest -m gguf to run only these tests + """ + return pytest.mark.gguf(test_case) + + +def is_modelopt(test_case): + """ + Decorator marking a test as a NVIDIA ModelOpt quantization test. These tests can be filtered using: + pytest -m "not modelopt" to skip + pytest -m modelopt to run only these tests + """ + return pytest.mark.modelopt(test_case) + + +def is_context_parallel(test_case): + """ + Decorator marking a test as a context parallel inference test. These tests can be filtered using: + pytest -m "not context_parallel" to skip + pytest -m context_parallel to run only these tests + """ + return pytest.mark.context_parallel(test_case) - Compile tests are skipped by default. Set the RUN_COMPILE environment variable to a truthy value to run them. +def is_cache(test_case): + """ + Decorator marking a test as a cache test. These tests can be filtered using: + pytest -m "not cache" to skip + pytest -m cache to run only these tests """ - return pytest.mark.skipif(not _run_compile_tests, reason="test is torch compile")(test_case) + return pytest.mark.cache(test_case) def require_torch(test_case): @@ -650,6 +846,16 @@ def decorator(test_case): return decorator +def require_modelopt_version_greater_or_equal(modelopt_version): + def decorator(test_case): + return pytest.mark.skipif( + not is_nvidia_modelopt_version(">=", modelopt_version), + reason=f"Test requires modelopt with version greater than {modelopt_version}.", + )(test_case) + + return decorator + + def deprecate_after_peft_backend(test_case): """ Decorator marking a test that will be skipped after PEFT backend @@ -663,7 +869,7 @@ def get_python_version(): return major, minor -def load_numpy(arry: Union[str, np.ndarray], local_path: Optional[str] = None) -> np.ndarray: +def load_numpy(arry: str | np.ndarray, local_path: str | None = None) -> np.ndarray: if isinstance(arry, str): if local_path is not None: # local_path can be passed to correct images of tests @@ -689,14 +895,14 @@ def load_numpy(arry: Union[str, np.ndarray], local_path: Optional[str] = None) - return arry -def load_pt(url: str, map_location: Optional[str] = None, weights_only: Optional[bool] = True): +def load_pt(url: str, map_location: str | None = None, weights_only: bool = True): response = requests.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT) response.raise_for_status() arry = torch.load(BytesIO(response.content), map_location=map_location, weights_only=weights_only) return arry -def load_image(image: Union[str, PIL.Image.Image]) -> PIL.Image.Image: +def load_image(image: str | PIL.Image.Image) -> PIL.Image.Image: """ Loads `image` to a PIL Image. @@ -737,7 +943,7 @@ def preprocess_image(image: PIL.Image, batch_size: int): return 2.0 * image - 1.0 -def export_to_gif(image: List[PIL.Image.Image], output_gif_path: str = None) -> str: +def export_to_gif(image: list[PIL.Image.Image], output_gif_path: str = None) -> str: if output_gif_path is None: output_gif_path = tempfile.NamedTemporaryFile(suffix=".gif").name @@ -831,7 +1037,7 @@ def export_to_obj(mesh, output_obj_path: str = None): f.writelines("\n".join(combined_data)) -def export_to_video(video_frames: List[np.ndarray], output_video_path: str = None) -> str: +def export_to_video(video_frames: list[np.ndarray], output_video_path: str = None) -> str: if is_opencv_available(): import cv2 else: @@ -1012,7 +1218,7 @@ def summary_failures_short(tr): # Adapted from https://github.com/huggingface/transformers/blob/000e52aec8850d3fe2f360adc6fd256e5b47fe4c/src/transformers..testing_utils.py#L1905 -def is_flaky(max_attempts: int = 5, wait_before_retry: Optional[float] = None, description: Optional[str] = None): +def is_flaky(max_attempts: int = 5, wait_before_retry: float | None = None, description: str | None = None): """ To decorate flaky tests (methods or entire classes). They will be retried on failures. @@ -1151,7 +1357,12 @@ def enable_full_determinism(): # variable 'CUDA_LAUNCH_BLOCKING' or 'CUBLAS_WORKSPACE_CONFIG' to be set, # depending on the CUDA version, so we set them both here os.environ["CUDA_LAUNCH_BLOCKING"] = "1" - os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" + # Use larger workspace size for PyTorch 2.10+ to avoid CUBLAS_STATUS_NOT_INITIALIZED errors + # (catches 2.11 dev versions which report as >= 2.10) + if is_torch_version(">=", "2.10"): + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + else: + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" torch.use_deterministic_algorithms(True) # Enable CUDNN deterministic mode @@ -1269,7 +1480,7 @@ def _is_torch_fp64_available(device): # This dispatches a defined function according to the accelerator from the function definitions. -def _device_agnostic_dispatch(device: str, dispatch_table: Dict[str, Callable], *args, **kwargs): +def _device_agnostic_dispatch(device: str, dispatch_table: dict[str, Callable], *args, **kwargs): if device not in dispatch_table: return dispatch_table["default"](*args, **kwargs) @@ -1327,7 +1538,7 @@ def backend_supports_training(device: str): # Guard for when Torch is not available if is_torch_available(): # Update device function dict mapping - def update_mapping_from_spec(device_fn_dict: Dict[str, Callable], attribute_name: str): + def update_mapping_from_spec(device_fn_dict: dict[str, Callable], attribute_name: str): try: # Try to import the function directly spec_fn = getattr(device_spec_module, attribute_name) @@ -1423,10 +1634,10 @@ def _get_expected_safetensors_files( module: torch.nn.Module, offload_to_disk_path: str, offload_type: str, - num_blocks_per_group: Optional[int] = None, - block_modules: Optional[List[str]] = None, + num_blocks_per_group: int | None = None, + block_modules: list[str] | None = None, module_prefix: str = "", - ) -> Set[str]: + ) -> set[str]: expected_files = set() def get_hashed_filename(group_id: str) -> str: @@ -1506,8 +1717,8 @@ def _check_safetensors_serialization( module: torch.nn.Module, offload_to_disk_path: str, offload_type: str, - num_blocks_per_group: Optional[int] = None, - block_modules: Optional[List[str]] = None, + num_blocks_per_group: int | None = None, + block_modules: list[str] | None = None, ) -> bool: if not os.path.isdir(offload_to_disk_path): return False, None, None diff --git a/utils/custom_init_isort.py b/utils/custom_init_isort.py index cc3bccb9bd63..2194b5a57063 100644 --- a/utils/custom_init_isort.py +++ b/utils/custom_init_isort.py @@ -63,7 +63,7 @@ def get_indent(line: str) -> str: def split_code_in_indented_blocks( - code: str, indent_level: str = "", start_prompt: Optional[str] = None, end_prompt: Optional[str] = None + code: str, indent_level: str = "", start_prompt: str | None = None, end_prompt: str | None = None ) -> List[str]: """ Split some code into its indented blocks, starting at a given level. diff --git a/utils/generate_model_tests.py b/utils/generate_model_tests.py new file mode 100644 index 000000000000..11acd2175e21 --- /dev/null +++ b/utils/generate_model_tests.py @@ -0,0 +1,592 @@ +#!/usr/bin/env python +# coding=utf-8 +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Utility script to generate test suites for diffusers model classes. + +Usage: + python utils/generate_model_tests.py src/diffusers/models/transformers/transformer_flux.py + +This will analyze the model file and generate a test file with appropriate +test classes based on the model's mixins and attributes. +""" + +import argparse +import ast +import sys +from pathlib import Path + + +MIXIN_TO_TESTER = { + "ModelMixin": "ModelTesterMixin", + "PeftAdapterMixin": "LoraTesterMixin", +} + +ATTRIBUTE_TO_TESTER = { + "_cp_plan": "ContextParallelTesterMixin", + "_supports_gradient_checkpointing": "TrainingTesterMixin", +} + +ALWAYS_INCLUDE_TESTERS = [ + "ModelTesterMixin", + "MemoryTesterMixin", + "TorchCompileTesterMixin", +] + +# Attention-related class names that indicate the model uses attention +ATTENTION_INDICATORS = { + "AttentionMixin", + "AttentionModuleMixin", +} + +OPTIONAL_TESTERS = [ + # Quantization testers + ("BitsAndBytesTesterMixin", "bnb"), + ("QuantoTesterMixin", "quanto"), + ("TorchAoTesterMixin", "torchao"), + ("GGUFTesterMixin", "gguf"), + ("ModelOptTesterMixin", "modelopt"), + # Quantization compile testers + ("BitsAndBytesCompileTesterMixin", "bnb_compile"), + ("QuantoCompileTesterMixin", "quanto_compile"), + ("TorchAoCompileTesterMixin", "torchao_compile"), + ("GGUFCompileTesterMixin", "gguf_compile"), + ("ModelOptCompileTesterMixin", "modelopt_compile"), + # Cache testers + ("PyramidAttentionBroadcastTesterMixin", "pab_cache"), + ("FirstBlockCacheTesterMixin", "fbc_cache"), + ("FasterCacheTesterMixin", "faster_cache"), + # Other testers + ("SingleFileTesterMixin", "single_file"), + ("IPAdapterTesterMixin", "ip_adapter"), +] + + +class ModelAnalyzer(ast.NodeVisitor): + def __init__(self): + self.model_classes = [] + self.current_class = None + self.imports = set() + + def visit_Import(self, node: ast.Import): + for alias in node.names: + self.imports.add(alias.name.split(".")[-1]) + self.generic_visit(node) + + def visit_ImportFrom(self, node: ast.ImportFrom): + for alias in node.names: + self.imports.add(alias.name) + self.generic_visit(node) + + def visit_ClassDef(self, node: ast.ClassDef): + base_names = [] + for base in node.bases: + if isinstance(base, ast.Name): + base_names.append(base.id) + elif isinstance(base, ast.Attribute): + base_names.append(base.attr) + + if "ModelMixin" in base_names: + class_info = { + "name": node.name, + "bases": base_names, + "attributes": {}, + "has_forward": False, + "init_params": [], + } + + for item in node.body: + if isinstance(item, ast.Assign): + for target in item.targets: + if isinstance(target, ast.Name): + attr_name = target.id + if attr_name.startswith("_"): + class_info["attributes"][attr_name] = self._get_value(item.value) + + elif isinstance(item, ast.FunctionDef): + if item.name == "forward": + class_info["has_forward"] = True + class_info["forward_params"] = self._extract_func_params(item) + elif item.name == "__init__": + class_info["init_params"] = self._extract_func_params(item) + + self.model_classes.append(class_info) + + self.generic_visit(node) + + def _extract_func_params(self, func_node: ast.FunctionDef) -> list[dict]: + params = [] + args = func_node.args + + num_defaults = len(args.defaults) + num_args = len(args.args) + first_default_idx = num_args - num_defaults + + for i, arg in enumerate(args.args): + if arg.arg == "self": + continue + + param_info = {"name": arg.arg, "type": None, "default": None} + + if arg.annotation: + param_info["type"] = self._get_annotation_str(arg.annotation) + + default_idx = i - first_default_idx + if default_idx >= 0 and default_idx < len(args.defaults): + param_info["default"] = self._get_value(args.defaults[default_idx]) + + params.append(param_info) + + return params + + def _get_annotation_str(self, node) -> str: + if isinstance(node, ast.Name): + return node.id + elif isinstance(node, ast.Constant): + return repr(node.value) + elif isinstance(node, ast.Subscript): + base = self._get_annotation_str(node.value) + if isinstance(node.slice, ast.Tuple): + args = ", ".join(self._get_annotation_str(el) for el in node.slice.elts) + else: + args = self._get_annotation_str(node.slice) + return f"{base}[{args}]" + elif isinstance(node, ast.Attribute): + return f"{self._get_annotation_str(node.value)}.{node.attr}" + elif isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitOr): + left = self._get_annotation_str(node.left) + right = self._get_annotation_str(node.right) + return f"{left} | {right}" + elif isinstance(node, ast.Tuple): + return ", ".join(self._get_annotation_str(el) for el in node.elts) + return "Any" + + def _get_value(self, node): + if isinstance(node, ast.Constant): + return node.value + elif isinstance(node, ast.Name): + if node.id == "None": + return None + elif node.id == "True": + return True + elif node.id == "False": + return False + return node.id + elif isinstance(node, ast.List): + return [self._get_value(el) for el in node.elts] + elif isinstance(node, ast.Dict): + return {self._get_value(k): self._get_value(v) for k, v in zip(node.keys, node.values)} + return "" + + +def analyze_model_file(filepath: str) -> tuple[list[dict], set[str]]: + with open(filepath) as f: + source = f.read() + + tree = ast.parse(source) + analyzer = ModelAnalyzer() + analyzer.visit(tree) + + return analyzer.model_classes, analyzer.imports + + +def determine_testers(model_info: dict, include_optional: list[str], imports: set[str]) -> list[str]: + testers = list(ALWAYS_INCLUDE_TESTERS) + + for base in model_info["bases"]: + if base in MIXIN_TO_TESTER: + tester = MIXIN_TO_TESTER[base] + if tester not in testers: + testers.append(tester) + + for attr, tester in ATTRIBUTE_TO_TESTER.items(): + if attr in model_info["attributes"]: + value = model_info["attributes"][attr] + if value is not None and value is not False: + if tester not in testers: + testers.append(tester) + + if "_cp_plan" in model_info["attributes"] and model_info["attributes"]["_cp_plan"] is not None: + if "ContextParallelTesterMixin" not in testers: + testers.append("ContextParallelTesterMixin") + + # Include AttentionTesterMixin if the model imports attention-related classes + if imports & ATTENTION_INDICATORS: + testers.append("AttentionTesterMixin") + + for tester, flag in OPTIONAL_TESTERS: + if flag in include_optional: + if tester not in testers: + testers.append(tester) + + return testers + + +def generate_config_class(model_info: dict, model_name: str) -> str: + class_name = f"{model_name}TesterConfig" + model_class = model_info["name"] + forward_params = model_info.get("forward_params", []) + init_params = model_info.get("init_params", []) + + lines = [ + f"class {class_name}:", + " @property", + " def model_class(self):", + f" return {model_class}", + "", + " @property", + " def pretrained_model_name_or_path(self):", + ' return "" # TODO: Set Hub repository ID', + "", + " @property", + " def pretrained_model_kwargs(self):", + ' return {"subfolder": "transformer"}', + "", + " @property", + " def generator(self):", + ' return torch.Generator("cpu").manual_seed(0)', + "", + " def get_init_dict(self) -> dict[str, int | list[int]]:", + ] + + if init_params: + lines.append(" # __init__ parameters:") + for param in init_params: + type_str = f": {param['type']}" if param["type"] else "" + default_str = f" = {param['default']}" if param["default"] is not None else "" + lines.append(f" # {param['name']}{type_str}{default_str}") + + lines.extend( + [ + " return {}", + "", + " def get_dummy_inputs(self) -> dict[str, torch.Tensor]:", + ] + ) + + if forward_params: + lines.append(" # forward() parameters:") + for param in forward_params: + type_str = f": {param['type']}" if param["type"] else "" + default_str = f" = {param['default']}" if param["default"] is not None else "" + lines.append(f" # {param['name']}{type_str}{default_str}") + + lines.extend( + [ + " # TODO: Fill in dummy inputs", + " return {}", + "", + " @property", + " def input_shape(self) -> tuple[int, ...]:", + " return (1, 1)", + "", + " @property", + " def output_shape(self) -> tuple[int, ...]:", + " return (1, 1)", + ] + ) + + return "\n".join(lines) + + +def generate_test_class(model_name: str, config_class: str, tester: str) -> str: + tester_short = tester.replace("TesterMixin", "") + class_name = f"Test{model_name}{tester_short}" + + lines = [f"class {class_name}({config_class}, {tester}):"] + + if tester == "TorchCompileTesterMixin": + lines.extend( + [ + " @property", + " def different_shapes_for_compilation(self):", + " return [(4, 4), (4, 8), (8, 8)]", + "", + " def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]:", + " # TODO: Implement dynamic input generation", + " return {}", + ] + ) + elif tester == "IPAdapterTesterMixin": + lines.extend( + [ + " @property", + " def ip_adapter_processor_cls(self):", + " return None # TODO: Set processor class", + "", + " def modify_inputs_for_ip_adapter(self, model, inputs_dict):", + " # TODO: Add IP adapter image embeds to inputs", + " return inputs_dict", + "", + " def create_ip_adapter_state_dict(self, model):", + " # TODO: Create IP adapter state dict", + " return {}", + ] + ) + elif tester == "SingleFileTesterMixin": + lines.extend( + [ + " @property", + " def ckpt_path(self):", + ' return "" # TODO: Set checkpoint path', + "", + " @property", + " def alternate_ckpt_paths(self):", + " return []", + "", + " @property", + " def pretrained_model_name_or_path(self):", + ' return "" # TODO: Set Hub repository ID', + ] + ) + elif tester == "GGUFTesterMixin": + lines.extend( + [ + " @property", + " def gguf_filename(self):", + ' return "" # TODO: Set GGUF filename', + "", + " def get_dummy_inputs(self) -> dict[str, torch.Tensor]:", + " # TODO: Override with larger inputs for quantization tests", + " return {}", + ] + ) + elif tester in ["BitsAndBytesTesterMixin", "QuantoTesterMixin", "TorchAoTesterMixin", "ModelOptTesterMixin"]: + lines.extend( + [ + " def get_dummy_inputs(self) -> dict[str, torch.Tensor]:", + " # TODO: Override with larger inputs for quantization tests", + " return {}", + ] + ) + elif tester in [ + "BitsAndBytesCompileTesterMixin", + "QuantoCompileTesterMixin", + "TorchAoCompileTesterMixin", + "ModelOptCompileTesterMixin", + ]: + lines.extend( + [ + " def get_dummy_inputs(self) -> dict[str, torch.Tensor]:", + " # TODO: Override with larger inputs for quantization compile tests", + " return {}", + ] + ) + elif tester == "GGUFCompileTesterMixin": + lines.extend( + [ + " @property", + " def gguf_filename(self):", + ' return "" # TODO: Set GGUF filename', + "", + " def get_dummy_inputs(self) -> dict[str, torch.Tensor]:", + " # TODO: Override with larger inputs for quantization compile tests", + " return {}", + ] + ) + elif tester in [ + "PyramidAttentionBroadcastTesterMixin", + "FirstBlockCacheTesterMixin", + "FasterCacheTesterMixin", + ]: + lines.append(" pass") + elif tester == "LoraHotSwappingForModelTesterMixin": + lines.extend( + [ + " @property", + " def different_shapes_for_compilation(self):", + " return [(4, 4), (4, 8), (8, 8)]", + "", + " def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]:", + " # TODO: Implement dynamic input generation", + " return {}", + ] + ) + else: + lines.append(" pass") + + return "\n".join(lines) + + +def generate_test_file(model_info: dict, model_filepath: str, include_optional: list[str], imports: set[str]) -> str: + model_name = model_info["name"].replace("2DModel", "").replace("3DModel", "").replace("Model", "") + testers = determine_testers(model_info, include_optional, imports) + tester_imports = sorted(set(testers) - {"LoraHotSwappingForModelTesterMixin"}) + + lines = [ + "# coding=utf-8", + "# Copyright 2025 HuggingFace Inc.", + "#", + '# Licensed under the Apache License, Version 2.0 (the "License");', + "# you may not use this file except in compliance with the License.", + "# You may obtain a copy of the License at", + "#", + "# http://www.apache.org/licenses/LICENSE-2.0", + "#", + "# Unless required by applicable law or agreed to in writing, software", + '# distributed under the License is distributed on an "AS IS" BASIS,', + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.", + "# See the License for the specific language governing permissions and", + "# limitations under the License.", + "", + "import torch", + "", + f"from diffusers import {model_info['name']}", + "from diffusers.utils.torch_utils import randn_tensor", + "", + "from ...testing_utils import enable_full_determinism, torch_device", + ] + + if "LoraTesterMixin" in testers: + lines.append("from ..test_modeling_common import LoraHotSwappingForModelTesterMixin") + + lines.extend( + [ + "from ..testing_utils import (", + *[f" {tester}," for tester in sorted(tester_imports)], + ")", + "", + "", + "enable_full_determinism()", + "", + "", + ] + ) + + config_class = f"{model_name}TesterConfig" + lines.append(generate_config_class(model_info, model_name)) + lines.append("") + lines.append("") + + for tester in testers: + lines.append(generate_test_class(model_name, config_class, tester)) + lines.append("") + lines.append("") + + if "LoraTesterMixin" in testers: + lines.append(generate_test_class(model_name, config_class, "LoraHotSwappingForModelTesterMixin")) + lines.append("") + lines.append("") + + return "\n".join(lines).rstrip() + "\n" + + +def get_test_output_path(model_filepath: str) -> str: + path = Path(model_filepath) + model_filename = path.stem + + if "transformers" in path.parts: + return f"tests/models/transformers/test_models_{model_filename}.py" + elif "unets" in path.parts: + return f"tests/models/unets/test_models_{model_filename}.py" + elif "autoencoders" in path.parts: + return f"tests/models/autoencoders/test_models_{model_filename}.py" + else: + return f"tests/models/test_models_{model_filename}.py" + + +def main(): + parser = argparse.ArgumentParser(description="Generate test suite for a diffusers model class") + parser.add_argument( + "model_filepath", + type=str, + help="Path to the model file (e.g., src/diffusers/models/transformers/transformer_flux.py)", + ) + parser.add_argument( + "--output", "-o", type=str, default=None, help="Output file path (default: auto-generated based on model path)" + ) + parser.add_argument( + "--include", + "-i", + type=str, + nargs="*", + default=[], + choices=[ + "bnb", + "quanto", + "torchao", + "gguf", + "modelopt", + "bnb_compile", + "quanto_compile", + "torchao_compile", + "gguf_compile", + "modelopt_compile", + "pab_cache", + "fbc_cache", + "faster_cache", + "single_file", + "ip_adapter", + "all", + ], + help="Optional testers to include", + ) + parser.add_argument( + "--class-name", + "-c", + type=str, + default=None, + help="Specific model class to generate tests for (default: first model class found)", + ) + parser.add_argument("--dry-run", action="store_true", help="Print generated code without writing to file") + + args = parser.parse_args() + + if not Path(args.model_filepath).exists(): + print(f"Error: File not found: {args.model_filepath}", file=sys.stderr) + sys.exit(1) + + model_classes, imports = analyze_model_file(args.model_filepath) + + if not model_classes: + print(f"Error: No model classes found in {args.model_filepath}", file=sys.stderr) + sys.exit(1) + + if args.class_name: + model_info = next((m for m in model_classes if m["name"] == args.class_name), None) + if not model_info: + available = [m["name"] for m in model_classes] + print(f"Error: Class '{args.class_name}' not found. Available: {available}", file=sys.stderr) + sys.exit(1) + else: + model_info = model_classes[0] + if len(model_classes) > 1: + print(f"Multiple model classes found, using: {model_info['name']}", file=sys.stderr) + print("Use --class-name to specify a different class", file=sys.stderr) + + include_optional = args.include + if "all" in include_optional: + include_optional = [flag for _, flag in OPTIONAL_TESTERS] + + generated_code = generate_test_file(model_info, args.model_filepath, include_optional, imports) + + if args.dry_run: + print(generated_code) + else: + output_path = args.output or get_test_output_path(args.model_filepath) + output_dir = Path(output_path).parent + output_dir.mkdir(parents=True, exist_ok=True) + + with open(output_path, "w") as f: + f.write(generated_code) + + print(f"Generated test file: {output_path}") + print(f"Model class: {model_info['name']}") + print(f"Detected attributes: {list(model_info['attributes'].keys())}") + + +if __name__ == "__main__": + main() diff --git a/utils/modular_auto_docstring.py b/utils/modular_auto_docstring.py new file mode 100644 index 000000000000..fc4a82f98ea1 --- /dev/null +++ b/utils/modular_auto_docstring.py @@ -0,0 +1,352 @@ +# coding=utf-8 +# Copyright 2025 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Auto Docstring Generator for Modular Pipeline Blocks + +This script scans Python files for classes that have `# auto_docstring` comment above them +and inserts/updates the docstring from the class's `doc` property. + +Run from the root of the repo: + python utils/modular_auto_docstring.py [path] [--fix_and_overwrite] + +Examples: + # Check for auto_docstring markers (will error if found without proper docstring) + python utils/modular_auto_docstring.py + + # Check specific directory + python utils/modular_auto_docstring.py src/diffusers/modular_pipelines/ + + # Fix and overwrite the docstrings + python utils/modular_auto_docstring.py --fix_and_overwrite + +Usage in code: + # auto_docstring + class QwenImageAutoVaeEncoderStep(AutoPipelineBlocks): + # docstring will be automatically inserted here + + @property + def doc(self): + return "Your docstring content..." +""" + +import argparse +import ast +import glob +import importlib +import os +import re +import subprocess +import sys + + +# All paths are set with the intent you should run this script from the root of the repo +DIFFUSERS_PATH = "src/diffusers" +REPO_PATH = "." + +# Pattern to match the auto_docstring comment +AUTO_DOCSTRING_PATTERN = re.compile(r"^\s*#\s*auto_docstring\s*$") + + +def setup_diffusers_import(): + """Setup import path to use the local diffusers module.""" + src_path = os.path.join(REPO_PATH, "src") + if src_path not in sys.path: + sys.path.insert(0, src_path) + + +def get_module_from_filepath(filepath: str) -> str: + """Convert a filepath to a module name.""" + filepath = os.path.normpath(filepath) + + if filepath.startswith("src" + os.sep): + filepath = filepath[4:] + + if filepath.endswith(".py"): + filepath = filepath[:-3] + + module_name = filepath.replace(os.sep, ".") + return module_name + + +def load_module(filepath: str): + """Load a module from filepath.""" + setup_diffusers_import() + module_name = get_module_from_filepath(filepath) + + try: + module = importlib.import_module(module_name) + return module + except Exception as e: + print(f"Warning: Could not import module {module_name}: {e}") + return None + + +def get_doc_from_class(module, class_name: str) -> str: + """Get the doc property from an instantiated class.""" + if module is None: + return None + + cls = getattr(module, class_name, None) + if cls is None: + return None + + try: + instance = cls() + if hasattr(instance, "doc"): + return instance.doc + except Exception as e: + print(f"Warning: Could not instantiate {class_name}: {e}") + + return None + + +def find_auto_docstring_classes(filepath: str) -> list: + """ + Find all classes in a file that have # auto_docstring comment above them. + + Returns list of (class_name, class_line_number, has_existing_docstring, docstring_end_line) + """ + with open(filepath, "r", encoding="utf-8", newline="\n") as f: + lines = f.readlines() + + # Parse AST to find class locations and their docstrings + content = "".join(lines) + try: + tree = ast.parse(content) + except SyntaxError as e: + print(f"Syntax error in {filepath}: {e}") + return [] + + # Build a map of class_name -> (class_line, has_docstring, docstring_end_line) + class_info = {} + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef): + has_docstring = False + docstring_end_line = node.lineno # default to class line + + if node.body and isinstance(node.body[0], ast.Expr): + first_stmt = node.body[0] + if isinstance(first_stmt.value, ast.Constant) and isinstance(first_stmt.value.value, str): + has_docstring = True + docstring_end_line = first_stmt.end_lineno or first_stmt.lineno + + class_info[node.name] = (node.lineno, has_docstring, docstring_end_line) + + # Now scan for # auto_docstring comments + classes_to_update = [] + + for i, line in enumerate(lines): + if AUTO_DOCSTRING_PATTERN.match(line): + # Found the marker, look for class definition on next non-empty, non-comment line + j = i + 1 + while j < len(lines): + next_line = lines[j].strip() + if next_line and not next_line.startswith("#"): + break + j += 1 + + if j < len(lines) and lines[j].strip().startswith("class "): + # Extract class name + match = re.match(r"class\s+(\w+)", lines[j].strip()) + if match: + class_name = match.group(1) + if class_name in class_info: + class_line, has_docstring, docstring_end_line = class_info[class_name] + classes_to_update.append((class_name, class_line, has_docstring, docstring_end_line)) + + return classes_to_update + + +def strip_class_name_line(doc: str, class_name: str) -> str: + """Remove the 'class ClassName' line from the doc if present.""" + lines = doc.strip().split("\n") + if lines and lines[0].strip() == f"class {class_name}": + # Remove the class line and any blank line following it + lines = lines[1:] + while lines and not lines[0].strip(): + lines = lines[1:] + return "\n".join(lines) + + +def format_docstring(doc: str, indent: str = " ") -> str: + """Format a doc string as a properly indented docstring.""" + lines = doc.strip().split("\n") + + if len(lines) == 1: + return f'{indent}"""{lines[0]}"""\n' + else: + result = [f'{indent}"""\n'] + for line in lines: + if line.strip(): + result.append(f"{indent}{line}\n") + else: + result.append("\n") + result.append(f'{indent}"""\n') + return "".join(result) + + +def run_ruff_format(filepath: str): + """Run ruff check --fix, ruff format, and doc-builder style on a file to ensure consistent formatting.""" + try: + # First run ruff check --fix to fix any linting issues (including line length) + subprocess.run( + ["ruff", "check", "--fix", filepath], + check=False, # Don't fail if there are unfixable issues + capture_output=True, + text=True, + ) + # Then run ruff format for code formatting + subprocess.run( + ["ruff", "format", filepath], + check=True, + capture_output=True, + text=True, + ) + # Finally run doc-builder style for docstring formatting + subprocess.run( + ["doc-builder", "style", filepath, "--max_len", "119"], + check=False, # Don't fail if doc-builder has issues + capture_output=True, + text=True, + ) + print(f"Formatted {filepath}") + except subprocess.CalledProcessError as e: + print(f"Warning: formatting failed for {filepath}: {e.stderr}") + except FileNotFoundError as e: + print(f"Warning: tool not found ({e}). Skipping formatting.") + except Exception as e: + print(f"Warning: unexpected error formatting {filepath}: {e}") + + +def get_existing_docstring(lines: list, class_line: int, docstring_end_line: int) -> str: + """Extract the existing docstring content from lines.""" + # class_line is 1-indexed, docstring starts at class_line (0-indexed: class_line) + # and ends at docstring_end_line (1-indexed, inclusive) + docstring_lines = lines[class_line:docstring_end_line] + return "".join(docstring_lines) + + +def process_file(filepath: str, overwrite: bool = False) -> list: + """ + Process a file and find/insert docstrings for # auto_docstring marked classes. + + Returns list of classes that need updating. + """ + classes_to_update = find_auto_docstring_classes(filepath) + + if not classes_to_update: + return [] + + if not overwrite: + # Check mode: only verify that docstrings exist + # Content comparison is not reliable due to formatting differences + classes_needing_update = [] + for class_name, class_line, has_docstring, docstring_end_line in classes_to_update: + if not has_docstring: + # No docstring exists, needs update + classes_needing_update.append((filepath, class_name, class_line)) + return classes_needing_update + + # Load the module to get doc properties + module = load_module(filepath) + + with open(filepath, "r", encoding="utf-8", newline="\n") as f: + lines = f.readlines() + + # Process in reverse order to maintain line numbers + updated = False + for class_name, class_line, has_docstring, docstring_end_line in reversed(classes_to_update): + doc = get_doc_from_class(module, class_name) + + if doc is None: + print(f"Warning: Could not get doc for {class_name} in {filepath}") + continue + + # Remove the "class ClassName" line since it's redundant in a docstring + doc = strip_class_name_line(doc, class_name) + + # Format the new docstring with 4-space indent + new_docstring = format_docstring(doc, " ") + + if has_docstring: + # Replace existing docstring (line after class definition to docstring_end_line) + # class_line is 1-indexed, we want to replace from class_line+1 to docstring_end_line + lines = lines[:class_line] + [new_docstring] + lines[docstring_end_line:] + else: + # Insert new docstring right after class definition line + # class_line is 1-indexed, so lines[class_line-1] is the class line + # Insert at position class_line (which is right after the class line) + lines = lines[:class_line] + [new_docstring] + lines[class_line:] + + updated = True + print(f"Updated docstring for {class_name} in {filepath}") + + if updated: + with open(filepath, "w", encoding="utf-8", newline="\n") as f: + f.writelines(lines) + # Run ruff format to ensure consistent line wrapping + run_ruff_format(filepath) + + return [(filepath, cls_name, line) for cls_name, line, _, _ in classes_to_update] + + +def check_auto_docstrings(path: str = None, overwrite: bool = False): + """ + Check all files for # auto_docstring markers and optionally fix them. + """ + if path is None: + path = DIFFUSERS_PATH + + if os.path.isfile(path): + all_files = [path] + else: + all_files = glob.glob(os.path.join(path, "**/*.py"), recursive=True) + + all_markers = [] + + for filepath in all_files: + markers = process_file(filepath, overwrite) + all_markers.extend(markers) + + if not overwrite and len(all_markers) > 0: + message = "\n".join([f"- {f}: {cls} at line {line}" for f, cls, line in all_markers]) + raise ValueError( + f"Found the following # auto_docstring markers that need docstrings:\n{message}\n\n" + f"Run `python utils/modular_auto_docstring.py --fix_and_overwrite` to fix them." + ) + + if overwrite and len(all_markers) > 0: + print(f"\nProcessed {len(all_markers)} docstring(s).") + elif not overwrite and len(all_markers) == 0: + print("All # auto_docstring markers have valid docstrings.") + elif len(all_markers) == 0: + print("No # auto_docstring markers found.") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Check and fix # auto_docstring markers in modular pipeline blocks", + ) + parser.add_argument("path", nargs="?", default=None, help="File or directory to process (default: src/diffusers)") + parser.add_argument( + "--fix_and_overwrite", + action="store_true", + help="Whether to fix the docstrings by inserting them from doc property.", + ) + + args = parser.parse_args() + + check_auto_docstrings(args.path, args.fix_and_overwrite) diff --git a/utils/tests_fetcher.py b/utils/tests_fetcher.py index abdc9fd409db..59ef5499cc2a 100644 --- a/utils/tests_fetcher.py +++ b/utils/tests_fetcher.py @@ -56,7 +56,7 @@ import re from contextlib import contextmanager from pathlib import Path -from typing import Dict, List, Optional, Tuple, Union +from typing import Dict, List, Tuple, Union from git import Repo @@ -726,7 +726,7 @@ def init_test_examples_dependencies() -> Tuple[Dict[str, List[str]], List[str]]: return test_example_deps, all_examples -def create_reverse_dependency_map() -> Dict[str, List[str]]: +def create_reverse_dependency_map() -> dict[str, List[str]]: """ Create the dependency map from module/test filename to the list of modules/tests that depend on it recursively. @@ -778,7 +778,7 @@ def create_reverse_dependency_map() -> Dict[str, List[str]]: return reverse_map -def create_module_to_test_map(reverse_map: Dict[str, List[str]] = None) -> Dict[str, List[str]]: +def create_module_to_test_map(reverse_map: Dict[str, List[str]] = None) -> dict[str, List[str]]: """ Extract the tests from the reverse_dependency_map and potentially filters the model tests. @@ -864,7 +864,7 @@ def update_test_map_with_core_pipelines(json_output_file: str): json.dump(test_map, fp, ensure_ascii=False) -def create_json_map(test_files_to_run: List[str], json_output_file: Optional[str] = None): +def create_json_map(test_files_to_run: List[str], json_output_file: str | None = None): """ Creates a map from a list of tests to run to easily split them by category, when running parallelism of slow tests. @@ -909,7 +909,7 @@ def create_json_map(test_files_to_run: List[str], json_output_file: Optional[str def infer_tests_to_run( output_file: str, diff_with_last_commit: bool = False, - json_output_file: Optional[str] = None, + json_output_file: str | None = None, ): """ The main function called by the test fetcher. Determines the tests to run from the diff. @@ -1023,7 +1023,7 @@ def filter_tests(output_file: str, filters: List[str]): f.write(" ".join(test_files)) -def parse_commit_message(commit_message: str) -> Dict[str, bool]: +def parse_commit_message(commit_message: str) -> dict[str, bool]: """ Parses the commit message to detect if a command is there to skip, force all or part of the CI.